% DROUGHT ANALYSIS FIGURES

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% LOAD RECONSTRUCTION DATA FIRST!!!!
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

addpath('../')

% ALL PROXY TYPES, LMR DATABASE
%efl='cesm_lme010_r12000_p998_state_t2qjozn_avg_JunAug_prxydtst_5_prxtp_tca_2896_swtchbld100_05-Jan-2018_14:49:45.mat';
efl='cesm_lme010_r12000_p998_state_t2qjozn_avg_AprMar_prxydtst_5_prxtp_tca_2896_swtchbld100_05-Jan-2018_14:33:07.mat';
% INDEX ONLY RECONSTRUCTION WITH FULL ENSEMBLES
%efl='cesm_lme010_r12000_p998_state_2ozn_avg_AprMar_prxydtst_5_prxtp_tca_2605_swtchbld100_16-Nov-2017_15:57:59.mat';


%pth='/d2/nsteiger/output-da/hydroclimate/'; load([pth,efl])
pth='/d2/nsteiger/output-da/hydroclimate/subrecons_5Jan18/';load([pth,efl])


%varpth=['/d6/haibo/lme/cam5/BLMTRC5CN/ALLFORCEINGS/010'];
%sftlf=permute(ncread([varpth '/QSOIL/b.e11.BLMTRC5CN.f19_g16.010.clm2.h0.QSOIL.085001-184912.nc'],'landmask'),[2 1]);
%lat=ncread([varpth '/QSOIL/b.e11.BLMTRC5CN.f19_g16.010.clm2.h0.QSOIL.085001-184912.nc'],'lat');
%lon=ncread([varpth '/QSOIL/b.e11.BLMTRC5CN.f19_g16.010.clm2.h0.QSOIL.085001-184912.nc'],'lon');



%-------------------------
% LOAD SPATIAL FIELDS
%-------------------------

% Temperature
k=find(state_tp=='t');
Xvar=Xa_m(xmeta{k}.id_X(1):xmeta{k}.id_X(2),:);
Xt2m=reshape(Xvar,length(xmeta{k}.lat),length(xmeta{k}.lon),reconYrs);


clearvars -except xa_amo xb_itcz_sa xb_itcz_atl xb_itcz_pac xa_itcz_sa xa_itcz_atl xa_itcz_pac Xt2m xa_nino xb_nino


% Drought atlases use JJA instead of annual
efl='cesm_lme010_r12000_p998_state_t2qjozn_avg_JunAug_prxydtst_5_prxtp_tca_2896_swtchbld100_05-Jan-2018_14:49:45.mat';
pth='/d2/nsteiger/output-da/hydroclimate/subrecons_5Jan18/';load([pth,efl])
%pth='/d2/nsteiger/output-da/hydroclimate/'; load([pth,efl])


% What years to look at?
analyrs1=800:1600;
analyrs2=1601:1925;
analyrs3=800:2000;
analyrs4=800:1925;
[~,~,ir1] = intersect(analyrs1,r_o:r_f); % indices of recon segment
[~,~,ir2] = intersect(analyrs2,r_o:r_f); % indices of recon segment
[~,~,ir3] = intersect(analyrs3,r_o:r_f); % indices of recon segment
[~,~,ir4] = intersect(analyrs4,r_o:r_f); % indices of recon segment
% monthly indices
ovrlprng=rng_intr([r_o,r_f],[analyrs1(1),analyrs1(end)]);
mon_r=r_o:(1/12):(r_f+11/12);
irm1_a=find(abs(ovrlprng(1)-mon_r)<1e-6);
irm1_b=find(abs(ovrlprng(2)-1/12-mon_r)<1e-6)+12;

ovrlprng=rng_intr([r_o,r_f],[analyrs2(1),analyrs2(end)]);
mon_r=r_o:(1/12):(r_f+11/12);
irm2_a=find(abs(ovrlprng(1)-mon_r)<1e-6);
irm2_b=find(abs(ovrlprng(2)-1/12-mon_r)<1e-6)+12;

ovrlprng=rng_intr([r_o,r_f],[analyrs3(1),analyrs3(end)]);
mon_r=r_o:(1/12):(r_f+11/12);
irm3_a=find(abs(ovrlprng(1)-mon_r)<1e-6);
irm3_b=find(abs(ovrlprng(2)-1/12-mon_r)<1e-6)+12;

