%----------------------------------------
% MAKE CORAL AND TREE RING PSEUDOPROXIES
%----------------------------------------
% Updates: Can handle ensemble members 11-13

addpath('../')
addpath('../regem-imputation/')
addpath('../vs-lite/')

% Do trees and/or corals?
tpsm='y';
cpsm='y';

mdl_prior='cesm_lme007';
disp(['Model = ' mdl_prior]);
p_yrs=1:1155; % full length of the data is 1156

timestp=datestr(now);
timestp=strrep(timestp,' ','_'); % remove space


if cpsm=='y'

   disp('Making corals...')

   %------------------------------
   % CORAL PSM parameters
   %------------------------------

   % Load data for getting the parameters of the coral PSMs
   load('pages2k_proxy_corals.mat','p_idx','proxy','proxy_names','plat','plon')
   coral_names=proxy_names(p_idx);
   c_lat=plat;
   c_lon=plon;

   B=zeros(3,length(proxy));
   for i=1:length(proxy)
      B(:,i)=proxy{i}.b;
   end

   % Hand correct proxy 8 to be median of others
   B(:,8)=median(B,2);

   %clearvars -except coral_names c_lat c_lon B


   %======================================
   %           CORAL PSM	
   %======================================

   %mon_avg_o=1;mon_avg_f=12;
   %mon_avg_o=12;mon_avg_f=2;
   mon_avg_o=4;mon_avg_f=3;
   obs_dtp1='soda_sst';% sea surface temperature
   obs_dtp2='soda_sss';% sea surface salinity

   % LOAD *MONTHLY* MODEL DATA
   [M_sst,~,~,~,olat,olon] = load_Smon_ex(mdl_prior,'s',mon_avg_o,mon_avg_f,p_yrs);
   [M_sss] = load_Smon_ex(mdl_prior,'y',mon_avg_o,mon_avg_f,p_yrs);


   % Locations of proxies in the model
   % mask out land lat/lon values
   mask=sum(~isnan(M_sst),2);
   mask(mask==0)=NaN;
   mask(~isnan(mask))=1;

   [pl_m] = nearest_latlon(olat,olon,c_lat,c_lon,mask);

   % Pull out model data at proxy locations
   M_sstpmon=M_sst(pl_m,:);
   M_ssspmon=M_sss(pl_m,:);
   clear M_sst M_sss

   % Take the annual mean
   if max(diff(p_yrs)>1);error('Annualizing SSS and SST data only applies for continuous, non-random-in-time prior');end
   M_sstbc0=mon2ann(M_sstpmon,mon_avg_o,mon_avg_f);
   M_sssbc0=mon2ann(M_ssspmon,mon_avg_o,mon_avg_f);

   % LOAD ANNUAL OBSERVATIONAL DATA
   [O_sst,xlat,xlon] = load_obs_S(obs_dtp1,mon_avg_o,mon_avg_f);
   [O_sss] = load_obs_S(obs_dtp2,mon_avg_o,mon_avg_f);

   % Get locations of proxies in obs
   % mask out land lat/lon values
   mask=sum(~isnan(O_sst),2);
   mask(mask==0)=NaN;
   mask(~isnan(mask))=1;
   [pl_o] = nearest_latlon(xlat,xlon,c_lat,c_lon,mask);


   % disp('Mean correction, no annual variance correction directly in the BL PSM...')
   M_sstbc1=bsxfun(@minus,M_sstbc0,nanmean(M_sstbc0,2));
   M_sstbc=bsxfun(@plus,M_sstbc1,nanmean(O_sst(pl_o,90:end),2));
   M_sssbc1=bsxfun(@minus,M_sssbc0,nanmean(M_sssbc0,2));
   M_sssbc=bsxfun(@plus,M_sssbc1,nanmean(O_sss(pl_o,90:end),2));


   %-------------------------------------------
   % APPLY LINEAR FIT TO MODEL DATA AND OUTPUT
   %-------------------------------------------


   coral_psms0=zeros(length(coral_names),length(p_yrs));
   for jj=1:length(coral_names)
       coral_psms0(jj,:)=B(1,jj) + B(2,jj)*M_sstbc(jj,:) + B(3,jj)*M_sssbc(jj,:);
   end

   % Add observational error 0.1 for corals
   coral_psms=coral_psms0+0.1*randn(size(coral_psms0));

   coral_snr=zeros(length(coral_names),1);
   for i=1:length(coral_names)
      cr1=corr(coral_psms(i,:)',M_sstbc(i,:)');
      coral_snr(i)=abs(cr1)/sqrt(1-cr1^2);
   end

   %=============================================
   %          SAVE PSMS
   %=============================================

   %save(['PSM_corals_JanDec_' mdl_prior '.mat'],'c_lat','c_lon','coral_names','coral_psms','c_lat','c_lon','coral_snr')
   %save(['PSM_corals_DecFeb_' mdl_prior '.mat'],'c_lat','c_lon','coral_names','coral_psms','c_lat','c_lon','coral_snr')
   save(['PSM_corals_AprMar_' mdl_prior '_' timestp '.mat'],'c_lat','c_lon','coral_names','coral_psms','c_lat','c_lon','coral_snr','coral_psms0','M_sstbc')

end


