
% CODE FOR FIG S1

addpath('../')




% 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])

% PDSI ONLY, ENSEMBLE RECONSTRUCTION
efl='cesm_lme010_r12000_p998_state_q_avg_JunAug_prxydtst_5_prxtp_tca_2896_swtchbld100_06-Aug-2018_17:32:52.mat';
pth='./output-recon/';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;


%======================
%======  PHYDA  =======
%======================

k=find(state_tp=='q');
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;


% Latitude weighting
A=cosd(repmat(xmeta{k}.lat,[1 length(xmeta{k}.lon)]));

% PDSI ensemble
Xa_di_ens=zeros(reconYrs,sub_ens);
for i=1:sub_ens
   Xvar=squeeze(Xa_ens(xmeta{k}.id_X(1):xmeta{k}.id_X(2),i,:));
   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);
   Xv_msk=bsxfun(@times,Xpdsi,nasw);
   Xa_di_ens(:,i)=wmean_a(Xv_msk,A); % RECONSTRUCTION DROUGHT INDEX
end

% Get the ensemble mean; identical to the original reconstruction
Xvar=squeeze(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);
Xv_msk=bsxfun(@times,Xpdsi,nasw);
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,:);

% Ensemble estimate of drought periods
dsens=struct();
for i=1:sub_ens
   dsens.ds{i}=mdrought(Xa_di_ens(ir4,i),1,'y');
end

% Find which years there is agreement
mds=zeros(length(ir4),1);
for i=1:sub_ens
   for j=1:length(dsens.ds{i})
      m0=zeros(length(ir4),1);
      m0(dsens.ds{i}(j,1):dsens.ds{i}(j,2))=1;
      mds=mds+m0; % add up droughts
   end
end

figure
plot(analyrs4,mds,'o')
hold on
plot(analyrs4,90*ones(length(analyrs4),1),'r','linewidth',2)
title('Ensemble agreement about drought year')
ylabel('Percent agreement','fontsize',16)
xlabel('Year','fontsize',16)
xlim([analyrs4(1) analyrs4(end)])
box on
set(gca,'fontsize',16)

%print(['figs/ens_drght_agreement.png'],'-dpng','-r300');


% Find years of 90% agreement
%dyrs=find(mds>=sub_ens*0.9);
dyrs=find(mds>=sub_ens*0.5);


% NOW FIGURE OUT WHICH YEARS ARE CONTINUOUS CHUNKS
drghts=struct([]); % initialize structure
negpos=ones(length(ir4),1);
negpos(dyrs)=0;
% Find the timing of droughts
i=1;j=1;flag=0;
while i+1 <= length(ir4)
% Find the start of a drought
if negpos(i)==0
   d_start=i;
   i=i+1;
   % Keep going to the end of a drought
   while negpos(i)~=1
       i=i+1;
       if i+1 > length(ir4)
	   d_end=i-1;
	   drghts(j).strtend=[d_start d_end];
	   flag=1;
	   break
       end
   end
   if flag==1;break;end % get out of nested loops
   
   d_end=i-1;
   drghts(j).strtend=[d_start d_end];
   
   j=j+1;
end

i=i+1;
end
   
% Pull out years from structure
d=zeros(length(drghts),2);
for kk=1:length(drghts)
  d(kk,1:2)=drghts(kk).strtend;
end

% Throw out droughts that aren't at least 10 years long
rmv=[];
ii=1;
for jj=1:length(drghts)
   if (d(jj,2)-d(jj,1)) < 10
      rmv(ii)=jj;
      ii=ii+1;
   end
end
d(rmv,:)=[];  

% FOR RIGHT NOW, JUST USE ALL DROUGHTS; NEED TO SORT OUT LATER...
ds4=d;
svmd=d;

%save('svmd_ens_50pagree.mat','svmd','Xa_di','Xa_di_ens')


%======================
%======  NADA  ========
%======================

da=permute(ncread('/d1/nsteiger/drought-atlas/nada_hd2_cl.nc','pdsi'),[2 3 1]);
%nada(nada<-99)=NaN; % put nans in fill spots
lat_da=ncread('/d1/nsteiger/drought-atlas/nada_hd2_cl.nc','lat');
lon_da=ncread('/d1/nsteiger/drought-atlas/nada_hd2_cl.nc','lon');
yrs_da=ncread('/d1/nsteiger/drought-atlas/nada_hd2_cl.nc','time');


% Coarsen NADA to match climate model grid
ndm=load('nada_mask.mat');
% create grids for interpolation
nlns=find(lon_da<0);lon_da(nlns)=lon_da(nlns)+360;
[xn,yn]=meshgrid(lon_da,lat_da);
[xc,yc]=meshgrid(lon,lat);
% note masks of both data types
nm=find(~isnan(da(:,:,end)));
cm=find(ndm.sftlf==1);
da(da<-99)=NaN; % put nans in fill spots but after determining mask

