%  This MATLAB script builds forecast and analysis ensembles
%  every "dT" time steps
%  with data collected every "dS" locations
%  (the user is prompted to input dT and dS)
%
%     fore1     forecast ensembles for ENKF
%     forec     forecast ensembles for cENKF
%     anly1     analysis ensembles for ENKF
%     anlyc     analysis ensembles for cENKF
%
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%  RELIES ON:
%     INIT.mat as generated by go.m
%	COV_init.m         (implicit through go.m) 
%	ENSEMBLE0_init.m   (implicit through go.m)
%     	PARAMS_init.m      (implicit through go.m)
%       enkf.m             (explicit)
%	forward_1d.m       (explicit)
%	   odetrap.m       (implicit through forward_1d)
%		nsolve.m   (implicit through forward_1d)
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%  Load information previously stored by running "go.m"
global PARAMS TRU;
load INIT;

% Save dimension sizes 
Nx  = size(TRU,1); 
Nt  = size(TRU,2);
Nx2 = Nx/2;

dT = input('Input number of iterations between Analysis Cycles (dT = 10)? ');
if(isempty(dT)); dT = 10; end;

dS =input('Input number of spatial locations between Data Points (dS = 10)? ');
if(isempty(dS)); dS = 10; end;

tme  = dT:dT:Nt;
spc1 = round(dS/2):dS:Nx2;
spc2 = [spc1, Nx2+spc1];
ns   = length(spc2);
ns2  = length(spc1);
tru  = TRU(:,tme);

% Build Observation Function ("H" in the paper)
HHH = eye(Nx); 
HHH = HHH(spc2,:);

%  Generate Data ("y" at each analysis time point
%  spatial locations are   msh(spc1)
%  Every TEMP data point has corresponding SUPPLY data point
dat = HHH*( tru + Sigc*randn(size(tru)));

% Deriviative Smoother Matrix ("A" in the Paper)
AAA = eye(Nx2-1,Nx)-triu(ones(Nx2-1,Nx),1)+triu(ones(Nx2-1,Nx),2);

% Ensembles
% Prior ensemble "ens0" is built in "go.m"
fore1 = repmat(ens0,[1,1,length(tme)]);  % the forecast for  EnkF
forec = fore1;                           % the forecast for cEnKF
anly1 = fore1;                           % the analysis for  EnKF
anlyc = fore1;                           % the analysis for cEnKF


warning off
tic; 
for j = 1:length(tme);
  for i = 1:Nens;
    % forward propigate each ensemble member
    dum = forward_1d(squeeze(anly1(:,i,j)),dT,0);
    % these become the forecasts
    fore1(:,i,j) = dum(:,end);
    % repeat for the Constrained Ensemble Kalman Filter
    dum = forward_1d(squeeze(anlyc(:,i,j)),dT,0);
    forec(:,i,j) = dum(:,end);
  end
  % Do the EnKF analysis update
  dum = enkf(dat(:,j),HHH,HHH*Sig,fore1(:,:,j));
  anly1(:,:,j+1) = dum;

  % Do the cEnKF analysis update
  tmp = mean(AAA*dum,2);
  ss2 = std(AAA*dum,1,2).^2;
  anlyc(:,:,j+1) = enkf(tmp,AAA,30*ones(Nx2-1,1),dum);

  % Plots to show progress
  disp([j,round(toc)])
  plot(msh(1:Nx2),anly1(1:Nx2,:,j+1),'b--',msh(1:Nx2),anlyc(1:Nx2,:,j+1),'r')
  hold on; plot(msh(1:Nx2),tru(1:Nx2,j),'y')
  plot(msh(spc1),dat(1:ns2,j),'k^'); 
  title(sprintf('TME: %d ',tme(j)*dT)); hold off;
  pause(.001);

end
warning on;
% save SCENE1


disp('Watch the movie? ')
dum = input('0 = No, t>0 is number of seconds between the 10 frames  ');
if(dum);
for j = 1:length(tme);
 plot(msh,anly1(1:Nx2,1,j+1),'b--',msh,anlyc(1:Nx2,1,j+1),'r')
 hold on; plot(msh,tru(1:Nx2,j),'y',msh(spc1),dat(1:ns2,j),'g^');
 plot(msh,anly1(1:Nx2,:,j+1),'b--',msh,anlyc(1:Nx2,:,j+1),'r')
 plot(msh,tru(1:Nx2,j),'y',msh(spc1),dat(1:ns2,j),'g^')

 axis([0,1,100,2500])
 hold off;
 legend('ENKF only','Smoothed Ens.','Truth','Data')
title(sprintf('Time = %1.3f',PARAMS.dt*tme(j)));
pause(dum)
end
disp('Done')
end