using PyPlot
using NtToolBox
# using Autoreload
# arequire("NtToolBox")

n = 256
name = "NtToolBox/src/data/hibiscus.png"
f0 = load_image(name, n);

imageplot(f0, "Image f0", [1, 2, 1])

s = 3;

include("NtToolBox/src/ndgrid.jl")
x = [collect(0:n/2 - 1); collect(-n/2:-1)]
(Y, X) = meshgrid(x, x)
h = exp((-X.^2 - Y.^2)/(2*s^2))
h = h/sum(h);

hF = real(plan_fft(h)*h);

imageplot(fftshift(h), "Filter", [1, 2, 1])
imageplot(fftshift(hF), "Fourier transform", [1, 2, 2])

Phi = (x, h) -> real(plan_ifft((plan_fft(x)*x) .* (plan_fft(h)*h))*((plan_fft(x)*x) .* (plan_fft(h)*h)))

y0 = Phi(f0, h);

imageplot(y0, "Observation without noise", [1, 2, 2])

sigma = .02;

y = y0 .+ randn(n, n).*sigma;

imageplot(clamP(y), "Observation with noise", [1, 2, 2])

imageplot(y, string("Observation, SNR = ", string(round(snr(f0, y),2)), "dB"))

yF = plan_fft(y)*y;

Lambda = 0.02;

fL2 = real(plan_ifft((yF .* hF ./ (abs(hF).^2 .+ Lambda)))*((yF .* hF ./ (abs(hF).^2 .+ Lambda))));

imageplot(clamP(fL2), string("L2 deconvolution, SNR = ", string(round(snr(f0, fL2), 2)), "dB"))

include("NtSolutions/inverse_2_deconvolution_variational/exo1.jl")

imageplot(clamP(fL2), string("L2 deconvolution, SNR = ", string(round(snr(f0, fL2), 2)), "dB") )

S = (X.^2 + Y.^2) .* (2/n).^2;

Lambda = 0.2;

fSob = real(plan_ifft((yF .* hF ./ (abs(hF).^2 .+ Lambda.*S)))*((yF .* hF ./ (abs(hF).^2 .+ Lambda.*S))));

imageplot(clamP(fSob), string("Sobolev deconvolution, SNR = ", string( round(snr(f0, fSob), 2) ), "dB") )

include("NtSolutions/inverse_2_deconvolution_variational/exo2.jl")

imageplot(clamP(fSob), string("Sobolev deconvolution, SNR = ", string( round(snr(f0, fSob), 2) ), "dB") )

epsilon = 0.4*1e-2;

Lambda = 0.06;

tau = 1.9 / (1 + Lambda * 8 / epsilon);

fTV = y;

niter = 300;

a = rand(n,n)
b = rand(n, n, 2)
dotp = (x, y) -> sum(x[:].*y[:])
print(string("Should be 0: ", string(dotp(grad(a),b) + dotp(a, NtToolBox.div(b)))) ) #We use NtToolBox.div instead of div because otherwise there will be a conflict with the function div of package Base.

repeat3 = x -> cat(3, x, x)
Gr = grad(fTV)
d = sqrt(epsilon^2 .+ sum(Gr.^2, 3))[:, :]
G = -NtToolBox.div(Gr ./ repeat3(d) ); #We use NtToolBox.div instead of div because otherwise there will be a conflict with the function div of package Base.

tv = sum(d)
print(string("Smoothed TV norm of the image: ", string(round(tv,2))) )

e = Phi(fTV, h) - y
fTV = fTV - tau*(Phi(e, h) + Lambda*G);

include("NtSolutions/inverse_2_deconvolution_variational/exo3.jl")

imageplot(clamP(fTV))

include("NtSolutions/inverse_2_deconvolution_variational/exo4.jl")

imageplot(clamP(fBest), string("TV deconvolution, SNR = ", string( round(snr(f0, fBest), 2) ), "dB") )
<script>
  $(document).ready(function(){
      $('div.prompt').hide();
  });
</script>