%% TV examples.
% first show some TV functions and their values
nx = 256;
xr = [0:1/(nx-1):1]';
f1 = zeros(size(xr)); 
f1(2:nx-1) = 1/(nx-2) *[1:nx-2];f1(nx) = 1;
f2 = zeros(size(xr)); 
f2(floor(0.15*nx):floor(0.2*nx)) = 0.2; f2(ceil(0.2*nx):floor(0.35*nx)) = 0.6; f2(ceil(0.35*nx):end) = 1;
f3 = zeros(size(xr)); 
f3(floor(0.0625*nx):floor(0.33*nx)) = 0.8; f3(ceil(0.33*nx):end) = 1;
f4 = zeros(size(xr)); 
f4(floor(0.12*nx):floor(0.89*nx)) = 0.1; f4(ceil(0.89*nx):end) = 1;

figure(11); clf;
subplot(3,4,1); plot(xr,f1);
subplot(3,4,2); plot(xr,f2);
subplot(3,4,3); plot(xr,f3);
subplot(3,4,4); plot(xr,f4);

Dxn = spdiags([-1*ones(nx,1),ones(nx,1)],0:1,nx,nx)/nx;
Dxn0 = Dxn(1:nx-1,1:nx);

df1 = abs(Dxn0*f1);
df2 = abs(Dxn0*f2);
df3 = abs(Dxn0*f3);
df4 = abs(Dxn0*f4);

subplot(3,4,5); plot(xr(1:nx-1),df1);title(['TV : ',num2str(sum(df1))]);
subplot(3,4,6); plot(xr(1:nx-1),df2);title(['TV : ',num2str(sum(df2))]);
subplot(3,4,7); plot(xr(1:nx-1),df3);title(['TV : ',num2str(sum(df3))]);
subplot(3,4,8); plot(xr(1:nx-1),df4);title(['TV : ',num2str(sum(df4))]);


subplot(3,4,9); plot(xr(1:nx-1),df1.^2/2);title(['Tk1 : ',num2str(sum(df1.^2)/2)]);
subplot(3,4,10); plot(xr(1:nx-1),df2.^2/2);title(['Tk1 : ',num2str(sum(df2.^2)/2)]);
subplot(3,4,11); plot(xr(1:nx-1),df3.^2/2);title(['Tk1 : ',num2str(sum(df3.^2)/2)]);
subplot(3,4,12); plot(xr(1:nx-1),df4.^2/2);title(['Tk1 : ',num2str(sum(df4.^2)/2)]);

%% Do regularised inversion using TV and quadratic programming
%
%addpath('../Lin1D/');
sigma = 0.1; % std.dev. of added white noise

f = zeros(nx,1);
f(floor(0.187*nx):floor(0.2*nx)) = -1; f(floor(0.312*nx):floor(0.35*nx)) = 1.5; f(floor(0.5*nx):floor(0.6*nx)) = 2; % arbitrary function

[y,K] = linblur(f,0.04);
yn = y + sigma*randn(nx,1); % add white noise, 


figure (1);clf; 
subplot(1,2,1);hold on;
plot(y,'k');
plot(yn,'g');
title('Data Space');
subplot(1,2,2);hold on;
plot(f,'k');
title('Solution Space');

%% set up matrices for quadratic programming
% construct difference operator as explicit (sparse) matrix
%Dx = spdiags([-1*ones(n,1),ones(n,1)],0:1,n,n)/n;
%Dx0 = Dx(1:n-1,:);
onen = ones(nx,1);
szeron = sparse(nx,nx); % sparse n X n zero

% choose regln parameter and set up matrices
alphaTVQ = 8e-1;
H = [2*K'*K zeros(nx,2*nx) ; zeros(2*nx,3*nx)];
b = [-2*K'*yn; alphaTVQ*onen; alphaTVQ*onen];
B = sparse(3*nx,3*nx);
B(nx+1:3*nx,nx+1:3*nx) = -speye(2*nx);
Ceq = sparse(nx,3*nx);
Ceq(1:nx,1:nx) = Dxn; 
Ceq(1:nx,nx+1:2*nx) = -speye(nx);
Ceq(1:nx,2*nx+1:3*nx) = speye(nx);
%[Dx, sparse(n,2*n); szeron(n), -speye(n), szeron(n); sparse(n,2*n),speye(n)]; 

%% sanity checks;
Df = Dxn*f;
vp = zeros(size(f)); vm = vp;
dfp = find(Df>0);
dfm = find(Df<0);
vp(dfp) = Df(dfp);
vm(dfm) = -Df(dfm);

h = [f;vp;vm];
Ch = Ceq*h;


%% solve it...
x0 = [yn;zeros(2*nx,1)]; % initial step
[xqp,fval,exitflag,output,lambda]= quadprog(H,b,B,zeros(3*nx,1),Ceq,zeros(nx,1));
xfmincon = fmincon(@(x) (0.5*x'*H*x + b'*x),x0,B,zeros(3*nx,1),Ceq,zeros(nx,1));

figure(1);
subplot(1,2,2);
plot(xqp(1:nx),'--m');
plot(xfmincon(1:nx),'c');
subplot(1,2,1);
plot(K*xqp(1:nx),'--m');
plot(K*xfmincon(1:nx),'c');

%% include TK1 as a comparison
alphaTK1 = 5e2*alphaTVQ;
xTK1 = (K'*K + alphaTK1*Dxn'*Dxn)\(K'*yn);
figure(1);
subplot(1,2,2);
plot(xTK1,'b');
legend('f_{true}','f_{QP}','f_{fmincon}','f_{TK1}');

subplot(1,2,1);
plot(K*xTK1,'b');
legend('g_{0}','g^{obs}','g_{QP}','g_{fmincon}','g_{TK1}');

%%
function [y,K] = linblur(f,s)
n = size(f,1);
K = zeros(n,n);
h = 1/n;
C = h/sqrt(2*pi*s*s);
x = [1:n];
k = C*exp(-x.^2 * h*h/(2*s*s));
for i = 1:n
    K(i,1:i) = k(i:-1:1);
    K(i,i+1:n) = k(2:n-i+1);
end
y = K*f;
end