da_i=nan(length(lat),length(lon),size(da,3));
for i=1:size(da,3)
   % linear interpolation and extrapolation for irregular data
   da1=da(:,:,i);
   F = scatteredInterpolant(xn(nm),yn(nm),da1(nm),'linear','linear');
   vq=F(xc(cm),yc(cm)); % do the interpolation
   % expand out the interpolated data points onto the climate model mask
   vd=nan(length(lat),length(lon));
   vd(cm)=vq; % full interpolated field
   da_i(:,:,i)=vd;
end



% 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;
A=cosd(repmat(xmeta{k}.lat,[1 length(xmeta{k}.lon)]));

DA_msk=bsxfun(@times,da_i,nasw);
DA_di=wmean_a(DA_msk,A); % DROUGHT ATLAS DROUGHT INDEX
% Indices of the drought atlas
analyrs5=1870:2000;
[~,~,id1] = intersect(analyrs1,yrs_da); % indices of recon segment
[~,~,id2] = intersect(analyrs2,yrs_da); % indices of recon segment
[~,~,id3] = intersect(analyrs3,yrs_da); % indices of recon segment
[~,~,id4] = intersect(analyrs4,yrs_da); % indices of recon segment
[~,~,id5] = intersect(analyrs5,yrs_da); % indices of recon segment

[~,~,ir5] = intersect(analyrs5,r_o:r_f); % indices of recon segment


% MEGADROUGHTS IN NADA	
nada_ds4=mdrought(DA_di(ir4),1,'y'); % 'y' indicates that the mean is removed

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


load distinct_colors_20.mat

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


% NADA AND HIGHLIGHT MEGADROUGHTS
figure
hold on
plot(analyrs3,DA_di(id3),'linewidth',1,'color',clrs.grey)
%plot(analyrs2,Xa_di(ir2),'linewidth',1,'color',clrs.red)
plot(analyrs3,smooth(DA_di(id3),11),'linewidth',2,'color',clrs.red)
%plot(analyrs2,smooth(Xa_di(ir2),10,'lowess'),'linewidth',2,'color',clrs.red)
plot(analyrs3,zeros(length(analyrs3),1),'--','linewidth',2,'color',clrs.yellow)
set(gca,'fontsize',16)
%ylim([-0.8 1])
xlim([analyrs3(1) analyrs3(end)])
xlabel('Year','fontsize',18)
ylabel(['NASW PDSI'],'fontsize',18)
box on
fig = gcf;fig.PaperUnits = 'inches';
fig.PaperPosition = [0 2 7.5 4.5]; % [left bottom width height]

y=Xa_di(ir3); % using 'l' from above
yMin=min(y)-0.5;yMax=max(y)+0.5;
ylim([yMin yMax])
% Add a patches for top five times from truth
for i=1:length(svmd)
    p=patch(analyrs3(1)+[ds4(i,1) ds4(i,2) ds4(i,2) ds4(i,1)],[yMin yMin yMax yMax],clrs.blue);
    p.FaceAlpha=0.35;p.EdgeColor='none';
end

legend({'NADA','NADA-11'},'location','southeast')
hold off

%print(['figs/drght_cmp_nada.png'],'-dpng','-r300');


% PHYDA AND HIGHLIGHT MEGADROUGHTS
figure
hold on
%plot(analyrs3,DA_di(id3),'linewidth',1,'color',clrs.grey)
plot(analyrs3,Xa_di(ir3),'linewidth',1,'color',clrs.grey)
%plot(analyrs3,smooth(DA_di(id3),11),'linewidth',2,'color',clrs.red)
plot(analyrs3,smooth(Xa_di(ir3),11),'linewidth',2,'color',clrs.purple)
plot(analyrs3,zeros(length(analyrs3),1),'--','linewidth',2,'color',clrs.yellow)
set(gca,'fontsize',16)
%ylim([-0.8 1])
xlim([analyrs3(1) analyrs3(end)])
xlabel('Year','fontsize',18)
ylabel(['NASW PDSI'],'fontsize',18)
box on
fig = gcf;fig.PaperUnits = 'inches';
fig.PaperPosition = [0 2 7.5 4.5]; % [left bottom width height]

%y=Xa_di(ir3); % using 'l' from above
%yMin=min(y)-0.5;yMax=max(y)+0.5;
ylim([yMin yMax])
% Add a patches for top five times from truth
for i=1:length(svmd)
    p=patch(analyrs3(1)+[ds4(i,1) ds4(i,2) ds4(i,2) ds4(i,1)],[yMin yMin yMax yMax],clrs.blue);
    p.FaceAlpha=0.35;p.EdgeColor='none';
end

legend({'PHYDA','PHYDA-11'},'location','southeast')
hold off


%print(['figs/drght_cmp_phyda.png'],'-dpng','-r300');

