% Add toolboxes 
addpath ../InverseModel
addpath ../MiscUtil
clear; %close all

% Set model run parameters
N=100;       % number of time to run the model
I=15; J=15;  % domain size
% set initial conditions
To=zeros(I,J);
To(11,5)=5;
To(6,8)=5;

% run model to obtain true field
[T_true, y_coord,x_coord]= im_adv_diff_model(To,N);


% make a figure of T0 and T100
figure(1); clf;
subplot(1,2,1)
utl_contourfill(x_coord,y_coord,To,30); 
colormap(utl_getpmap(98)); colorbar
title 'Initial Condition (T=0)'
subplot(1,2,2)
utl_contourfill(x_coord,y_coord,T_true,30); 
colormap(utl_getpmap(98)); colorbar
title 'Final Condition (T=100)'

% convert output into a vector
T_true=T_true(:); 

%----------------------------------------------------------------------------
% 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 = 5000;  
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;

figure(1);
subplot(1,2,1)
utl_contourfill(x_coord,y_coord, reshape(T_obs, [I J]) ,30); 
colormap(utl_getpmap(98)); colorbar
title (['Final Condition + Measurment Error (sn=',num2str(signal_noise),')'])

%==========================================================
%	how to build E,  for the model y = E*x + n
% NOTE: we will build E by perturbing the model 
% initial conditions with impulses over 1 grid point
%==========================================================
[I,J]=size(x_coord);
load E
if ~exist('E')
   n=0;
   E = zeros(I*J, I*J);
   for j=1:J
      for i=1:I
         n=n+1
         To_pert=zeros(I,J,1);
         % perturb an inital patch of 1 grid points
         To_pert(i,j)=1; 
	     T100= im_adv_diff_model(To_pert,N);
         % save the column of E associated with the patch
         E(:,n)=T100(:);
      end
   end
end

% verify that E is correct 
T_verify=E*To(:);
figure(1);
subplot(1,2,1)
utl_contourfill(x_coord,y_coord, reshape(T_verify, [I J]) ,30); 
colormap(utl_getpmap(98)); colorbar
title (['Final Condition from E derived numerically'])

%----------------------------------------------------------------------------
% Begin LSQ Procedure
case_ids ={'normal' 'weighted' 'tapered'};
case_id=3;

% 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

%==========================================================
% build W weights to constrain model-data misfit
%==========================================================

W = 1/(rms_err^2);

%==========================================================
% build S weights to constrain model parameters
%==========================================================

% size of sources (assume about 10 x bigger than the maximum [])
rms_x = 0.5;                    

S=1/(rms_x^2); 


% additional flags to account for no_boundary sources
% and for knwon location of the sources. (set to 1 to activate)
no_boundary=0; known_sources =0;
S_default=S; 
n=0; clear S
for i=1:I
   for j=1:J
      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
%==========================================================

% redefine weigths according to type of LSQ selected
if strcmp(case_ids{case_id},'normal')
    W=1; S=0;
end
if strcmp(case_ids{case_id},'weighted')
    S=0;
end


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


% once you find your estimates of the model parameters x
% you can remap these on the model intial conditions with
% the reverse loop
To_hat=reshape(xhat,[I J]);

% run model with new initial conditons
T_hat = im_adv_diff_model(To_hat,N);
T_hat=T_hat(:);



% make a figure of T0 and T100
figure(2); clf;
subplot(2,2,1)
utl_contourfill(x_coord,y_coord,To,30); 
colormap(utl_getpmap(98)); colorbar
title 'Initial Condition (T=0)'
subplot(2,2,2)
utl_contourfill(x_coord,y_coord,reshape(T_true, [I J]),30); 
colormap(utl_getpmap(98)); colorbar
title 'Final Condition (T=100)'

% make a figure of T0_hat and T100_hat
subplot(2,2,3)
utl_contourfill(x_coord,y_coord,reshape(To_hat, [I J]),30); 
colormap(utl_getpmap(98)); colorbar
title 'Initial Condition estimate (T=0)'
subplot(2,2,4)
utl_contourfill(x_coord,y_coord, reshape(T_hat, [I J]),30); 
colormap(utl_getpmap(98)); colorbar
title 'Final Condition estimate (T=100)'

return

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

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

To=zeros(I,J);
To(11,5)=5;
To(6,8)=5;

figure(1); clf
colormap(utl_getpmap(98));
subplot(2,2,1)
utl_contourfill(x_coord,y_coord,To,30); 
colormap(utl_getpmap(98)); colorbar

subplot(2,2,3)
rnt_contourfill(reshape(T_obs,[30 30])',30); colorbar ; 


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


figure(2)
colormap(getpmap(98));
subplot(1,2,2)
rnt_contourfill(reshape(T_obs,[30 30])',30); colorbar ; shading flat
title 'TRUE + measurment error'
subplot(1,2,1)
rnt_contourfill(reshape(T_true,[30 30])',30); colorbar ; shading flat
title 'TRUE dispersion'
rnt_font
lprps fig3.eps