using PyPlot

N = 1024;

s = 5;

t = -N/2 : N/2 - 1
h = (1-t.^2./s^2).*exp(-(t.^2)./(2*s^2))
h = h - mean(h);

h_tf = fft(fftshift(h))
opA = u -> real(ifft(fft(u) .* h_tf));

plot(t, h)
xlim(-100, 100)

plot(t, fftshift(abs(h_tf)))

srand(80);

s = Int(round(N*.01)) # number of nonzero elements of xsharp
sel = randperm(N)
sel = sel[1 : s]   # indices of the nonzero elements of xsharp
xsharp = zeros(N)
xsharp[sel] = sign(randn(s)) .* (1 - 0.3*rand(s));

noiselevel = 0.2;

y = opA(xsharp) + noiselevel * randn(N);

fig = figure(figsize = (14, 5))
stem(sel, xsharp[sel])
xlim(0, N - 1)
title(L"signal $x^\sharp$")

fig = figure(figsize = (14, 5))
plot(0 : N - 1, y)
xlim(0, N - 1)
title(L"signal $y$")

Lambda = 3;

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

grad_f = x -> opA(opA(x)-y);

beta = maximum(abs(fft(h)))^2
beta

gamma = 1.9 / beta;

nbiter = 2000
x = y
En_array = zeros(nbiter + 1)

En_array[1] = vecnorm(opA(x) - y)^2/2 + Lambda*maximum(sum(abs(x), 1))
for iter in 1 : nbiter  # iter goes from 1 to nbiter
    x = prox_gamma_g(x - gamma.*grad_f(x), gamma, Lambda)
    En_array[iter + 1] = vecnorm(opA(x) - y)^2/2 + Lambda*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)
subfig2[:stem](x_restored)
subfig1[:set_title](L"$x^\sharp$")
subfig2[:set_title](L"$x_\mathrm{restored}$")

plot(log10((En_array[1 : 1800] - En_array[end])./(En_array[1] - En_array[end])))

gamma = 1/beta
nbiter = 1700
rho = 0.95
x = y
En_array_overrelaxed = zeros(nbiter + 1)
En_array_overrelaxed[1] = vecnorm(opA(x) - y)^2/2 + Lambda*maximum(sum(abs(x), 1))
for iter in 1 : nbiter
    xprevious = x
    x = prox_gamma_g(x - gamma.*grad_f(x), gamma, Lambda)
    x += rho .* (x - xprevious)
    En_array_overrelaxed[iter + 1] = vecnorm(opA(x) - y)^2/2 + Lambda*maximum(sum(abs(x), 1)) 
end
x_restored = x;

plot(log10((En_array[1 : 1700] - En_array[end])./(En_array[1] - En_array[end])))
plot(log10((En_array_overrelaxed[1 : 1700] - En_array[end])./(En_array[1] - En_array[end])))
# I changed 1800 to 1700 because En_array_overrelaxed is of length 1701

gamma = 1/beta
nbiter = 1700
a = 10
x = y
En_array_fista = zeros(nbiter + 1)
En_array_fista[1] = vecnorm(opA(x) - y)^2/2 + Lambda*maximum(sum(abs(x), 1))
z = x
for iter in 1 : nbiter
    xprevious = x
    x = prox_gamma_g(z - gamma.*grad_f(z), gamma, Lambda)
    alpha = iter/(iter + 1 + a)
    z = x + alpha .* (x - xprevious)
    En_array_fista[iter + 1] = norm(opA(x) - y)^2/2 + Lambda*maximum(sum(abs(x), 1))
end
x_restored = x;

plot(log10((En_array[1 : 1700] - En_array[end])./(En_array[1] - En_array[end])))
plot(log10((En_array_overrelaxed[1 : 1700] - En_array[end])./(En_array[1] - En_array[end])))
plot(log10((abs(En_array_fista[1 : 1700] - En_array[end]))./(En_array[1] - En_array[end])))

<script>
  $(document).ready(function(){
      $('div.prompt').hide();
  });
</script>