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

n = 256;

f0 = rescale(load_image("NtToolBox/src/data/boat.png", n));

figure(figsize = (5,5))
imageplot(f0)

sigma = .08;

using Distributions
f = f0 + sigma.*rand(Normal(), n, n);

figure(figsize = (5,5))
imageplot(clamP(f))

Jmin = 4;

wav  = f -> NtToolBox.perform_wavelet_transf(f, Jmin, +1)
iwav = fw -> NtToolBox.perform_wavelet_transf(fw, Jmin, -1);

figure(figsize = (10, 10))
plot_wavelet(wav(f), Jmin)
show()

psi = (s, T) -> max(1 - T^2./max(abs(s).^2, 1e-9.*ones(size(s))), zeros(size(s)));

s = linspace(-3, 3, 1024)

plot(s, s.*psi(s, 1))
plot(s, s, "r--")

show()

theta = (x, T) -> psi(x, T).*x
ThreshWav = (f, T) -> iwav(theta(wav(f), T));

T = 1.5*sigma

figure(figsize = (5, 5))
imageplot(clamP(ThreshWav(f, T)))

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

## Insert your code here.

figure(figsize = (5, 5))
imageplot(clamP(fThresh), @sprintf("SNR = %.1f dB", snr(f0, fThresh)))

w = 4
n = 256;

include("NtToolBox/src/ndgrid.jl")
(X, Y, dX, dY) = ndgrid(1:w:n-w+1, 1:w:n-w+1, 0:w-1, 0:w-1)

I = (X + dX) + (Y + dY - 1).*n
for k in 1:Base.div(n, w)
    for l in 1:Base.div(n, w)
        I[k, l, :, :] = transpose(I[k, l, :, :])
    end
end

block = x -> x[I];

function assign(M,I,H)
    M_temp = M
    M_temp[I] = H
    return reshape(M_temp, n, n)
end

iblock = H -> assign(zeros(n, n), I, H);

print(string("Should be 0:", string(norm(f - iblock(block(f))))))

function energy(H)
    H_tmp = copy(H)
    for i in 1:Base.div(n, w)
        for j in 1:Base.div(n, w)
            H_tmp[i, j, :, :] = sqrt(mean(H_tmp[i, j, :, :].^2))#*np.ones([1,1])
        end
    end
    return H_tmp
end;

Thresh = (H, T) -> psi(energy(H), T).*H
ThreshBlock = (x, T) -> iblock(Thresh(block(x), T));

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

## Insert your code here.

T = 1.25*sigma

figure(figsize = (10, 10))
plot_wavelet(ThreshBlock(wav(f), T), Jmin)
show()

ThreshWav = (f, T) -> iwav(ThreshBlock(wav(f), T));

figure(figsize = (5, 5))
imageplot(clamP(ThreshWav(f, T)))

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

## Insert your code here.

figure(figsize = (5, 5))
imageplot(clamP(fBlock), @sprintf("SNR = %.1f dB", snr(f0, fBlock)))

wav  = f -> NtToolBox.perform_wavelet_transf(f, Jmin, + 1, "9-7", 0, 1)
iwav = fw -> NtToolBox.perform_wavelet_transf(fw, Jmin, -1,"9-7", 0, 1);

fw = wav(f)
n = 256;

(X, Y, J, dX, dY) = ndgrid(1 : w : n - w + 1, 1 : w : n - w + 1, 1 : size(fw)[3], 0 : w - 1, 0 : w - 1)

I = (X + dX) + (Y + dY - 1).*n + (J - 1).*n^2
for k in 1 : Base.div(n, w)
    for l in 1 : Base.div(n, w)
        for m in 1 : size(fw)[3]
            I[k, l, m, :, :] = transpose(I[k, l, m, :, :])
        end
    end
end

block = x -> x[I]


function assign(M, I, H)
    M_temp = M
    M_temp[I] = H
    return reshape(M_temp,n, n, size(fw)[3])
end

iblock = H -> assign(zeros(n, n, size(fw)[3]), I, H);

function energy(H)
    H_tmp = copy(H)
    for i in 1 : Base.div(n, w)
        for j in 1 : Base.div(n, w)
            for k in 1 : size(fw)[3]
                H_tmp[i, j, k, :, :] = sqrt(mean(H_tmp[i, j, k, :, :].^2))
            end
        end
    end
    return H_tmp
end;

Thresh = (H, T) -> psi(energy(H), T).*H
ThreshBlock = (x, T) -> iblock(Thresh(block(x), T));

ThreshWav = (f, T) -> iwav(ThreshBlock(wav(f), T));

T = 1.25*sigma

figure(figsize = (5, 5))
imageplot(clamP(ThreshWav(f, T)))

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

# tlist = linspace(.5, 2, 20).*sigma

# snr_stein = collect(snr(f0, ThreshWav(f,t)) for t in Tlist)

# plot(collect(T/sigma for T in Tlist), snr_stein, linewidth = 2)
# xlabel(L"$T/sigma$")
# ylabel("SNR")

# show()

# Tmax = maximum(snr_stein)
# fTI = ThreshWav(f, T)

## Insert your code here.

figure(figsize = (5,5))
imageplot(clamP(fTI), @sprintf("SNR = %.1f dB", snr(f0, fTI)));
<script>
  $(document).ready(function(){
      $('div.prompt').hide();
  });
</script>