
%==========================================================================
% THIS IS THE DRIVING SCRIPT FOR PDSI CALCULATIONS
% Nathan Steiger, Apr 2016
%==========================================================================


varpth='/d1/nsteiger/climate-data/gfdl-esm2m/';

lat=ncread([varpth 'tas_Amon_GFDL-ESM2M_historical_r1i1p1_186101-200512.nc'],'lat');
lon=ncread([varpth 'tas_Amon_GFDL-ESM2M_historical_r1i1p1_186101-200512.nc'],'lon');

nlat=repmat(lat,length(lon),1);
nlon=reshape(repmat(lon,1,length(lat))',length(lon)*length(lat),1);

syr=1861;
eyr=2005;
lyrs=length(syr:eyr);
%startstandardize=syr;
%endstandardize=eyr;

% make a year vector for calculations
yr_colm = repmat(syr:eyr,[12 1]);
yr_col=repmat(yr_colm(:),[length(lat)*length(lon),1]);

% get land area excluding Antarctica
sftlf=permute(ncread([varpth 'sftlf_fx_GFDL-ESM2M_historical_r0i0p0.nc'],'sftlf'),[2 1]);
lndidx=find((sftlf(:)>50));
%lndidx=find((sftlf(:)>50)&(nlat>-60)); % get grid cells with < 50% land & not Antarctica

% make a lat vector for calculations
l1=repmat(repmat(lat,1,length(lon)),[1,1,12*lyrs]);
l2=reshape(l1,length(lat)*length(lon),12*lyrs)';
l2=l2(:,lndidx);
lat_col=l2(:);



%nlon=reshape(repmat(lon,1,length(lat))',length(lon)*length(lat),1);

disp('Loading data...')

% Surface net radiation (W/M^2):
rsds=reshape(permute(ncread([varpth 'rsds_Amon_GFDL-ESM2M_historical_r1i1p1_186101-200512.nc'],'rsds'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
rsus=reshape(permute(ncread([varpth 'rsus_Amon_GFDL-ESM2M_historical_r1i1p1_186101-200512.nc'],'rsus'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
rlds=reshape(permute(ncread([varpth 'rlds_Amon_GFDL-ESM2M_historical_r1i1p1_186101-200512.nc'],'rlds'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
rlus=reshape(permute(ncread([varpth 'rlus_Amon_GFDL-ESM2M_historical_r1i1p1_186101-200512.nc'],'rlus'),[2 1 3]),length(lat)*length(lon),12*lyrs)';

netrad=rsds-rsus+rlds-rlus;
%netrad=reshape(netrad,length(lat)*length(lon),12*lyrs);
%Not an inherent model output caclulate by taking (rsds-rsus)+(rlds-rlus)
% Radiation balance = SW_net + LW_net = SW_incoming - SW_outgoing + LW_incoming - LW_outgoing
%%
% t_emperature (C):
t_c=reshape(permute(ncread([varpth 'tas_Amon_GFDL-ESM2M_historical_r1i1p1_186101-200512.nc'],'tas'),[2 1 3]),length(lat)*length(lon),12*lyrs)'-273.15;
%t_f = t_c*33.8; % C to  F

% Surface pressure (Pa):
ps=reshape(permute(ncread([varpth 'ps_Amon_GFDL-ESM2M_historical_r1i1p1_186101-200512.nc'],'ps'),[2 1 3]),length(lat)*length(lon),12*lyrs)';

% Vapor pressure (Pa):
huss=reshape(permute(ncread([varpth 'huss_Amon_GFDL-ESM2M_historical_r1i1p1_186101-200512.nc'],'huss'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
vp=huss.*ps./(.622+.378*huss);
%e=huss*ps/(.622+.378*huss)Not generally inherent model ou2tpu2t so mu2st be calcu2lated

% Surface wind (m/s) (CMIP5 output is at 10m)
u10=reshape(permute(ncread([varpth 'sfcWind_Amon_GFDL-ESM2M_historical_r1i1p1_186101-200512.nc'],'sfcWind'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
% estimate u2m from u10m using wind profile power law
u2=u10*(2/10)^0.143;
%%
% pripitation needs to be (mm/day):
pr=reshape(permute(ncread([varpth 'pr_Amon_GFDL-ESM2M_historical_r1i1p1_186101-200512.nc'],'pr'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
%An example of u2sing opendap to read in pr from BCC
%u2rl='http://strega.ldeo.colu2mbia.edu2:81/expert/CMIP5/.byScenario/.past1000/.atmos/.mon/.pr/.bcc-csm1-1/.r1i1p1/.pr/dods';
%pr=dou2ble(ncread(u2rl,'pr')); 
pr(isnan(pr))=0;

% convert
k1 =[31 28 31 30 31 30 31 31 30 31 30 31]; % days in month, nonleap
F=repmat(k1(:),[lyrs 1]);
%pr=pr*86400;% convert from kg*m-2*s-1 to mm/day
pr_inmon=bsxfun(@times,pr,F)*86400*0.03937;% convert from kg*m-2*s-1 to in/mon





% CONSTANT AWC  
% dimensions: AWC =[locations];
awc_c=5;
AWC=awc_c*ones(length(lat),length(lon));
AWC=AWC(:);

% VARIABLE AWC

% % LDEO code assu2mes that the lower layer of soil has AWC = 5, whereas mean
% % of u2S climate locations is 7, with a range of 4-11 in
% 
% % compu2te AWC by finding the maximu2m mrso then mapping this distribu2tion
% % onto the scale 4-11
% varpth='/d1/nsteiger/climate-data/gfdl-esm2m/';
% mrso=ncread([varpth 'mrso_Lmon_GFDL-ESM2M_historical_r1i1p1_186101-200512.nc'],'mrso');
% maxmrso=max(mrso,[],3,'omitnan');
% % cu2tt off low valu2es of the distribu2tion
% maxmrso(maxmrso<500)=nan;
% % linear mapping to 4-11
% %y= (ymax-ymin)*(x-xmin)/(xmax-xmin) + ymin;
% ymin=2;ymax=11.5;
% xmax=max(maxmrso(:));xmin=min(maxmrso(:));
% awc= (ymax-ymin)*(maxmrso-xmin)/(xmax-xmin) + ymin;



% % Use test data instead
% load('pdsi_test_data.mat','GA02TempAndPrecip');
% lat_col=GA02TempAndPrecip(:,1);
% yr_col=GA02TempAndPrecip(:,2);
% t_f = GA02TempAndPrecip(:,3);
% P = GA02TempAndPrecip(:,4);
% AWC=6;

% lat_col is a column of latitu2de locations that are repeated for each
% month of each location, so for 20 years worth of monthly data at one 
% location gives a vector of length 12*20 = 240, all of the same value.
% Other locations are concatenated after this

[beg_row_col,end_row_col,beg_row_mat,end_row_mat,count_loc,lat_list,lat_mat] = Count_Loc(lat_col);

% use the entire period of record as the calibration period (cf. Karl,
% 1986; Journal of Climate and Applied Meteorology, Vol. 25, No. 1,
% January 1986).

beg_calyr_col=zeros(length(count_loc),1);
end_calyr_col=zeros(length(count_loc),1);
beg_calyr_mat=zeros(length(count_loc),1);
end_calyr_mat=zeros(length(count_loc),1);
for j = 1:count_loc
    beg_calyr_col(j) = 1;
    end_calyr_col(j) = (end_row_col(j) - beg_row_col(j) + 1);
    beg_calyr_mat(j) = (beg_calyr_col(j) - 1)/12 + 1;
    end_calyr_mat(j) = end_calyr_col(j)/12;
end


%=================
% PET CALCULATION
%=================



%PET = Thornthwaite_PET(t_f(:),lat_list,count_loc,beg_row_col,lat_col);

% Calculating the PET (mm/day)
%PET_mm=penmont_PET(t_c,netrad,vp,ps,u2);
PET_mm=penmont_PET(t_c,netrad,vp,ps,[]); % uses u2=1 m/s everywhere

PET_mm_mon=bsxfun(@times,PET_mm,F);%Convert mm/day to mm/month.
PET_in=PET_mm_mon*0.03937;%Convert mm/month to inches/month.


%---------------------
% JUST TAKE LAND AREA
%---------------------

AWC=AWC(lndidx);
P=pr_inmon(:,lndidx);
P=P(:);

disp('Reshaping PET...')
% reshape PET [locations*years,monthss]
PET0=PET_in(:,lndidx)';
PET=zeros(length(lndidx)*lyrs,12);
n=1;
for i=1:length(lndidx)
    pet_o=repmat(PET0(i,:),lyrs,1);
    j=1;k=12;
    for l=1:lyrs
        PET(n,:)=pet_o(l,j:k);
        j=j+12;k=k+12;n=n+1;
    end
end

%PET_col = reshape((PET'),length(PET)*12,1);

disp('Computing WaterBalance...')
[ET,PR,R,RO,PRO,L,PL] = WaterBalance(AWC,PET,P,beg_row_col,end_row_col,count_loc);

disp('Computing Z_Index...')
[Z_all] = Z_Index(P,PET,ET,PR,R,RO,PRO,L,PL,beg_row_mat,end_row_mat, ...
          count_loc,beg_calyr_mat,end_calyr_mat,beg_row_col,end_row_col);

disp('Computing PDSI...');tic
% table = [1='Latitu2de',2='Year',3='PET',4='Z-Index',5='PPe',6='PX1',7='PX2',8='PX3',9='X',10='PDSI',11='PHDI');]      
[table] = PDSI_Central(Z_all,count_loc,beg_row_col,end_row_col,lat_col,yr_col,PET);
disp('Time to compute PDSI: ');toc


pdsi_locs=table(:,10);
      
tmstmp=datestr(now);tmstmp=strrep(tmstmp,' ','_'); % remove space
save(['pdsi_output_u2_1_AWC' num2str(awc_c) '_' tmstmp '.mat'],'-v7.3')























