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)
x = prox_f(s,y)
s += rho.*(prox_gamma_g(2.*x - s, gamma) - x)
En_array[1] = maximum(sum(abs(x), 1))
for iter in 2 : nbiter  # iter goes from 1 to nbiter
    x = prox_f(s,y)
    s += rho.*(prox_gamma_g(2.*x - s, gamma) - x)
    En_array[iter] = maximum(sum(abs(x), 1))
end
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

gamma = 1
rho = 1   
nbiter = 1000
s = zeros(N)
En_array = zeros(nbiter)
x = prox_f(s,y)
s += rho.*(prox_gamma_g(2*x - s, gamma) - x)
En_array[1] = norm(x, 1)
for iter in 2 : nbiter
    x = prox_f(s,y)
    s += rho.*(prox_gamma_g(2*x - s, gamma) - x)
    En_array[iter] = norm(x, 1)
end
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)
subfig1[:set_title](L"$x^\sharp$")
subfig2[:set_title](L"$x_\mathrm{restored}$")

plot(log10(En_array - minimum(En_array)))
<script>
  $(document).ready(function(){
      $('div.prompt').hide();
  });
</script>