import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
%matplotlib inline
plt.rcParams['font.size']=15

def plt_legend_out(frameon=True):
    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left', borderaxespad=0, frameon=frameon)

from scipy import sparse
from scipy.sparse.linalg import spsolve

def baseline_als(y, lam, p, niter=10):
  L = len(y)
  D = sparse.diags([1,-2,1],[0,-1,-2], shape=(L,L-2))
  w = np.ones(L)
  for i in range(niter):
    W = sparse.spdiags(w, 0, L, L)
    Z = W + lam * D.dot(D.transpose())
    z = spsolve(Z, w*y)
    w = p * (y > z) + (1-p) * (y < z)
  return z

x = np.linspace(-1,20,100)

y1 = np.sin(x) * 100
y2 = x**2 + 80
spectra = y1 + y2

plt.plot(x,spectra,label='raw spectra',color='k')
plt.xlabel('x')
plt.ylabel('y')
plt.grid()
plt_legend_out()
plt.show()

baseline = baseline_als(spectra,1000,0.001)

plt.plot(x,spectra,label='raw spectra',color='k')
plt.plot(x,baseline,label='baseline',color='r',ls='dashed')
plt.xlabel('x')
plt.ylabel('y')
plt.grid()
plt_legend_out()
plt.show()

corrected = spectra - baseline

plt.plot(x,spectra,label='raw spectra',color='k')
plt.plot(x,baseline,label='baseline',color='r',ls='dashed')
plt.plot(x,corrected,label='corrected spectra',color='g')
plt.xlabel('x')
plt.ylabel('y')
plt.grid()
plt_legend_out()
plt.show()