clear all;
close all;
clc;

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

xl = -1; xu = 1;
np = 100; % number of points for plot
xp = linspace(xl,xu,np)'; % points forplot
% kernel parameters
s2eps = 1e-9; %0.1; % 0.25; %
s2 = 1; l = 0.5;

%% first case: 20 samples
t = 'first case';
rng(1);
ns = 20;
js = sort(randperm(np,ns));
xs = xp(js); 
mu = 0.0;
R = correlationfunc(l,xs,xs);
Sigma = s2*R + s2eps*eye(ns);
[TSigma,p] = chol(Sigma);
fs = mu + TSigma' * randn(ns,1); % ns points randomly generated from GP

lb = (mu - 2*sqrt(s2))*ones(np,1); % lower bound prior confidence interval
ub = (mu + 2*sqrt(s2))*ones(np,1); % upper bound prior confidence interval

% posterior
[fp,s2p_sk,lower95_sk,upper95_sk] = posterior(s2,l,ns,xs,np,xp,fs,TSigma);

% plot
plot_prior_and_posterior(xs,fs,lb,ub,xp,fp,lower95_sk,upper95_sk,t)

clear ll ss2 p1 p2 llogLik;
% ll = 0.47:0.001:0.49;
% ss2 = 0.65:0.001:0.75;
ll = 0.01:0.01:0.9;
ss2 = 0.01:0.01:0.9;
[p1,p2] = meshgrid(ll,ss2);
for i=1:size(p1,1)
    for j=1:size(p2,2)
        llogLik(i,j) = logLik(s2eps,p2(i,j),p1(i,j),ns,xs,fs);
    end
end
llogLikmax = max(llogLik,[],'all')
[imax,jmax]=find(llogLik==llogLikmax);
lmax = p1(imax,jmax);
s2max = p2(imax,jmax);

figure;
hold on;
contourf(p1,p2,llogLik,50)
plot(p1(imax,jmax),p2(imax,jmax),'o')
line([p1(imax,jmax) p1(imax,jmax)],[min(ss2),max(ss2)])
line([min(ll),max(ll)],[p2(imax,jmax) p2(imax,jmax)])
title(t)
figure;
hold on
surf(p1,p2,llogLik);
plot3(p1(imax,jmax),p2(imax,jmax),llogLikmax,'o');
xlabel('l');
ylabel('s2');
zlabel('loglikelihood');
title(t)
view(45,45)

% posterior
Rmax = correlationfunc(lmax,xs,xs);
Sigmamax = s2max*Rmax + s2eps*eye(ns);
[TSigmamax,p] = chol(Sigmamax);
[fpmax,s2p_skmax,lower95_skmax,upper95_skmax] = posterior(s2max,lmax,ns,xs,np,xp,fs,TSigmamax);
% plot
plot_prior_and_posterior(xs,fs,lb,ub,xp,fpmax,lower95_skmax,upper95_skmax,t)


%% second case: 10 samples subset
t = 'second case';
js = sort(randperm(ns,10));
ns = 10;
xs = xs(js); 
fs = fs(js); 
R = correlationfunc(l,xs,xs);
Sigma = s2*R + s2eps*eye(ns);
[TSigma,p] = chol(Sigma);

lb = (mu - 2*sqrt(s2))*ones(np,1); % lower bound prior confidence interval
ub = (mu + 2*sqrt(s2))*ones(np,1); % upper bound prior confidence interval

% posterior
[fp,s2p_sk,lower95_sk,upper95_sk] = posterior(s2,l,ns,xs,np,xp,fs,TSigma);

% plot
plot_prior_and_posterior(xs,fs,lb,ub,xp,fp,lower95_sk,upper95_sk,t)

clear ll ss2 p1 p2 llogLik;
% ll = 0.44:0.001:0.46;
% ss2 = 0.60:0.001:0.70;
ll = 0.01:0.01:0.9;
ss2 = 0.01:0.01:0.9;
[p1,p2] = meshgrid(ll,ss2);
for i=1:size(p1,1)
    for j=1:size(p2,2)
        llogLik(i,j) = logLik(s2eps,p2(i,j),p1(i,j),ns,xs,fs);
    end
