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

n = 128*2
N = n^2;

f0 = load_image("NtToolBox/src/data/hibiscus.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), @sprintf("Noisy, SNR = %.1f dB", snr(f0, f)))

convol = (f, g) -> real(plan_ifft((plan_fft(f)*f).*(plan_fft(g)*g))
*((plan_fft(f)*f).*(plan_fft(g)*g)))

include("NtToolBox/src/ndgrid.jl")
normalize = f -> f./sum(f)
x = [collect(0 : Base.div(n, 2)); collect(-Base.div(n, 2) + 1 : -1)]
(Y, X) = meshgrid(x, x)
g = lambd -> normalize(exp(-(X.^2 .+ Y.^2)/(2*lambd^2)));

h = (f, lambd) -> convol(f, g(lambd));

lambd = 1.5

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

df = lambd -> real(sum(plan_fft(g(lambd))*g(lambd)));

SURE = (f ,hf, lambd) -> -N*sigma^2 + vecnorm(hf - f)^2 + 2*sigma^2*df(lambd); # vecnorm is for Frobenius norm

include("Exos/denoisingadv_9_sure/exo1.jl") #It takes time to run

# ntrials = 100
# nlaunch = 20
# E0 = []
# E = []

# for i in 1:nlaunch

#     f = repeat(f0, inner = [1, 1, ntrials]) + sigma.*rand(Normal(), n, n, ntrials)
#     hf = h(f, lambd)

#     #quadratic error
#     e = sum((hf - repeat(f0, inner = [1, 1, ntrials])).^2, (1, 2))
#     E0 = [E0; e[:]]

#     #sure error
#     e = -N*sigma^2 + sum((hf - f).^2, (1, 2)) + 2*sigma^2*df(lambd)
#     E = [E; e[:]]
# end


# v_true = mean(E0)
# v_sure = mean(E)
# a = v_true - 8*stdm(E0, mean(E0))
# b = v_true + 8*stdm(E0, mean(E0))
# t = linspace(a, b, 31)
# mybar = e -> hist(e[collect((i > a) & (i < b) for i in E0)], t)

# figure(figsize = (10, 7))

# subplot(2,1,1)
# s = mybar(E0)[2]
# s = [s; 0]
# bar(t[1 : end], s, width = (b-a)/31, color = "darkblue", edgecolor = "white")
# axvline(v_true, color = "red", linewidth = 3)

# subplot(2,1,2)
# s = mybar(E)[2]
# s = [s; 0]
# bar(t[1 : end], s, width = (b-a)/31, color = "darkblue",edgecolor = "white")
# axvline(v_sure, color = "red", linewidth = 3)

# show()



## Insert your code here.

include("Exos/denoisingadv_9_sure/exo2.jl")

## Insert your code here.

include("Exos/denoisingadv_9_sure/exo3.jl")

## Insert your code here.

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

h_daub = compute_wavelet_filter("Daubechies", 4)
W  = f1 -> NtToolBox.perform_wavortho_transf(f1,0,+1,h_daub)
Ws = x -> NtToolBox.perform_wavortho_transf(x,0,-1,h_daub);

figure(figsize = (10,10))
plot_wavelet(W(f0), 1)
show()

S = (x, lambd) -> max(0, 1 - lambd./max(1e-9, abs(x)) ) .* x;

h = (f1, lambd) -> Ws(S(W(f1), lambd));

lambd = 3*sigma/2

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

df = (hf, lambd) -> sum(abs(W(hf)) .> 1e-8);

SURE = (f, hf, lambd) -> -N*sigma^2 + vecnorm(hf - f)^2 + 2*sigma^2*df(hf, lambd);

include("Exos/denoisingadv_9_sure/exo4.jl")

## Insert your code here.

include("Exos/denoisingadv_9_sure/exo5.jl")

## Insert your code here.

include("Exos/denoisingadv_9_sure/exo6.jl")

## Insert your code here.

q = 4;

include("NtToolBox/src/ndgrid.jl")

(X, Y, dX, dY) = ndgrid(1:q:n-q+1, 1:q:n-q+1, 0:q-1, 0:q-1)
I = X + dX + (Y + dY - 1).*n
I

for i in 1:Base.div(n, q)
    for j in Base.div(n, q)
        I[i,j, :, :] = transpose(I[i,j, :, :])
    end
end


blocks = fw -> fw[I]





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

unblock = H -> assign(zeros(n,n), I, H)

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

S = (H,lambda) -> max(1 - lambda^2 ./ energy(H), 0) .* H

h = (f, lambd) -> Ws(unblock(S(blocks(W(f)), lambd)))

lambd = 1.1*sigma

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