clear; clearf
N=100;       % number of time to run the model
K=0.04;      % diffusion coefficient
I=30; J=30;  % domain size

% Velocity field is given below (DO NOT CHANGE)
u=zeros(I,J,1);
v=zeros(I,J,1);
[y_coord,x_coord]=meshgrid(1:I,1:J);
u_adv=1.6;
u=u_adv*cos(3*y_coord/J).*x_coord/(I/2)/1.5;
v=-1.2*u_adv*cos(2*14*x_coord/J);

% 
% 
% % load true dispersion pattern at t=100
source=1;
[T_true, To]=hw3_get_sources(source);
T_true = load(['pattern',num2str(source),'.dat']);
T_true_gridded = reshape(T_true, [30 30]);
figure(1); clf;
pcolor(x_coord, y_coord, T_true_gridded); shading flat; colorbar;
title 'TRUE DISPERSION PATTERN at t=100'

% add Guassian noise with zero mean and a given signal to noise ratio
% to the true dispersion data to generate some synthetic observations
signal_noise = 5;
rms_T = std(T_true);
rms_err = rms_T / signal_noise;
err = randn( length(T_true), 1);
err = err*rms_err/std(err);
T_obs  = T_true + err;


%return
% HINTS below
%----------------------------------------------------------------------------

% define cost function and weigths for least square problem
% J = (y - Ex)'*W*(y - Ex) + x'*S*x
% J = n'*W*n + x'*S*x
% where y = Ex + n   is my model and E is the adv-diff integral solution

%rms_err= rms_T/0.1;
W = 1/(rms_err^2);
rms_x = 0.1;
S=1/(rms_x^2);
no_boundary=1; known_sources =1;
%S=0;
%W=1;

%==========================================================
%	how to build E,  for the model y = E*x + n
% NOTE: becuase you know that sources occupy a square of
% 4 gridpoints, we will build E by perturbing the model
% initial conditions with impulses over 4 grid points
% square only.
%==========================================================
[I,J]=size(x_coord);
%load E
%if ~exist('E')
    n=0;
    E = zeros(I*J, I*J/4);
    for i=1:2:I-1
        for j=1:2:J-1
            n=n+1;
            To_pert=zeros(I,J,1);
            % perturb an inital patch of 4 gridpoints
            To_pert(i:i+1,j:j+1)=1;
            % run model
            T=adv_diff_model(To_pert,K,N,u,v);
            % get time = 100
            T100=T(:,:,100);
            % save the column of E associated with the patch
            E(:,n)=T100(:);
        end
    end
%end

%==========================================================
% build S weights for model parameters
%==========================================================

S_default=S;
n=0; clear S
for i=1:2:I-1
    for j=1:2:J-1
        n=n+1   ;
        S(n)=S_default;
        
        if no_boundary == 1
            if (i == 1 | i == I-1) |   (j == 1 | j == J-1)
                S(n) = 1/0.00001;
            end
            if (i == 3 | i == I-3) |   (j == 3 | j == J-3)
                S(n) = 1/0.00001;
            end
        end
        
        if known_sources == 1
            S(n) = 1/0.00001;
            if i==23 & j==21
                S(n) = 1/(rms_x*rms_x);
            end
            if i==9 & j==7
                S(n) = 1/(rms_x*rms_x);
            end
        end
    end
end
S= diag(S);

%==========================================================
% do inversion
% J = (y - Ex)'*W*(y - Ex) + x'*S*x
% J = n'*W*n + x'*S*x
% where y = Ex + n   is my model and E is the adv-diff integral solution
%==========================================================

y = T_obs;

CME = E'*W*E + S ;
[I1,J1]=size(CME);
cond = [1:I1]*0 + 1.0e-10;
%cond=0;
CME = CME + diag(cond);
xhat = CME \ E'*W*y;

% remap model parameters to the initial conditions of adv-diff model
n=0; To_hat=zeros(I,J,1);
for i=1:2:I-1
    for j=1:2:J-1
        n=n+1;
        To_hat(i:i+1,j:j+1)=xhat(n);
    end
end

% run model with new initial conditons
T = adv_diff_model(To_hat,K,N,u,v);
T_hat = T(:,:,100); T_hat = T_hat(:);


%posterior stats
n=T_obs-T_hat;
J = n'*W*n + xhat'*S*xhat;

n_std = std(n);
corr=corrcoef(T_obs,T_hat); corr=corr(2,1);


figure(1); clf
colormap(getpmap(98));
subplot(2,2,1)
pcolor(To'); colorbar ; shading flat
subplot(2,2,3)
pcolor(reshape(T_true,[30 30])'); colorbar ; shading flat


subplot(2,2,2)
pcolor(To_hat'); colorbar ; shading flat
title(['J = ',num2str(J),'  STD Nhat = ',num2str(n_std),'  assumed RMS = ',num2str(std(err)) ]);
subplot(2,2,4)
pcolor(reshape(T_hat,[30 30])'); colorbar ; shading flat
title(['Correlation between Yobs and Yhat ',num2str(corr*100)]);


figure(3)
colormap(getpmap(98));
subplot(1,2,2)
pcolor(reshape(T_obs,[30 30])'); colorbar ; shading flat
title 'TRUE + measurment error'
subplot(1,2,1)
pcolor(reshape(T_true,[30 30])'); colorbar ; shading flat
title 'TRUE dispersion'