clear all;
close all;
clc;

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

rng(1);
ns = 55;
xs = linspace(0,1,ns)'; 
s2eps = 1e-9; %0.1; % 0.25; %
s2 = 1; t = 0.5;
mu = 0.0;
R = exp(-0.5*pdist2(xs,xs).^2/(t^2));
Sigma = s2*R + s2eps*eye(ns);
[TSigma,p] = chol(Sigma);
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)'; 
s = s2*exp(-0.5*pdist2(xs,xp).^2/(t^2));

F = ones(ns,1);
x = TSigma'\fs;
M = TSigma'\F;
b = (M'*M)\(M'*x);  % b = (F'*Sigma^(-1)*F)^(-1)*F'*Sigma^(-1)*fs;
z = x - M * b;
zz = z'*z; % zz = (fs-F*b)'*Sigma^(-1)*(fs-F*b);

f = ones(np,1);
trend = f*b;
complement = (TSigma'\s)' * z; % complement = s'*Sigma^(-1)*(fs-F*b);
fp = trend + complement;

s21 = sum((TSigma'\s) .* (TSigma'\s), 1); % compute s'*Sigma^(-1)*s 
T_M = chol(M' * M);   % equivalently : qrR <- qr.R(qr(M))
m = (T_M)' \ (f - (TSigma'\s)' * M)';
s22 = sum(m .* m,1);
s2p_sk= s2 - s21';
s2p_uk= s2 - 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 = - sum(log(abs(diag(TSigma)))) - 0.5*ns*log(2*pi) - 0.5*zz;

% varying parameters t and s2
clear logLikt;
tt = 0.05:0.001:0.15;
s2st = 0.5:0.1:3.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 + s2eps*eye(ns);
        [TSigmat,p] = chol(Sigmat);
        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;
        zzt = zt'*zt; % zz = (fs-F*b)'*Sigma^(-1)*(fs-F*b);
        logLikt(i,j) =  - sum(log(abs(diag(TSigmat)))) - 0.5*ns*log(2*pi) - 0.5*zzt;
    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 + s2eps*eye(ns);
[TSigmat,p] = chol(Sigmat);
st = s2st2*exp(-0.5*pdist2(xs,xp).^2/(tt2^2));
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;
zzt = zt'*zt; % zz = (fs-F*b)'*Sigma^(-1)*(fs-F*b);
lmax =  - sum(log(abs(diag(TSigmat)))) - 0.5*ns*log(2*pi) - 0.5*zzt;


trend = f*bt;
complement = (TSigmat'\st)' * zt; % complement = s'*Sigma^(-1)*(fs-F*b);
fpt2 = trend + complement;

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

q95 = norminv(0.975);
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)])

%
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,fpt2);
color = 'g';
edgecolor ='k';
transparency = 0.2;
fill([xp; flipud(xp)], [lower95t2_sk; flipud(upper95t2_sk)], color, 'EdgeColor',edgecolor,'FaceAlpha',transparency,'EdgeAlpha',transparency);