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

n = 128
f0 = load_image("NtToolBox/src/data/lena.png")
f0 = rescale(f0[256 - Base.div(n, 2) + 1 : 256 + Base.div(n, 2), 256 - Base.div(n, 2) + 1 : 256 + Base.div(n, 2)]);

figure(figsize = (6, 6))
imageplot(f0, L"Image $f_0$")

rho = .7;

Omega = zeros(n, n)
sel = randperm(n^2)
Omega[sel[1:Int(round(rho*n^2))]] = 1;

Phi = (f, Omega) -> f.*(1 - Omega);

y = Phi(f0, Omega);

figure(figsize = (6, 6))
imageplot(y, "Observations y")

SoftThresh = (x,T) -> x.*max( 0, 1 - T./max(abs(x), 1e-10) );

x = linspace(-1, 1, 1000)

figure(figsize = (7, 5))
plot(x, SoftThresh(x, .5))
show()

Jmax = log2(n) - 1
Jmin = (Jmax - 3);

Psi = a -> NtToolBox.perform_wavelet_transf(a, Jmin, -1, "9-7", 0, 0)
PsiS = f -> NtToolBox.perform_wavelet_transf(f, Jmin, +1, "9-7", 0, 0);

SoftThreshPsi = (f, T) -> Psi(SoftThresh(PsiS(f), T));

figure(figsize = (6, 6))
imageplot(clamP(SoftThreshPsi(f0, 0.1)))

lambd = .03;

ProjC = (f, Omega) -> Omega.*f + (1 - Omega).*y;

fSpars = y;

fSpars = ProjC(fSpars, Omega);

fSpars = SoftThreshPsi(fSpars, lambd);

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

## Insert your code here.

figure(figsize = (6, 6))
imageplot(clamP(fSpars))

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

## Insert your code here.

J = Jmax - Jmin + 1;

u = reshape([4^(-J); 4.^(-floor(J+2/3:-1/3:1))], (1, 1, 13))

U = repeat(u, inner = [n, n, 1])


lambd = .01;

Xi = a -> NtToolBox.perform_wavelet_transf(a, Jmin, -1, "9-7", 0, 1)
PsiS = f -> NtToolBox.perform_wavelet_transf(f, Jmin, + 1, "9-7", 0, 1)
Psi = a -> Xi(a./U);

tau = 1.9*minimum(u);

a = U.*PsiS(fSpars) ;

fTI = Psi(a)
a = a + tau.*PsiS(Phi(y-Phi(fTI, Omega), Omega));

a = SoftThresh(a, lambd*tau);

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

# niter = 1000
# a = U.*PsiS(fSpars)
# E = []
# for i in 1 : niter
#     fTI = Psi(a)
#     d = y - Phi(fTI, Omega)
#     E = [E; 1/2.*vecnorm(d )^2 + lambd*sum(abs(a))]
#     # step 
#     a = SoftThresh(a + tau.*PsiS(Phi(d, Omega)), lambd*tau)
# end

# figure(figsize = (7, 5))    
# plot(E)
# show()

## Insert your code here.

fTI = Psi(a);

figure(figsize = (6, 6))
imageplot(clamP(fTI))

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

## Insert your code here.

HardThresh = (x, t) -> x.*(abs(x) .> t)

x = linspace(-1, 1, 1000)

figure(figsize = (7, 5))
plot(x, HardThresh(x, .5))
show()

niter = 500

lambda_list = linspace(1, 0, niter)

fHard = y

fHard = ProjC(fHard, Omega)

fHard = Xi(HardThresh(PsiS(fHard), tau*lambda_list[1]));

include("NtSolutions/inverse_5_inpainting_sparsity/exo5.jl")

# lambda_list = linspace(1, 0, niter)
# fHard = y

# for i in 1 : niter
#     fHard = Xi(HardThresh(PsiS(ProjC(fHard, Omega)), lambda_list[i]))
# end

    
# figure(figsize = (6, 6))
# imageplot(clamP(fSpars), @sprintf("Inpainting hard thresh., SNR = %.1f dB", snr(f0, fHard)))


## Insert your code here.
<script>
  $(document).ready(function(){
      $('div.prompt').hide();
  });
</script>