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

n = 128;

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

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

sigma = .04;

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

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

w = 3
w1 = 2*w + 1;

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

(Y, X, dX, dY) = ndgrid(1 : n, 1 : n, -w : w, -w : w)

X = X + dX
Y = Y + dY;

X[X .< 1] = 2 .- X[X .< 1] 
Y[Y .< 1] = 2 .- Y[Y .< 1]
X[X .> n] = 2*n .- X[X .> n]
Y[Y .> n] = 2*n .- Y[Y .> n];

I = X .+ (Y .- 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

function patch(f)
    return f'[I]
end;


P = patch(f);

figure(figsize = (5,5))
for i in 1:16
    x = rand(1 : n)
    y = rand(1 : n)
    imageplot(P[x, y, :, :], "", [4, 4, i])
end

d = 25;

resh = P -> transpose((reshape(P, (n*n,w1*w1))));

remove_mean = Q -> Q - repeat(mean(Q, 1), inner = (w1*w1, 1));

P1 = remove_mean(resh(P))
C = P1*transpose(P1);

(D, V) = eig(C)
D = D[end : -1 : 1]
V = V[:, end : -1 : 1];

plot(D, linewidth = 2)
ylim(0, maximum(D))
show()

figure(figsize = (5,5))
for i in 1 : 16
    imageplot(abs(reshape(V[:, i], (w1,w1))'), "", [4, 4, i])
end

iresh = Q -> reshape(Q', (n, n, d))
descriptor = f -> iresh(V[: , 1 : d]'*remove_mean(resh(P)));

H = descriptor(f);

distance = i -> sum((H - repeat(reshape(H[i[1], i[2], :], (1, 1, 25)), inner = [n, n, 1])).^2, 3)./(w1*w1);

normalize = K -> K./sum(K)
kernel = (i, tau) -> normalize(exp(-distance(i)./(2*tau^2)));

tau = .05
i = [84, 73]
D = distance(i)
K = kernel(i, tau);

figure(figsize = (10,10))
imageplot(D[:, :], "D", [1, 2, 1])
imageplot(K[:, :], "K", [1, 2, 2])

q = 14;

selection = i -> [clamP(i[1] - q : i[1] + q, 1, n)'; clamP(i[2] - q : i[2] + q, 1, n)'];

function distance_0(i,sel)
    H1 = (H[sel[1, :], :, :])
    H2 = (H1[:, sel[2, :], :])
    return sum((H2 - repeat(reshape(H[i[1], i[2], :], (1, 1, length(H[i[1], i[2], :]))), inner = [length(sel[1, :]), length(sel[1, :]), 1])).^2, 3)/w1*w1
end

distance = i -> distance_0(i, selection(i))
kernel = (i, tau) -> normalize(exp(-distance(i)./(2*tau^2)));


sel = selection(i)
D = distance(i)
K = kernel(i, tau);

figure(figsize = (10,10))

imageplot(D[:, :], "D", [1, 2, 1])
imageplot(K[:, :], "K", [1, 2, 2])

function NLval_0(K,sel)
    f_temp = f[sel[1, :], :]
    return sum(K.*f_temp[:, sel[2, :]])
end

NLval = (i, ta) -> NLval_0(kernel(i, tau), selection(i));


(Y, X) = meshgrid(0 : n - 1, 0 : n - 1)

function arrayfun(f, X, Y)
    n = size(X)[1]
    p = size(Y)[1]
    R = zeros(n, p)
    for k in 1:n
        for l in 1:p
            R[k,l] = f(k,l)
        end
    end
    return R
end

NLmeans = tau -> arrayfun((i1, i2) -> NLval([i1,i2], tau), X, Y);

tau = .03

figure(figsize = (5,5))
imageplot(NLmeans(tau))

tau = .03;

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

## Insert your code here.

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

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

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