from __future__ import division
from nt_toolbox.general import *
from nt_toolbox.signal import *
%pylab inline
%matplotlib inline
%load_ext autoreload
%autoreload 2

n = 256
name = 'nt_toolbox/data/hibiscus.bmp'
f0 = load_image(name, n);

imageplot(f0, 'Image f0', [1, 2, 1])

s = 3

x = concatenate( (arange(0,n/2), arange(-n/2,0)) );
[Y, X] = meshgrid(x, x)
h = exp((-X**2-Y**2)/ (2*s**2))
h = h/sum(h.flatten())

hF = real(fft2(h))

imageplot(fftshift(h), 'Filter', [1, 2, 1])
imageplot(fftshift(hF), 'Fourier transform', [1, 2, 2])

Phi = lambda x,h: real(ifft2(fft2(x) * fft2(h)))

y0 = Phi(f0, h)

imageplot(y0, 'Observation without noise', [1, 2, 2])

sigma = .02

y = y0 + randn(n,n)*sigma

imageplot(clamp(y), 'Observation with noise', [1, 2, 2])

imageplot(y, 'Observation, SNR = ' + str(round(snr(f0, y),2)) + 'dB')

yF = fft2(y)

Lambda = 0.02

fL2 = real(ifft2(yF * hF / (abs(hF)**2 + Lambda)))

imageplot(clamp(fL2), 'L2 deconvolution, SNR = ' + str(round(snr(f0, fL2),2)) + 'dB')

run -i nt_solutions/inverse_2_deconvolution_variational/exo1

imageplot(clamp(fL2), 'L2 deconvolution, SNR = ' + str(round(snr(f0, fL2), 2)) + 'dB' )

S = (X**2 + Y**2) * (2/n)**2

Lambda = 0.2

fSob = real(ifft2(yF * hF / (abs(hF)**2 + Lambda*S)))

imageplot(clamp(fSob), 'Sobolev deconvolution, SNR = ' + str( round(snr(f0, fSob), 2) ) + 'dB' )

run -i nt_solutions/inverse_2_deconvolution_variational/exo2

imageplot(clamp(fSob), 'Sobolev deconvolution, SNR = ' + str( round(snr(f0, fSob), 2) ) + 'dB' )

epsilon = 0.4*1e-2

Lambda = 0.06

tau = 1.9 / (1 + Lambda * 8 / epsilon)

fTV = y

niter = 300

a = randn(n,n)
b = randn(n,n,2)
dotp = lambda x,y: sum(x.flatten()*y.flatten())
print("Should be 0: " + str(dotp(grad(a),b) + dotp(a,div(b))) )

repeat3 = lambda x,k: resize( repeat( x, k, axis=1), [n, n, k])
Gr = grad(fTV)
d = sqrt(epsilon**2 + sum(Gr**2, axis=2))
G = -div(Gr / repeat3(d,2) )

tv = sum(d.flatten())
print('Smoothed TV norm of the image: ' + str(round(tv,2)) )

e = Phi(fTV, h)-y
fTV = fTV - tau*(Phi(e, h) + Lambda*G)

run -i nt_solutions/inverse_2_deconvolution_variational/exo3

imageplot(clamp(fTV))

run -i nt_solutions/inverse_2_deconvolution_variational/exo4

imageplot(clamp(fBest), 'TV deconvolution, SNR = ' + str( round(snr(f0, fBest), 2) ) + 'dB' )