using NtToolBox
using PyPlot

n=256
f0 = load_image("NtToolBox/src/data/lena.png")
s1,s2=size(f0)
f0 = f0[Base.div(s1-n,2):Base.div(s1-n,2)+n-1, Base.div(s2-n,2):Base.div(s2-n,2)+n-1];

imageplot(f0);

rho = .8
Lambda = rand(n,n).>rho;

Phi = f -> f.*Lambda;

y = Phi(f0);

imageplot(y);

K  = f -> Grad(f)
KS = u -> -Div(u[:,:,1],u[:,:,2]);

Amplitude = u -> sqrt(sum(u.^2,3));
F = u -> sum(sum(Amplitude(u)));

ProxF = (u,lambda) -> max(0,1-lambda./repeat(Amplitude(u), outer=(1, 1, 2))).*u;

figure(figsize=(10,10))
ax = gca(projection="3d")
t = -linspace(-2,2, 201);
(Y,X) = meshgrid(t,t);
U = cat(3,Y,X);
V = ProxF(U,1);
# 3D display
surf(V[:,:,1]); 
set_cmap("jet")
#ax[:view](-150, 40)
#gca()[:invert_zaxis]()
ax[:view_init](-150, 40)

#axis("tight");
#camlight;
#shading interp;

ProxFS = (y,sigma) -> y-sigma*ProxF(y/sigma,1/sigma);

V = ProxFS(U,1);
# display
surf(V[:,:,1]);
set_cmap("jet")
gca()[:view_init](-150, 40)
axis("tight");
#camlight; shading interp;

ProxG = (f,tau) -> f + Phi(y - Phi(f));

L = 8;
sigma = 10;
tau = .9/(L*sigma);
theta = 1;

f = y;
g = K(y)*0;
f1 = f;

fold = f;
g = ProxFS( g+sigma*K(f1), sigma);
f = ProxG(  f-tau*KS(g), tau);
f1 = f + theta * (f-fold);

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

# Insert your code here.

imageplot(f);

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

# Insert your code here.

n=64
radius = 0.6;
x = linspace(-1,1,n)
Y,X = meshgrid(x,x)
f0 = Int.(max( abs(X),abs(Y) ).<radius);

a = 4
Lambda = ones(n)
Lambda[Base.div(end,2)-a:Base.div(end,2)+a,:] = 0
Phi = f -> f.*Lambda;

imageplot(f0, "Original", [1,2,1])
imageplot(Phi(f0), "Damaged", [1,2,2]);

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

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