if tpsm=='y'

   %------------------------------
   disp('Making trees...')
   %------------------------------

   %clb_prd=[1901 2000];
   %[T1,T2,M1,M2,MSE,vsl_cr,trw_names]=get_vslite_params(clb_prd);
   %[T1,T2,M1,M2,MSE,vsl_cr,trw_names,plat,plon,archv_tp]=get_vslite_params(clb_prd);

   %disp('Saving TRW parameters...')
   %load('vs_lite_historical_params_2Feb2018.mat','M1','M2','T1','T2')
   load('vsl_hist_params_expanded_9Feb2018.mat','M1','M2','T1','T2','plat','plon','trw_names')

   t_lat=plat;t_lon=plon;


   %======================================
   %           TREE RING PSM	
   %======================================

   mon_avg_o=1;mon_avg_f=12;
   %obs_dtp1='soda_sst';% sea surface temperature
   %obs_dtp2='soda_sss';% sea surface salinity

   % LOAD *MONTHLY* MODEL DATA
   [M_t,mlat,mlon] = load_Smon_ex(mdl_prior,'t',mon_avg_o,mon_avg_f,p_yrs);
   [M_p] = load_Smon_ex(mdl_prior,'p',mon_avg_o,mon_avg_f,p_yrs);

   % Get the proxy locations in model space
   [pl_m] = nearest_latlon(mlat,mlon,t_lat,t_lon);

   %---------------------------
   % BIAS CORRECT MODEL OUTPUT
   %---------------------------

   % Load monthly observational data
   obs_dtp1='cru_ts3_tmp';% temperature
   [Od_t,xlat,xlon,x_yrs] = load_obs_S(obs_dtp1,[],[]);

   obs_dtp1='cru_ts3_pre';% precipitation
   [Od_p] = load_obs_S(obs_dtp1,[],[]);

   % Get locations of proxies in obs
   mask=ones(size(Od_t(:,end))); % mask out ocean lat/lon values
   mask(isnan(Od_t(:,end)))=NaN;
   [pl_o] = nearest_latlon(xlat,xlon,t_lat,t_lon,mask);

   % Bias-correct temperature
   % infill random missing values in obs data with regem
   disp('---Using 1950:1999 for bias correction---')
   c2_yrs=1950:1999;[~,io2] = intersect(x_yrs,c2_yrs);

   options.maxit=10; options.regress='iridge'; %options.regress='ttls';
   %Od_tf=regem(Od_t(pl_o,:),options); % bias correct based on all data
   Od_tf=regem(Od_t(pl_o,io2),options); % bias correct on recent data
   [M_tbc] = bias_correct(Od_tf,M_t(pl_m,:),'n');

   % Bias-correct precipitation
   %odp=Od_p(pl_o,:); % all precip data
   odp=Od_p(pl_o,io2); % just modern data
   mdp=double(M_p(pl_m,:));
   odp(odp==0)=0.01; % replace zero exactly with value less than error in obs
   mdp(mdp<1e-13)=1e-13; % replace a dozen or so really low values
   % infill missing values with regem
   options.maxit=10; options.regress='iridge'; %options.regress='ttls';
   odp_f=regem(odp,options);
   odp_f(odp_f<0)=0.01; % make sure there aren't any zeros
   [M_pbc] = bias_correct(odp_f,mdp,'g');

   %-------------
   % RUN VS-LITE
   %-------------

   % Reshape data
   k=1;l=12;
   t_cal=zeros(length(t_lat),12,length(p_yrs));
   p_cal=zeros(length(t_lat),12,length(p_yrs));
   for n=1:length(p_yrs)
       t_cal(:,:,n)=M_tbc(:,k:l);
       p_cal(:,:,n)=M_pbc(:,k:l);
       k=k+12; l=l+12;
   end

   M_trw=zeros(length(t_lat),length(p_yrs));
   %gT=zeros(pNum,12,length(p_yrs)); % growth from T and M
   %gM=zeros(pNum,12,length(p_yrs));%,gT(q,:,:),gM(q,:,:)
   for q=1:length(t_lat)
       [M_trw(q,:)] = VSLite_v2_3(p_yrs(1),p_yrs(end),plat(q),T1(q),T2(q),...
	   M1(q),M2(q),squeeze(t_cal(q,:,:)),squeeze(p_cal(q,:,:)),...
	   'lbparams',[0.76,0.01,0.093,4.886,5.80,1000,0.2,1]);
   end

   % Throw out locations that are too cold to model
   tr_nanids=find(isnan(M_trw(:,2))); % 2nd year of all

   % Add observational error 0.1 for corals
   tree_psms=M_trw+0.01*randn(size(M_trw));

   M_tann=mon2ann(M_t(pl_m,:),mon_avg_o,mon_avg_f);

   tree_snr=zeros(length(trw_names),1);
   for i=1:length(trw_names)
      cr1=corr(tree_psms(i,:)',M_tann(i,:)');
      tree_snr(i)=abs(cr1)/sqrt(1-cr1^2);
   end



   %=============================================
   %          SAVE PSMS
   %=============================================

   tree_psms0=M_trw;
   save(['PSM_trees_' mdl_prior '_v0_Jul2020.mat'],'t_lat','t_lon','trw_names','tr_nanids','tree_psms','tree_snr','tree_psms0','M_tann')


end







