using SymPy

@vars k

θ = 2/(k+2)
θp = 2/(k+3)

expr = (1-θp)/θp^2 - 1/θ^2

simplify(expr)

tnext(t) = (1+sqrt(1+4t^2))/2

tseq = function(n)
    vals = zeros(n); vals[1] = 1
    for i=2:n 
        vals[i] = tnext(vals[i-1]) 
    end
    return vals
end

using SymPy, PyPlot

@vars t
tt = (1+sqrt(1+4t^2))/2
ttt = (1+sqrt(1+4tt^2))/2
θ = tt/(t-1)
θ2 = ttt/(tt-1)

expr = (1-θ2)/θ2^2 - 1/θ^2
simplify(expr)
tvals = tseq(100)
expr_vals = map(expr, tvals)
plot(tvals, N(expr_vals))
# xlabel("t")
# ylabel("expr")


#
# Noisy Compressed sensing
#
using StatsBase, PyPlot

include("../../julia/lasso.ISTA.jl")
include("../../julia/lasso.FISTA.jl")

n=1000
s=10
x=zeros(n)
x[sample(1:n,s)] = rand([-1,1], s).*sqrt(log(n))

k=round(Int, 3*s*log(n/s)) #3*s*log(n/s))
A=randn(k,n)
println("A: k=$k, n=$n\n")

# noisy measurement
y=A*x + .1*randn(k)

lam=10
maxiter=1000
eps = 1e-4

# reconstruction
ista = lassoISTA(A, y, step="const", lam=lam, maxit=maxiter, eps=eps, repf=true, repg=true);
simple = lassoFISTA(A, y, step="simple", lam=lam, maxit=maxiter, eps=eps, repf=true, repg=true);
fista = lassoFISTA(A, y, step="const", lam=lam, maxit=maxiter, eps=eps, repf=true, repg=true);

z = ista[1]
println("True signal: sparsity=$(countnz(x)), dim=$n")
println("Recovered signal: sparsity=$(countnz(z)), mse=$(vecnorm(x-z,2)^2/n)")
sx=find(abs(x).>0)
sz=find(abs(z).>0)
print(sx)
print(sz)

subplot(411); plot(1:n, x); ylabel("True signal")
subplot(412); plot(1:n, ista[1]); ylabel("Recoverd (PGD)")
subplot(413); plot(1:n, simple[1]); ylabel("Recoverd (AGD:simple)")
subplot(414); plot(1:n, fista[1]); ylabel("Recoverd (AGD:fista)")
PyPlot.subplots_adjust(hspace=.5)

subplot(131)
PyPlot.semilogy(1:maxiter, ista[2], color="blue");
PyPlot.semilogy(1:maxiter, simple[2], color="green");
PyPlot.semilogy(1:maxiter, fista[2], color="red", linestyle="dashed");
xlabel("Iterations")
ylabel("Objective function value")
ax = gca()
ax[:set_ylim]([100, 1000])
ax[:legend](("PGD", "AGD (simple)", "AGD (fista)"), loc=1)

subplot(132)
PyPlot.semilogy(1:maxiter, ista[3], color="blue");
PyPlot.semilogy(1:maxiter, simple[3], color="green");
PyPlot.semilogy(1:maxiter, fista[3], color="red", linestyle="dashed");
xlabel("Iterations")
ylabel("Gradient inf norm")
ax = gca()
ax[:set_ylim]([0, 100])

subplot(133)
PyPlot.semilogy(1:maxiter, ista[4], color="blue");
PyPlot.semilogy(1:maxiter, simple[4], color="green");
PyPlot.semilogy(1:maxiter, fista[4], color="red", linestyle="dashed");
xlabel("Iterations")
ylabel("Sparsity")
ax = gca()
ax[:set_ylim]([0, n])

PyPlot.subplots_adjust(right=1.2, wspace=.4)


using Images, ImageView, Colors, PyPlot
pic = load("./tu32.tiff");
pic

y = separate(pic);
picx,picy = size(y)
y = y[:,:,1];
y = float(y[:]);

# noise
r = .05
y[rand(picx*picy) .> 1-r/2] = 1;  # salt
y[rand(picx*picy) .< r/2] = 0;    # pepper

print("picture dim = ($picx, $picy)\n")

yim = reshape(y, picx, picy);
Image(map(RGB, yim, yim, yim)')

n = picx*picy
m = 2n - (picx + picy-2) - 2

print("m,n = $m, $n, length img vec = $(length(y))")

D = spzeros(m,n);
e=1;

valid(coord) = (coord[1] >= 1 && coord[1] <= picx && coord[2] >= 1 && coord[2] <= picy)? true:false
c2i(coord) = (coord[2]-1)*picx + coord[1]

for i=1:picx, j=1:picy

    coords = ( (i,j+1), (i+1,j) )
    
    for c in coords
        if(valid(c)) 
            D[e, c2i((i,j))] = 1;
            D[e, c2i(c)] = -1;
            e=e+1;
        end
    end
end

print("total added edges = $(e-1)")

# Soft-thresholding function
soft(v, lam) = (sign(v) .* max(abs(v) - lam, 0))

DD = D'*D

ρ = 1000; 
δ = 0;


@time L = chol(ρ*DD + (1+δ)*eye(n), Val{:L});


maxiter = 25000
λ = .3

x = y
z = D*x;
α = zeros(m)
Dalp = zeros(n)

prevX, prevZ, prevAlp = x, z, α

@printf "iter   primal    infeas      rho      diff_x    diff_z   diff_alp   pd_gap\n"

for k=1:maxiter
   
    tmp = y - Dalp + ρ*(D'*z);
    x = L'\(L\tmp);
    Dx = D*x;
    z = soft(Dx + (1/ρ)*α, λ/ρ)
    r = Dx-z
    α = α + ρ*r
    Dalp = D'*α
        
    if(k%1000==0)
        
        primal = .5*vecnorm(y-x)^2 + λ*vecnorm(z,1)
        pd_gap = (λ*vecnorm(Dx,1) + dot(Dalp,x))/max(1,primal)
        infeas = vecnorm(r,Inf)

        reld_x = vecnorm(x-prevX)/max(1,vecnorm(x))
        reld_z = vecnorm(z-prevZ)/max(1,vecnorm(z))
        reld_alp = vecnorm(α-prevAlp)/max(1,vecnorm(α))
        
        ρ = ρ*1.2;
        L = chol(ρ*DD + (1+δ)*eye(n), Val{:L});
        
        @printf "%04d: %8.3e %8.3e %8.3e %8.3e %8.3e %8.3e %8.3e\n" k primal infeas ρ reld_x reld_z reld_alp pd_gap
    end
   
    prevX, prevZ, prevAlp = x, z, α
    
end

xim = reshape(x, picx, picy);
recovered = Image(map(RGB, xim, xim, xim)')