% PHYDA + NADA AND HIGHLIGHT MEGADROUGHTS
figure
hold on
plot(analyrs3,prctile(Xa_di_ens(ir3,:)',5),'linewidth',1,'color',clrs.grey)
h3=plot(analyrs3,prctile(Xa_di_ens(ir3,:)',95),'linewidth',1,'color',clrs.grey);
h1=plot(analyrs3,DA_di(id3),'linewidth',1.5,'color',clrs.red);
h2=plot(analyrs3,Xa_di(ir3),'linewidth',1.5,'color',clrs.purple);
%plot(analyrs3,smooth(DA_di(id3),11),'linewidth',2,'color',clrs.red)
%plot(analyrs3,smooth(Xa_di(ir3),11),'linewidth',2,'color',clrs.purple)
plot(analyrs3,zeros(length(analyrs3),1),'--','linewidth',2,'color',clrs.yellow)
%legend({'NADA','PHYDA'},'location','southeast')
%legend({'PHYDA','NADA'},'location','southeast')
legend([h1,h2,h3],{'NADA','PHYDA','5^{th}-95^{th} pctl'},'location','southeast')
set(gca,'fontsize',16)
%ylim([-0.8 1])
xlim([analyrs3(1) analyrs3(end)])
xlabel('Year','fontsize',18)
ylabel(['NASW PDSI'],'fontsize',18)
box on
fig = gcf;fig.PaperUnits = 'inches';
fig.PaperPosition = [0 2 7.5 4.5]; % [left bottom width height]

% y=DA_di(id2); % using 'l' from above
% yMin=min(y)-0.5;yMax=max(y)+0.5;
ylim([yMin yMax])
% % Add a patches for top five times from truth
% for i=1:9
%     p=patch(analyrs3(1)+[ds4(i,1) ds4(i,2) ds4(i,2) ds4(i,1)],[yMin yMin yMax yMax],clrs.blue);
%     p.FaceAlpha=0.40;p.EdgeColor='none';
% end
% 
% hold off

%print(['figs/drght_cmp_annual.png'],'-dpng','-r300');


% PHYDA + NADA 
figure
hold on
plot(analyrs3,smooth(prctile(Xa_di_ens(ir3,:)',5),11),'linewidth',1,'color',clrs.grey)
h3=plot(analyrs3,smooth(prctile(Xa_di_ens(ir3,:)',95),11),'linewidth',1,'color',clrs.grey);
h2=plot(analyrs3,smooth(Xa_di(ir3),11),'linewidth',1.5,'color',clrs.purple);
h1=plot(analyrs3,smooth(DA_di(id3),11),'linewidth',1.5,'color',clrs.red);
%plot(analyrs3,smooth(DA_di(id3),11),'linewidth',2,'color',clrs.red)
%plot(analyrs3,smooth(Xa_di(ir3),11),'linewidth',2,'color',clrs.purple)
plot(analyrs3,zeros(length(analyrs3),1),'--','linewidth',2,'color',clrs.yellow)
%legend({'NADA','PHYDA'},'location','southeast')
legend([h1,h2,h3],{'NADA-11','PHYDA-11','5^{th}-95^{th} pctl'},'location','southeast')
set(gca,'fontsize',16)
%ylim([-0.8 1])
xlim([analyrs3(1) analyrs3(end)])
xlabel('Year','fontsize',18)
ylabel(['NASW PDSI'],'fontsize',18)
box on
fig = gcf;fig.PaperUnits = 'inches';
fig.PaperPosition = [0 2 7.5 4.5]; % [left bottom width height]

% y=DA_di(id2); % using 'l' from above
% yMin=min(y)-0.5;yMax=max(y)+0.5;
ylim([yMin yMax])

%print(['figs/drght_cmp_dcdl.png'],'-dpng','-r300');






% % NADA AND HIGHLIGHT MEGADROUGHTS
% figure
% hold on
% %plot(analyrs3,DA_di(id3),'linewidth',1,'color',clrs.grey)
% %plot(analyrs3,Xa_di(ir3),'linewidth',1,'color',clrs.red)
% plot(analyrs3,smooth(DA_di(id3),11),'linewidth',2,'color',clrs.grey)
% plot(analyrs3,smooth(Xa_di(ir3),11),'linewidth',2,'color',clrs.purple)
% plot(analyrs3,zeros(length(analyrs3),1),'--','linewidth',2,'color',clrs.yellow)
% %legend('NADA','PHYDA','location','southeast')
% set(gca,'fontsize',16)
% %ylim([-0.8 1])
% xlim([analyrs3(1) analyrs3(end)])
% xlabel('Year','fontsize',18)
% ylabel(['NASW PDSI'],'fontsize',18)
% box on
% fig = gcf;fig.PaperUnits = 'inches';
% fig.PaperPosition = [0 2 7.5 4.5]; % [left bottom width height]
% 
% y=DA_di(id2); % using 'l' from above
% yMin=min(y)-0.5;yMax=max(y)+0.5;
% ylim([yMin yMax])
% % Add a patches for top five times from truth
% for i=1:9
%     p=patch(analyrs3(1)+[ds4(i,1) ds4(i,2) ds4(i,2) ds4(i,1)],[yMin yMin yMax yMax],clrs.blue);
%     p.FaceAlpha=0.40;p.EdgeColor='none';
% end
% 
% hold off






