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 = lambda x, T: x*np.maximum(1-T/np.maximum(abs(x), 1e-10*np.ones(np.shape(x))), np.zeros(np.shape(x)))
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, ti=0)
PsiS = f -> NtToolBox.perform_wavelet_transf(f, Jmin, +1, ti=0)


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

figure(figsize = (6, 6))
imageplot(clamP(SoftThreshPsi(f0, 0.1))) #There is still a problem with perform_wavelet_transf especially for dir=-1
# SoftThreshPsi(f0, 0.1)


lambd = .03

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

fSpars = y

fSpars = ProjC(fSpars, Omega)

fSpars = SoftThreshPsi(fSpars, lambd) #Problem: perform_wavelet_transf especially for dir=-1

run -i nt_solutions/inverse_5_inpainting_sparsity/exo1

## Insert your code here.

plt.figure(figsize=(6,6))
imageplot(clamp(fSpars))

run -i nt_solutions/inverse_5_inpainting_sparsity/exo2

## Insert your code here.

# J = Jmax - Jmin + 1
# u = np.hstack(([4**(-J)], 4**(-np.floor(np.arange(J + 2./3,1,-1./3)))))
# U = np.transpose(np.tile(u, (n,n,1)),(2,0,1))

# U = repmat( reshape(u,[1 1 length(u)]), [n n 1] );

J = Jmax - Jmin + 1;


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


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




lambd = .01

Xi = a -> NtToolBox.perform_wavelet_transf(a, Jmin, -1, ti = 1)
PsiS = f -> NtToolBox.perform_wavelet_transf(f, Jmin, + 1, ti = 1)
Psi = a -> Xi(a/U)

tau = 1.9*minimum(u)

a = U.*PsiS(fSpars) #Problem: perform_wavelet_transf especially for ti=-1

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

a = SoftThresh(a, lambd*tau)

run -i nt_solutions/inverse_5_inpainting_sparsity/exo3

## Insert your code here.

fTI = Psi(a)

plt.figure(figsize=(6,6))
imageplot(clamp(fTI))

run -i nt_solutions/inverse_5_inpainting_sparsity/exo4

## 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])) #Problem: perform_wavelet_transf especially for ti=-1

run -i nt_solutions/inverse_5_inpainting_sparsity/exo5

## Insert your code here.