% Expand out PDSI land-only fields from reconstruction
% NEED TO MAKE 'LNDIDX' AVAILABLE IN THE OUTPUT FILE!!! MAYBE LOAD IT AS A GLOBAL VARIABLE SOMEWHERE? OR JUST LOAD IT IN PARAMS SCRIPT?
%load(['/d1/nsteiger/climate-data/cesm-lme/cesm_lme_010_pdsi_output_biascorr_AWC_c_7_u2_gcm.mat'],'lndidx')
%varpth=['/d6/haibo/lme/cam5/BLMTRC5CN/ALLFORCEINGS/010'];
%sftlf=permute(ncread([varpth '/QSOIL/b.e11.BLMTRC5CN.f19_g16.010.clm2.h0.QSOIL.085001-184912.nc'],'landmask'),[2 1]);

% PDSI
k=find(state_tp=='q');
Xvar=Xa_m(xmeta{k}.id_X(1):xmeta{k}.id_X(2),:);
Xpdsi=NaN(length(xmeta{k}.lat)*length(xmeta{k}.lon),reconYrs);
Xpdsi(xmeta{k}.lndidx,:)=Xvar;
Xpdsi=reshape(Xpdsi,length(xmeta{k}.lat),length(xmeta{k}.lon),reconYrs);

lat=xmeta{k}.lat;
lon=xmeta{k}.lon;



% PULL OUT SPATIAL MEAN OVER AMERICAN SOUTHWEST
%load('/d1/nsteiger/climate-data/cesm-lme/Nrth_A_Swthwst_mask.mat')
varpth=['/d6/haibo/lme/cam5/BLMTRC5CN/ALLFORCEINGS/010'];
sftlf=permute(ncread([varpth '/QSOIL/b.e11.BLMTRC5CN.f19_g16.010.clm2.h0.QSOIL.085001-184912.nc'],'landmask'),[2 1]);
%lat=ncread([varpth '/QSOIL/b.e11.BLMTRC5CN.f19_g16.010.clm2.h0.QSOIL.085001-184912.nc'],'lat');
%lon=ncread([varpth '/QSOIL/b.e11.BLMTRC5CN.f19_g16.010.clm2.h0.QSOIL.085001-184912.nc'],'lon');
% NASW = box of 235:255 E (125:105 W); 31:42 N (just the US, not Mexico)
% Find nearest lat/lon equal to or outside bounds
lonext=[235 255];latext=[31 42];
ln_a=find((abs(lon-lonext(1))-min(abs(lon-lonext(1))))<1e-5);
ln_b=find((abs(lon-lonext(2))-min(abs(lon-lonext(2))))<1e-5);
lt_a=find((abs(lat-latext(1))-min(abs(lat-latext(1))))<1e-5);
lt_b=find((abs(lat-latext(2))-min(abs(lat-latext(2))))<1e-5);
nasw=zeros(size(sftlf));
nasw(lt_a:lt_b,ln_a:ln_b)=sftlf(lt_a:lt_b,ln_a:ln_b);

nasw(nasw==0)=NaN;
Xv_msk=bsxfun(@times,Xpdsi,nasw);
A=cosd(repmat(xmeta{k}.lat,[1 length(xmeta{k}.lon)]));
Xa_di=wmean_a(Xv_msk,A); % RECONSTRUCTION DROUGHT INDEX

% How to compute the megadroughts? 1 => 11 year running mean, 2 => 2S2E
%ds1=mdrought(Xa_di(ir1),1,'y'); % 'y' indicates that the mean is removed
%ds2=mdrought(Xa_di(ir1),2,'y');

% How to compute the megadroughts? 1 => 11 year running mean, 2 => 2S2E
ds4=mdrought(Xa_di(ir4),1,'y'); % 'y' indicates that the mean is removed
% Most severe droughts have severity < -10 and all occur before 1600 (last one ends in 1590)
si=find(ds4(:,3)<-10);
%si=find(ds4(:,2)+analyrs1(1)<analyrs1(end));
svmd=ds4(si,:);

%===========================================================================

load distinct_colors_20.mat

% Analysis options
plot_p='y'; % plot proxies?
prntfg='n'; % Save figures?

