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

n = 256
name = "NtToolBox/src/data/flowers.png"
f0 = load_image(name, n)
imageplot(f0, "Original image")

sigma = .08;

using Distributions
f = f0 + sigma.*rand(Normal(), size(f0)[1], size(f0)[2]);

imageplot(clamP(f), string("Noisy, SNR = ", string(snr(f0,f)) ))

thresh_hard = (u, t) -> u.*(abs(u) .> t)
alpha = linspace(-3, 3, 1000)
plot(alpha, thresh_hard(alpha, 1))
axis("equal")

h = [0, .482962913145, .836516303738, .224143868042, -.129409522551]
h = h./norm(h)
Jmin = 2;

a = perform_wavortho_transf(f, Jmin, +1, h);

plot_wavelet(a, Jmin);

T = 3*sigma;

aT = thresh_hard(a, T);

plot_wavelet(aT, Jmin);

fHard = perform_wavortho_transf(aT, Jmin, -1, h);

imageplot(clamP(fHard), string("Hard, SNR = ", string(snr(f0, fHard))) )

thresh_soft = (u, t) -> max(1 - t./abs(u), 0).*u
alpha = linspace(-3, 3, 1000)
plot(alpha, thresh_soft(alpha, 1))
axis("equal");

T = 3/2*sigma;

aT = thresh_soft(a, T);

aT[1:2^Jmin, 1:2^Jmin] = a[1:2^Jmin, 1:2^Jmin];

fSoft = perform_wavortho_transf(aT, Jmin, -1, h);

imageplot(clamP(fSoft), string("Soft, SNR=", string(snr(f0, fSoft))) )

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

m = 4;

include("NtToolBox/src/ndgrid.jl")
(dY, dX) = meshgrid(0:m-1, 0:m-1)
delta = cat(2, reshape(dX, m*m, 1), reshape(dY, m*m, 1));

fTI = zeros(n, n)
T = 3*sigma
index = 1
for i in 1:m*m
    fS = circshift(f, delta[i, :])
    a = perform_wavortho_transf(fS, Jmin, 1, h)
    aT = thresh_hard(a, T)
    fS = perform_wavortho_transf(aT, Jmin, -1, h)
    fS = circshift(fS, -delta[i, :])
    fTI = i/(i + 1.0).*fTI + 1.0/(i + 1).*fS
    index += 1
end

imageplot(clamP(fTI), string("TI, SNR = ", string(snr(f0,fTI))) )
<script>
  $(document).ready(function(){
      $('div.prompt').hide();
  });
</script>