%% First create an ensemble of random fields
Y = zeros(10000,1000);
Z = zeros(10000,1000);
figure(1);subplot(121);
for i = 1:1000,
    [z,x] = randfield(0.25,100);
    Y(:,i) = x(:);
    Z(:,i) = z(:);
    if (rem(i,10)==0)
        imagesc(x,[-2,2]);colorbar;drawnow;
    end
end
subplot(122);
imagesc(reshape(mean(Y,2),[100 100]),[-2 2]);colorbar;drawnow; title('Ensemble Mean');

%% Training and testing split
TrainY = Y(:,1:500);
TestY = Y(:,501:1000);
TestZ = Z(:,501:1000);

% Calculate the distance between pairs of vectors. These are "sphered"
O1 = TestZ(:,1:250);O2 = TestZ(:,251:500);
eO = O1 - O2;
oerror = trace(eO'*eO)/10000/250


%% Reduced model tests
% calculate mean

my = mean(TrainY,2); %% this must be near zero
MY = repmat(my,[1,500]);
DY = TrainY-MY;

% The SVD of DY yields the eigenvectors/values of the covariance
[u,s,v] = svd(DY,0);

figure(2); subplot(121);plot(diag(s));

% Check reconstruction error.
for i = 1:250,
    ss = diag(s(1:i,1:i))/sqrt(i-1); %% Keep the top i
    IS = diag(1./(ss)); % produces ixi diagonal matrix.
    U = u(:,1:i); % Keep i singular vectors
    Red=(IS*(U'*TestY)); %Reduction
    
    e1 = Red(:,1:250); e2 = Red(:,251:500);
    e = e1 - e2;
    
    reconerror(i) = oerror-trace(e'*e)/i/250;
    
end
% http://www.stat.cmu.edu/~cshalizi/350/lectures/13/lecture-13.pdf
% QUOTE
% 1. How many do you need? Of course, there is
% nothing which says that accuracy wouldn?t drop like a stone as soon as q
% < p (the optimal number of components, ED). It?s always possible that the very smallest component, carrying the
% least variance, is also the only one carrying any information about what
% we actually want to make inferences about.

figure(2); subplot(122); plot(1:250,reconerror);

%% Ensemble Filtering example

Truth = reshape(TestY(:,end),[100 100]);
SparseMeasure = 1:5:100;
[xc,yc]=meshgrid(SparseMeasure,SparseMeasure);
figure(3); subplot(121);
imagesc(Truth);title('Truth');

stdobs = 0.1;
Obs = Truth+randn(size(Truth))*stdobs;
figure(3); subplot(122);imagesc(Obs);hold on;plot(xc,yc,'k.'); hold off; title('Observations');

obs = Obs(SparseMeasure,SparseMeasure);
obs = obs(:);
R = stdobs*stdobs*eye(length(SparseMeasure)^2);

Ensemble = reshape(TrainY(:,1:250),[100 100 250]);
EnsObs = Ensemble(SparseMeasure,SparseMeasure,:);

EnsObs = reshape(EnsObs,[length(SparseMeasure)^2, 250]);
Ensemble = reshape(Ensemble,[10000 250]);

EnsObsMean = sum(EnsObs,2);
DEnsObs = EnsObs - repmat(EnsObsMean,[1 250]);
[u,s,v]  = svd(DEnsObs,0);
s= diag(diag(s));

mix = pinv(s.*s/stdobs^2+eye(size(s)))/stdobs^2; %% This is a diagonal matrix!

X1 = u'*(repmat(obs,[1 250]) - EnsObs);
X2 = mix*X1;
X3 = u*X2;
X4 = DEnsObs'*X3;
X5 = eye(size(X4))+X4;
AnaEns = Ensemble*X5;
AnaEns = reshape(AnaEns,[100 100 250]);
Ensemble = reshape(Ensemble,[100 100 250]);
figure(4);

for i = 1:250,
    subplot(121); imagesc(Ensemble(:,:,i));
    subplot(122);imagesc(AnaEns(:,:,i));drawnow;
end
subplot(121); imagesc(mean(Ensemble,3));colorbar;title('Prior Mean');
subplot(122); imagesc(mean(AnaEns,3));colorbar;title('Posterior Mean');
%% Filtering again, but with a part observed
Truth = reshape(TestY(:,end),[100 100]);
SparseMeasure = 51:5:100; [xc,yc]=meshgrid(SparseMeasure,SparseMeasure);
figure(5); subplot(121); imagesc(Truth);

stdobs = 0.1; Obs = Truth+randn(size(Truth))*stdobs;
subplot(122); imagesc(Obs);hold on;plot(xc,yc,'k.'); hold off;

obs = Obs(SparseMeasure,SparseMeasure); obs = obs(:);
R = stdobs*stdobs*eye(length(SparseMeasure)^2);
nens = 100;
Ensemble = reshape(TrainY(:,1:nens),[100 100 nens]);
EnsObs = Ensemble(SparseMeasure,SparseMeasure,:);

EnsObs = reshape(EnsObs,[length(SparseMeasure)^2, nens]);
Ensemble = reshape(Ensemble,[10000 nens]);

EnsObsMean = mean(EnsObs,2);
DEnsObs = EnsObs - repmat(EnsObsMean,[1 nens]);
[u,s,v]  = svd(DEnsObs,0);
s= diag(diag(s));

mix = pinv(s.*s/stdobs^2+eye(size(s)))/stdobs^2; %% This is a diagonal matrix!

X1 = u'*(repmat(obs,[1 nens]) - EnsObs);
X2 = mix*X1;
X3 = u*X2;
X4 = DEnsObs'*X3;
X5 = eye(size(X4))+X4;
AnaEns = Ensemble*X5;
AnaEns = reshape(AnaEns,[100 100 nens]);
Ensemble = reshape(Ensemble,[100 100 nens]);
figure(6);

for i = 1:nens,
    subplot(121); imagesc(Ensemble(:,:,i));
    subplot(122);imagesc(AnaEns(:,:,i));drawnow;
end

subplot(121); imagesc(mean(Ensemble,3));colorbar;title('Prior Mean');
subplot(122); imagesc(mean(AnaEns,3));colorbar;title('Posterior Mean');

%% Split Solutio

Truth = reshape(TestY(:,end),[100 100]);
node(1,1).truth = Truth(1:50,1:50); node(1,2).truth =Truth(1:50,51:100);
node(2,1).truth = Truth(51:100,1:50);   node(2,2).truth = Truth(51:100,51:100);
stdobs = 0.1;Obs = Truth+randn(size(Truth))*stdobs;
node(1,1).Obs = Obs(1:50,1:50); node(1,2).Obs = Obs(1:50,51:100);
node(2,1).Obs = Obs(51:100,1:50);   node(2,2).Obs = Obs(51:100,51:100);

SparseMeasure = 1:5:50;

for i = 1:2,
    for j = 1:2,
        figure(8);        subplot(2,2,(i-1)*2+j);
        [xc,yc]=meshgrid(SparseMeasure,SparseMeasure);
        imagesc(node(i,j).Obs);hold on;plot(xc,yc,'k.'); hold off;
        node(i,j).obs = node(i,j).Obs(SparseMeasure,SparseMeasure);
        node(i,j).obs = node(i,j).obs(:);
        node(i,j).R =  stdobs*stdobs*eye(length(SparseMeasure)^2);
    end
end
figure(7);
imagesc([node(1,1).Obs node(1,2).Obs; node(2,1).Obs node(2,2).Obs]);

Ens = reshape(TrainY(:,1:100),[100 100 100]);
node(1,1).ens = Ens(1:50,1:50,:);   node(1,2).ens = Ens(1:50,51:100,:);
node(2,1).ens = Ens(51:100,1:50,:); node(2,2).ens = Ens(51:100,51:100,:);

for i = 1:2
    for j = 1:2
        node(i,j).ensobs = node(i,j).ens(SparseMeasure,SparseMeasure,:);
        node(i,j).ensobsvec = reshape(node(i,j).ensobs,[length(SparseMeasure)^2,100]);
        node(i,j).ensvec = reshape(node(i,j).ens,[2500,100]);
        EnsObsMean = mean(node(i,j).ensobsvec,2);
        DEnsObs = node(i,j).ensobsvec - repmat(EnsObsMean,[1 100]);
        [u,s,v]  = svd(DEnsObs,0);
        s= diag(diag(s));
        size(s)
        mix = pinv(s.*s/stdobs^2+eye(size(s)))/stdobs^2; %% This is a diagonal matrix!
        X1 = u'*(repmat(node(i,j).obs,[1 100]) - node(i,j).ensobsvec);
        X2 = mix*X1;
        X3 = u*X2;
        X4= DEnsObs'*X3;
        X5 = eye(size(X4))+X4;
        node(i,j).anaensvec = node(i,j).ensvec*X5;
        node(i,j).anaens = reshape(node(i,j).anaensvec,[50 50 100]);
        node(i,j).X4 = X4;
        
    end
end
figure(10);

for k = 1:100,
    ana = [node(1,1).anaens(:,:,k) node(1,2).anaens(:,:,k);...
        node(2,1).anaens(:,:,k) node(2,2).anaens(:,:,k)];
    
    subplot(1,2,1); imagesc(ana);
    fore = [node(1,1).ens(:,:,k) node(1,2).ens(:,:,k);...
        node(2,1).ens(:,:,k) node(2,2).ens(:,:,k)];
    subplot(122); imagesc(fore);drawnow;
end
%% Message Passing
X4sum = node(1,1).X4 + node(1,2).X4 + node (2,1).X4 + node(2,2).X4;
X5 = eye(size(X4sum))+X4sum;
for i = 1:2
    for j = 1:2,
        node(i,j).anaensvec = node(i,j).ensvec*X5;
        node(i,j).anaens = reshape(node(i,j).anaensvec,[50 50 100]);
        
    end
end
