using PyPlot

srand(0);

N = 400;

P = Int(round(N/4))

A = randn(P, N) ./ sqrt(P)

S = 17

sel = randperm(N)
sel = sel[1 : S]   # indices of the nonzero elements of xsharp
xsharp = zeros(N)
xsharp[sel] = 1;

y = A*xsharp;

prox_gamma_g = (x, gamma) -> x - x./max(abs(x)./gamma, 1)

figsize = (9, 6)
t = -1 : 0.001 : 1
plot(t, prox_gamma_g(t, 0.3))
axis("equal")

pA = pinv(A) # pseudo-inverse. Equivalent to pA = A.T.dot(inv(A.dot(A.T)))
prox_f = (x, y) -> x + pA*(y - A*x)

gamma = 0.1 # try 1, 10, 0.1
rho = 1     # try 1, 1.5, 1.9

nbiter = 700

s = zeros(N)
En_array = zeros(nbiter)
for iter in 1 : nbiter  # iter goes from 1 to nbiter
    # put your code here
    En_array[iter] = maximum(sum(abs(x), 1))
x_restored = x

fig, (subfig1, subfig2) = subplots(1, 2, figsize = (16, 7)) # one figure with two horizontal subfigures
subfig1[:stem](xsharp)
subfig1[:set_ylim](0, 1.1)
subfig2[:stem](x_restored)
subfig2[:set_ylim](0, 1.1)
subfig1[:set_title](L"$x^\sharp$")
subfig2[:set_title](L"$x_\mathrm{restored}$")

plot(log10(En_array - minimum(En_array)))

S = 31
srand(0)
sel = randperm(N)
sel = sel[1 : S]   # indices of the nonzero elements of xsharp
xsharp = zeros(N)
xsharp[sel] = 1

y = A*xsharp

# put your code here

fig, (subfig1, subfig2) = subplots(1, 2, figsize = (16, 7)) # one figure with two horizontal subfigures
subfig1[:stem](xsharp)
subfig1[:set_ylim](0, 1.1)
subfig2[:stem](x_restored)
subfig1[:set_title](L"$x^\sharp$")
subfig2[:set_title](L"$x_\mathrm{restored}$")

plot(log10(En_array - minimum(En_array)))