end
llogLikmax = max(llogLik,[],'all')
[imax,jmax]=find(llogLik==llogLikmax);
lmax = p1(imax,jmax);
s2max = p2(imax,jmax);

figure;
hold on;
contourf(p1,p2,llogLik,50)
plot(p1(imax,jmax),p2(imax,jmax),'o')
line([p1(imax,jmax) p1(imax,jmax)],[min(ss2),max(ss2)])
line([min(ll),max(ll)],[p2(imax,jmax) p2(imax,jmax)])
title(t)
figure;
hold on
surf(p1,p2,llogLik);
plot3(p1(imax,jmax),p2(imax,jmax),llogLikmax,'o');
xlabel('l');
ylabel('s2');
zlabel('loglikelihood');
title(t)
view(45,45)

% posterior
Rmax = correlationfunc(lmax,xs,xs);
Sigmamax = s2max*Rmax + s2eps*eye(ns);
[TSigmamax,p] = chol(Sigmamax);
[fpmax,s2p_skmax,lower95_skmax,upper95_skmax] = posterior(s2max,lmax,ns,xs,np,xp,fs,TSigmamax);
% plot
plot_prior_and_posterior(xs,fs,lb,ub,xp,fpmax,lower95_skmax,upper95_skmax,t)


%% third case: 5 samples subset
t = 'third case';
js = sort(randperm(ns,5));
ns = 5;
xs = xs(js); 
fs = fs(js); 
R = correlationfunc(l,xs,xs);
Sigma = s2*R + s2eps*eye(ns);
[TSigma,p] = chol(Sigma);

lb = (mu - 2*sqrt(s2))*ones(np,1); % lower bound prior confidence interval
ub = (mu + 2*sqrt(s2))*ones(np,1); % upper bound prior confidence interval

% posterior
[fp,s2p_sk,lower95_sk,upper95_sk] = posterior(s2,l,ns,xs,np,xp,fs,TSigma);

% plot
plot_prior_and_posterior(xs,fs,lb,ub,xp,fp,lower95_sk,upper95_sk,t)

clear ll ss2 p1 p2 llogLik;
% ll = 0.15:0.001:0.22;
% ss2 = 0.07:0.001:0.15;
ll = 0.01:0.01:0.9;
ss2 = 0.01:0.01:0.9;
[p1,p2] = meshgrid(ll,ss2);
for i=1:size(p1,1)
    for j=1:size(p2,2)
        llogLik(i,j) = logLik(s2eps,p2(i,j),p1(i,j),ns,xs,fs);
    end
end
llogLikmax = max(llogLik,[],'all')
[imax,jmax]=find(llogLik==llogLikmax);
lmax = p1(imax,jmax);
s2max = p2(imax,jmax);

figure;
hold on;
contourf(p1,p2,llogLik,50)
plot(p1(imax,jmax),p2(imax,jmax),'o')
line([p1(imax,jmax) p1(imax,jmax)],[min(ss2),max(ss2)])
line([min(ll),max(ll)],[p2(imax,jmax) p2(imax,jmax)])
title(t)
figure;
hold on
surf(p1,p2,llogLik);
plot3(p1(imax,jmax),p2(imax,jmax),llogLikmax,'o');
xlabel('l');
ylabel('s2');
zlabel('loglikelihood');
title(t)
view(45,45)

% posterior
Rmax = correlationfunc(lmax,xs,xs);
Sigmamax = s2max*Rmax + s2eps*eye(ns);
[TSigmamax,p] = chol(Sigmamax);
[fpmax,s2p_skmax,lower95_skmax,upper95_skmax] = posterior(s2max,lmax,ns,xs,np,xp,fs,TSigmamax);
% plot
plot_prior_and_posterior(xs,fs,lb,ub,xp,fpmax,lower95_skmax,upper95_skmax,t)