using PyPlot
using NtToolBox

N = [300,200];

d = 2;

x = rand(2,N[1])-.5;

theta = 2*pi*rand(1,N[2])
r = .8 + .2*rand(1,N[2])
y = [cos(theta).*r; sin(theta).*r];

plotp = (x,col) -> scatter(x[1,:], x[2,:], s=200, edgecolors="k", c=col, linewidths=2);

figure(figsize=(10,10))

plotp(x, "b")
plotp(y, "r")

axis("off")
xlim(minimum(y[1,:])-.1,maximum(y[1,:])+.1)
ylim(minimum(y[2,:])-.1,maximum(y[2,:])+.1);

x2 = sum(x.^2,1)
y2 = sum(y.^2,1)
C = repeat(y2,outer=(N[1],1)) + repeat(x2',outer=(1,N[2])) - 2*x'*y;

p = ones(N[1])/N[1]
q = ones(N[2])/N[2];

gamma = .01;

xi = exp(-C/gamma);

b = ones(N[2]);

a = p./(xi*b)
b = q./(xi'*a);

include("NtSolutions/optimaltransp_5_entropic/exo1.jl");

# Insert your code here.

Pi = (diagm(a)*xi)*diagm(b);

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

include("NtSolutions/optimaltransp_5_entropic/exo2.jl");

# Insert your code here.

Pi = (diagm(a)*xi)*diagm(b);

figure(figsize=(10,10))

plotp(x, "b")
plotp(y, "r")

A = Pi .* (Pi .> minimum(1./Array(N))*.7)
i,j = findn(A)
plot([x[1,i],y[1,j]],[x[2,i],y[2,j]],"k",lw = 2)

A = Pi .* (Pi .> minimum(1./Array(N))*.3)
i,j = findn(A)
plot([x[1,i],y[1,j]],[x[2,i],y[2,j]],"k:",lw = 1)

axis("off")
xlim(minimum(y[1,:])-.1,maximum(y[1,:])+.1)
ylim(minimum(y[2,:])-.1,maximum(y[2,:])+.1);

N = 200;

t = collect(0:N-1)/N;

Gaussian = (t0,sigma) -> exp(-(t-t0).^2/(2*sigma.^2))
normalize = p -> p/sum(p)

sigma = .06;
p = Gaussian(.25,sigma)
q = Gaussian(.8,sigma);

vmin = .02;
p = normalize( p+maximum(p)*vmin)
q = normalize( q+maximum(q)*vmin);

figure(figsize = (10,7))

subplot(2, 1, 1)
bar(t, p, width = 1/length(t), color = "darkblue")
subplot(2, 1, 2)
bar(t, q, width = 1/length(t), color = "darkblue");

gamma = .03^2;

Y,X =meshgrid(t,t)
xi = exp(-(X-Y).^2/gamma);

b = ones(N);

a = p./(xi*b)
b = q./(xi'*a);

include("NtSolutions/optimaltransp_5_entropic/exo3.jl");

# Insert your code here.

Pi = (diagm(a)*xi)*diagm(b);
figure(figsize=(5,5))
imageplot(log(Pi+1e-5))

s = xi*(b.*t).*a./p;

figure(figsize=(5,5))

imageplot(log(Pi+1e-5))
plot(s*N,t*N, "r", linewidth=3)
xlim(0,N)
ylim(N,0);

N = 70;

names = ["disk","twodisks","letter-x","letter-z"]
vmin = .05
P = zeros(N,N,length(names))
for i in 1:length(names)
    p = load_image(string("NtToolBox/src/data/" , names[i] , ".png"),N) 
    p = normalize(rescale(p)+vmin)
    P[:,:,i] = p
end
    
K = length(names);

figure(figsize=(5,5))
for i in 1:K
    imageplot(P[:,:,i], "",[2,2,i])
    end

gamma = .04^2;

n = 41 # width of the convolution kernel
t = collect(linspace(-n/(2*N),n/(2*N),n))
g = exp(-t.^2/gamma)'
g2 = g*g'  
xi = x -> conv2(conv2(x, g)',g)'[Base.div(n,2)+1:Base.div(n,2)+size(x,1),Base.div(n,2)+1:Base.div(n,2)+size(x,2)];

xi2 = x -> g*x*g;

figure(figsize=(10,5))
imageplot(P[:,:,1],"",[1,2,1])
imageplot(xi(P[:,:,1]),"",[1,2,2])

figure(figsize=(10,5))
imageplot(P[:,:,1],"",[1,2,1])
x= P[:,:,1]
(n,m)= size(x)
N = min(n,m);

# std in pixel
sigma = .05; 

xi = function(x)
    Y,X = meshgrid( collect(1:n)/N, collect(1:n)/N ); 
    kx = exp(-(X-Y).^2/gamma);
    return kx*x*kx
end

imageplot(xi(P[:,:,1]),"",[1,2,2])

lambd = ones(K)/K;

b = ones(N,N,K)
a = copy(b);

for k in 1:K
    a[:,:,k] = P[:,:,k]./xi(b[:,:,k])
end

q = zeros(N,N)

for k in 1:K
    q = q + lambd[k] * log(max(1e-19, b[:,:,k].*xi(a[:,:,k])))
end 
q = exp(q);

for k in 1:K
    b[:,:,k] = q./xi(a[:,:,k])
end

include("NtSolutions/optimaltransp_5_entropic/exo4.jl");

gamma

# Insert your code here.

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

include("NtSolutions/optimaltransp_5_entropic/exo5.jl");

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