% Load full ensembles for the nino and amo variables
efl='cesm_lme010_r12000_p998_state_2ozn_avg_AprMar_prxydtst_5_prxtp_tca_2896_swtchbld100_05-Jan-2018_14:24:35.mat';
pth='/d2/nsteiger/output-da/hydroclimate/subrecons_5Jan18/';
load([pth,efl],'state_tp','xmeta','Xa_ens','reconYrs','p_yrs')


k=find(state_tp=='o');
xa_amo_ens=permute(squeeze(Xa_ens(xmeta{k}.id_X(1),:,:)),[2 1]);

k=find(state_tp=='2');
xa_gmt_ens=permute(squeeze(Xa_ens(xmeta{k}.id_X(1),:,:)),[2 1]);

k=find(state_tp=='n'); % specific nino variable
if strcmp(xmeta{k}.indinfo{3}.indxnm,'Nino 3.4'); i=3; else; error('Wrong Nino specification'); end
xa_nino_ens=reshape(permute(Xa_ens((xmeta{k}.id_X(1)+(i-1)*12):(xmeta{k}.id_X(1)+i*12-1),:,:),[1 3 2]),12*reconYrs,length(p_yrs));

% Annualize Nino
xa_nino_ann=mon2ann(xa_nino_ens',4,3)';
xa_nino_ann(end+1,:)=xa_nino_ann(end,:); % repeat end year because of annualizing



% Get an offset value for anomalies
gmt_off=mean(mean(xa_gmt_ens(ir2,:),2));

% Get an offset value for anomalies
amo_off=mean(mean(xa_amo_ens(ir2,:),2));

% Get an offset value for anomalies
nino_off=mean(mean(xa_nino_ens(irm2_a:irm2_b,:),2));

%==================================================================================================r
%                        REGRESSION MODELS (LINEAR & LOGISTIC)
%===================================================================================================

ds=svmd;
ird_ann=[];
for i=1:size(ds,1)
  [~,~,irab] = intersect((ds(i,1)+analyrs1(1)-1):(ds(i,2)+analyrs1(1)-1),r_o:r_f); 
  % Collect all years of drought
  ird_ann=cat(1,ird_ann,irab);
end

% Bivariate logistic regression
X(:,1)=mean(xa_nino_ann(ir4,:)-nino_off,2);
X(:,2)=mean(xa_amo_ens(ir4,:)-amo_off,2);


y0=zeros(reconYrs,1);
y0(ird_ann)=1;
Y=categorical(y0(ir4));
[B,dev,stats] = mnrfit(X,Y);

% Bivariate logistic regression with interaction
%Xi(:,1)=mean(xa_nino_ann(ir4,:)-nino_off,2);
%Xi(:,2)=mean(xa_amo_ens(ir4,:)-amo_off,2);
%Xi(:,3)=Xi(:,1).*Xi(:,2);
%[Bi,devi,statsi] = mnrfit(Xi,Y);



% Explore the predicted probabilities using the range of ENSO and AMO

% Define a grid of all possible Nino and AMO values
%ns=xa_nino_ann(ir4,:)-nino_off;as=xa_amo_ens(ir4,:)-amo_off;
% Just pick the mean values over which the model is built
ns=mean(xa_nino_ann(ir4,:)-nino_off,2);as=mean(xa_amo_ens(ir4,:)-amo_off,2);
[mx,my] = meshgrid(min(ns(:)):0.01:max(ns(:)),min(as(:)):0.01:max(as(:)));
Xm(:,1)=mx(:);
Xm(:,2)=my(:);

% Get probability for random observations
pihat = mnrval(B,Xm);
[pihat,dlow,dhi] = mnrval(B,Xm,stats);% get uncertainties
% NOTE THAT UPPER AND LOWER BOUNDS ARE THE SAME!!


% Reshape to match the grid
pM=reshape(pihat(:,2),size(mx));
pL=reshape(dlow(:,2),size(mx));
%pU=reshape(dhi(:,2),size(mx));

% figure
% [Cc,h]=contour(mx,my,pM,'ShowText','on','fill','on','linecolor','r');
% colormap(flip(gray))
% clabel(Cc,h,'FontSize',14,'Color','red')
% title('Predicted megadrought probability','fontsize',20)
% set(gca,'fontsize',16)
% xlabel('Nino3.4 (annual anom. deg. C)','fontsize',18)
% ylabel('AMO (anom. deg. C)','fontsize',18)
% 
% if prntfg=='y';print(['./figs/logreg_mdrought_prob_nino_amo.png'],'-dpng','-r300');end
% 
% 
% figure
% [Cc,h]=contour(mx,my,pL,'ShowText','on','linecolor','k');
% clabel(Cc,h,'FontSize',14,'Color','black')
% title('Probability 95% conf.','fontsize',20)
% set(gca,'fontsize',16)
% xlabel('Nino3.4 (annual anom. deg. C)','fontsize',18)
% ylabel('AMO (anom. deg. C)','fontsize',18)
% 
% if prntfg=='y';print(['./figs/logreg_mdrought_prob_nino_amo_unc.png'],'-dpng','-r300');end


% Alternative plot with both probability and uncertainty
figure
hold on
[C1,h1]=contour(mx,my,pM,'fill','on');
set(gca,'tickdir','both')
colormap(flip(gray))
c=colorbar;
c.Label.String = 'Est. megadrought probability';
c.Label.FontSize = 16;
[Cc,h]=contour(mx,my,pL,'showtext','on','linecolor',clrs.red,'linestyle','-.');
clabel(Cc,h,'Color',clrs.red)
set(gca,'fontsize',16)
xlabel('Nino3.4 (annual anom. deg. C)','fontsize',18)
ylabel('AMO (anom. deg. C)','fontsize',18)
box on

if prntfg=='y';print(['./figs/logreg_mdrought_prob_nino_amo_data&unc.png'],'-dpng','-r300');end


% %-------------------------
% % TEST GMT INSTEAD OF AMO
% %-------------------------
% 
% clear X Xm Xi
% 
% % Bivariate logistic regression
% X(:,1)=mean(xa_nino_ann(ir4,:)-nino_off,2);
% X(:,2)=mean(xa_gmt_ens(ir4,:)-gmt_off,2);
% 
% y0=zeros(reconYrs,1);
% y0(ird_ann)=1;
% Y=categorical(y0(ir4));
% [B,dev,stats] = mnrfit(X,Y);
% 
% % Bivariate logistic regression with interaction
% Xi(:,1)=mean(xa_nino_ann(ir4,:)-nino_off,2);
% Xi(:,2)=mean(xa_gmt_ens(ir4,:)-gmt_off,2);
% Xi(:,3)=Xi(:,1).*Xi(:,2);
% [Bi,devi,statsi] = mnrfit(Xi,Y);
% 
% 
% 
% % Explore the predicted probabilities using the range of ENSO and AMO
% 
% % Define a grid of all possible Nino and AMO values
% %ns=xa_nino_ann(ir4,:)-nino_off;as=xa_gmt_ens(ir4,:)-gmt_off;
% % Just pick the mean values over which the model is built
% ns=mean(xa_nino_ann(ir4,:)-nino_off,2);as=mean(xa_gmt_ens(ir4,:)-gmt_off,2);
% [mx,my] = meshgrid(min(ns(:)):0.01:max(ns(:)),min(as(:)):0.01:max(as(:)));
% Xm(:,1)=mx(:);
% Xm(:,2)=my(:);
% 
% % Get probability for random observations
% pihat = mnrval(B,Xm);
% [pihat,dlow,dhi] = mnrval(B,Xm,stats);% get uncertainties
% % NOTE THAT UPPER AND LOWER BOUNDS ARE THE SAME!!
% 
% % Reshape to match the grid
% pM=reshape(pihat(:,2),size(mx));
% pL=reshape(dlow(:,2),size(mx));
% 
% 
% % Alternative plot with both probability and uncertainty
% figure
% hold on
% [C1,h1]=contour(mx,my,pM,'fill','on');
% set(gca,'tickdir','both')
% colormap(flip(gray))
% c=colorbar;
% c.Label.String = 'Est. megadrought probability';
% c.Label.FontSize = 16;
% [Cc,h]=contour(mx,my,pL,'showtext','on','linecolor',clrs.red,'linestyle','-.');
% clabel(Cc,h,'Color',clrs.red)
% set(gca,'fontsize',16)
% xlabel('Nino3.4 (annual anom. deg. C)','fontsize',18)
% ylabel('GMT (anom. deg. C)','fontsize',18)
% box on
% 
% if prntfg=='y';print(['./figs/logreg_mdrought_prob_nino_gmt.png'],'-dpng','-r300');end
% 
% 


% TRY TO FIX THE AUTOCORRELATIONS






