clear all;
close all;
clc;

% GP Z(x) with m=0; k(x1,x2)= s2*exp(-(x1-x2)^2/(2*t^2))
% s2=2, t=0.1
% we sample a realization of Z(x) on 10 points

rng(1);
ns = 40;
xs = linspace(0,1,ns)'; 
keps = 1e-9;
s2=2; t=0.1;
mu = 0.0;
R = exp(-0.5*pdist2(xs,xs).^2/(t^2));
Sigma = s2*R;
[TR,p] = chol(R + keps*eye(ns));
TSigma = sqrt(s2)*TR;
fs = mu + TSigma' * randn(ns,1);

lb = mu - 2*sqrt(diag(Sigma)) ;
ub = mu + 2*sqrt(diag(Sigma)) ;

% we plot the sampled realization of Z(x)
np = 100;
xp = linspace(0,1,np)'; 
r = exp(-0.5*pdist2(xs,xp).^2/(t^2));

F = ones(ns,1);
x = TR'\fs;
M = TR'\F;
b = (M'*M)\(M'*x);  % b = inv(F'*R^(-1)*F)*F'*iR*fs;
z = x - M * b;
s2s = z'*z/ns; % s2s = (fs-F*b)'*R^(-1)*(fs-F*b)/ns;

f = ones(np,1);
trend = f*b;
complement = (TR'\r)' * z; % complement = r'*R^(-1)*(fs-F*b);
fp = trend + complement;

s21 = sum((TR'\r) .* (TR'\r), 1); % compute r'*R^(-1)*r 
T_M = chol(M' * M);   % equivalently : qrR <- qr.R(qr(M))
m = (T_M)' \ (f - (TR'\r)' * M)';
s22 = sum(m .* m,1);
s2p_sk= s2s * (1 - s21');
s2p_uk= s2s * (1 - s21' + s22');

q95 = norminv(0.975);
lower95_sk = fp - q95*sqrt(s2p_sk);
upper95_sk = fp + q95*sqrt(s2p_sk);
lower95_uk = fp - q95*sqrt(s2p_uk);
upper95_uk = fp + q95*sqrt(s2p_uk);

figure;
hold on;
plot(xs,fs,'o');
color = 'y';
edgecolor ='k';
transparency = 0.2;
fill([xs; flipud(xs)], [lb; flipud(ub)], color, 'EdgeColor',edgecolor,'FaceAlpha',transparency,'EdgeAlpha',transparency);
plot(xp,fp);
color = 'b';
edgecolor ='k';
transparency = 0.2;
fill([xp; flipud(xp)], [lower95_sk; flipud(upper95_sk)], color, 'EdgeColor',edgecolor,'FaceAlpha',transparency,'EdgeAlpha',transparency);

%likelihood
logLik = -0.5*ns*s2s - sum(log(abs(diag(TR)))) - 0.5*ns*log(2*pi) - 0.5*ns;

% varying parameter t
tt=0.05:0.001:0.15;
for i=1:length(tt)
    Rt = exp(-0.5*pdist2(xs,xs).^2./(tt(i)^2));
    Sigmat = s2*Rt;
    [TRt,p] = chol(Rt + keps*eye(ns));
    xt = TRt'\fs;
    Mt = TRt'\F;
    bt = (Mt'*Mt)\(Mt'*xt);  % b = inv(F'*R^(-1)*F)*F'*iR*fs;
    zt = xt - Mt * bt;
    s2st(i) = zt'*zt/ns; % s2s = (fs-F*b)'*R^(-1)*(fs-F*b)/ns;
    logLikt(i) = -0.5*ns*s2st(i) - sum(log(abs(diag(TRt)))) - 0.5*ns*log(2*pi) - 0.5*ns;
end
[lmax,imax]=max(logLikt);
tt1=tt(imax);
s2st1=s2st(imax);
Rt = exp(-0.5*pdist2(xs,xs).^2./(tt1^2));
Sigmat = s2*Rt;
[TRt,p] = chol(Rt + keps*eye(ns));
xt = TRt'\fs;
Mt = TRt'\F;
bt = (Mt'*Mt)\(Mt'*xt);  % b = inv(F'*R^(-1)*F)*F'*iR*fs;
zt = xt - Mt * bt;
s2st1 = zt'*zt/ns; % s2s = (fs-F*b)'*R^(-1)*(fs-F*b)/ns;
lmax = -0.5*ns*s2st1 - sum(log(abs(diag(TRt)))) - 0.5*ns*log(2*pi) - 0.5*ns;
rt = exp(-0.5*pdist2(xs,xp).^2/(tt1^2));
trend = f*bt;
complement = (TRt'\rt)' * zt; % complement = r'*R^(-1)*(fs-F*b);
fpt1 = trend + complement;

s21 = sum((TRt'\rt) .* (TRt'\rt), 1); % compute r'*R^(-1)*r 
s2pt1_sk= s2st1 * (1 - s21');

lower95t1_sk = fpt1 - q95*sqrt(s2pt1_sk);
upper95t1_sk = fpt1 + q95*sqrt(s2pt1_sk);


figure;
hold on
plot(tt,logLikt)
plot(tt(imax),lmax,'o')
grid on;

% varying parameters t and s2
clear logLikt;
tt = 0.05:0.001:0.15;
s2st = 1:0.1:5;
[p1,p2] = meshgrid(tt,s2st);
for i=1:size(p1,1)
    for j=1:size(p1,2)
        Rt = exp(-0.5*pdist2(xs,xs).^2./(p1(i,j)^2));
        Sigmat = p2(i,j)*Rt;
        [TSigmat,p] = chol(Sigmat + keps*eye(ns));
        xt = TSigmat'\fs;
        Mt = TSigmat'\F;
        bt = (Mt'*Mt)\(Mt'*xt);  % b = inv(F'*S^(-1)*F)*F'*S^(-1)*fs;
        zt = xt - Mt * bt;
        logLikt(i,j) = -0.5* zt'*zt - sum(log(abs(diag(TSigmat)))) - 0.5*ns*log(2*pi) ;
    end
end
lmax = max(logLikt,[],'all')
[imax,jmax]=find(logLikt==lmax);
tt2=p1(imax,jmax);
s2st2=p2(imax,jmax);

Rt = exp(-0.5*pdist2(xs,xs).^2./(tt2^2));
Sigmat = s2st2*Rt;
[TSigmat,p] = chol(Sigmat + keps*eye(ns));
xt = TSigmat'\fs;
Mt = TSigmat'\F;
bt = (Mt'*Mt)\(Mt'*xt);  % b = inv(F'*S^(-1)*F)*F'*S^(-1)*fs;
zt = xt - Mt * bt;
lmax = -0.5* zt'*zt - sum(log(abs(diag(TSigmat)))) - 0.5*ns*log(2*pi) ;
st = s2st2*exp(-0.5*pdist2(xs,xp).^2/(tt2^2));
trend = f*bt;
complement = (TSigmat'\st)' * zt; % complement = r'*R^(-1)*(fs-F*b);
fpt2 = trend + complement;

s21 = sum((TSigmat'\st) .* (TSigmat'\st), 1); % compute r'*R^(-1)*r 
s2pt2_sk= s2st2 - s21';

lower95t2_sk = fpt2 - q95*sqrt(s2pt2_sk);
upper95t2_sk = fpt2 + q95*sqrt(s2pt2_sk);


figure;
hold on;
contourf(p1,p2,logLikt,50)
plot(p1(imax,jmax),p2(imax,jmax),'o')
line([p1(imax,jmax) p1(imax,jmax)],[min(s2st),max(s2st)])
line([min(tt),max(tt)],[p2(imax,jmax) p2(imax,jmax)])

% plot realization

figure;
hold on;
plot(xs,fs,'o');
plot(xp,fp);
color = 'b';
edgecolor ='k';
transparency = 0.2;
fill([xp; flipud(xp)], [lower95_sk; flipud(upper95_sk)], color, 'EdgeColor',edgecolor,'FaceAlpha',transparency,'EdgeAlpha',transparency);

plot(xp,fpt1,'k');
color = 'y';
edgecolor ='k';
transparency = 0.2;
fill([xp; flipud(xp)], [lower95t1_sk; flipud(upper95t1_sk)], color, 'EdgeColor',edgecolor,'FaceAlpha',transparency,'EdgeAlpha',transparency);

plot(xp,fpt2,'y');
color = 'g';
edgecolor ='k';
transparency = 0.2;
fill([xp; flipud(xp)], [lower95t2_sk; flipud(upper95t2_sk)], color, 'EdgeColor',edgecolor,'FaceAlpha',transparency,'EdgeAlpha',transparency);