#!/usr/bin/env python3
import numpy as np
import cvxpy as cp
import matplotlib.pyplot as plt

np.random.seed(2020)

# dimension of x
n = 200

# number of noisy samples
m = 80

# bound on the noise magnitudes
c = 0.5

# Generate the "true" vector x
# Here, the mask induces sparsity!
xmask = np.random.binomial(1,0.05,(n,1))
x = xmask*(np.random.uniform(3,10,(n,1)))

# Generate the A matrix (randomly)
A = np.random.normal(0,1,(m,n))


# observations
# (modeling the noises as uniformly distributed between -c and c)
y = np.squeeze(A@x + np.random.uniform(-c,c,(m,1)))



# Now, solve (a couple of different ways)

# Least squares
# Now, solve the optimization!
xhat_ls = cp.Variable(n)
obj_fun = cp.sum_squares(y-A@xhat_ls)
prob_ls = cp.Problem(cp.Minimize(obj_fun))
prob_ls.solve()

# Minimizing the l_2 norm of x with a constraint that
# describes the noise characteristics
xhat_l2 = cp.Variable(n)
#obj_fun = cp.sum_squares(y-A@xhat_l2)
prob_l2 = cp.Problem(cp.Minimize(cp.norm2(xhat_l2)),\
                                 [cp.abs(y-A@xhat_l2)<=c])
prob_l2.solve()

# Minimizing the l_1 norm
xhat_l1 = cp.Variable(n)
#obj_fun = cp.sum_squares(y-A@xhat_l1)
prob_l1 = cp.Problem(cp.Minimize(cp.norm1(xhat_l1)),\
                                 [cp.abs(y-A@xhat_l1)<=c])
prob_l1.solve()




plt.figure()
plt.plot(x, 'g', label='Original Vector')#, use_line_collection = True)
plt.plot(xhat_ls.value, 'r', label='LS Estimate')#, use_line_collection = True)
plt.plot(xhat_l2.value, 'k', label='Min l_2 Estimate')
plt.plot(xhat_l1.value, 'm', label='Min l_1 Estimate')
plt.legend()
plt.savefig("oct14a.jpg", format="jpeg", dpi=300, bbox_inches="tight")
plt.close()


# Just show the original and the minimum l_1 solutions
plt.figure()
plt.plot(x, 'g', label='Original Vector')#, use_line_collection = True)
plt.plot(xhat_l1.value, 'm', label='Min l_1 Estimate')
plt.legend()
plt.savefig("oct14b.jpg", format="jpeg", dpi=300, bbox_inches="tight")
plt.close()