# -*- coding: utf-8 -*-
import numpy as np
from sklearn.linear_model import Lasso
import matplotlib.pyplot as plt
from scipy.fftpack import dct, idct

N       = 1200
M       = 120
noise   = 0.1
x       = np.linspace(0,1,N)

# signal original
y_clear = np.cos(2*np.pi*10*x) + np.cos(5*np.pi*x)

# signal bruité
y = y_clear + np.random.normal(0,noise,N)

# on choisit des indices au hasard, M parmi N, M << N
r = np.random.choice(N-1,M,replace=False)
r.sort()

# on sample le signal
y2 = y[r]

# A = matrice de DCT inverse
I = np.identity(N)
A = idct(I,norm='ortho',axis=0)
A = A[r]

# y_lasso = argmin ||Ax - y||^2_2 + alpha* ||x||_1
#           x € R^n

lasso   = Lasso(alpha=0.001)
lasso.fit(A, y2)
y_lasso = lasso.coef_

# Affichage

plt.plot(idct(y_lasso,norm='ortho',axis=0),color='blue',label='Reconstruction L_1, M = ' + str(M))
plt.plot(y_clear,color='green',label='Signal original N = ' + str(N))
plt.plot(y,color='purple',label='Signal avec bruit', alpha=0.45)
plt.legend()
plt.savefig('cs_lasso.png',bbox_inches='tight')
plt.show()