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

% LOAD THE MODEL DATA

disp('Loading data...')

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

%syr=1861;eyr=2005;lyrs=length(syr:eyr);
syr=1871;eyr=2011;lyrs=length(syr:eyr);

%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');
load('/d1/nsteiger/climate-data/echam5/e5_latlon.mat')

% land area
%sftlf=permute(ncread([varpth 'sftlf_fx_GFDL-ESM2M_historical_r0i0p0.nc'],'sftlf'),[2 1]);
load([varpth 'sftlf_fx_ECHAM5-wiso.mat'])

% 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)';
trads=reshape(permute(flip(ncread([varpth 'trads_mon_1871_2011.nc'],'trads'),2),[2 1 3]),length(lat)*length(lon),12*lyrs)';
srads=reshape(permute(flip(ncread([varpth 'srads_mon_1871_2011.nc'],'srads'),2),[2 1 3]),length(lat)*length(lon),12*lyrs)';

% 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_c=reshape(permute(flip(ncread([varpth 't2m_mon_1871_2011.nc'],'temp2'),2),[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)';
ps=reshape(permute(flip(ncread([varpth 'ps_mon_1871_2011.nc'],'aps'),2),[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)';
huss=reshape(permute(flip(squeeze(ncread([varpth 'qsurf_mon_1871_2011.nc'],'q')),2),[2 1 3]),length(lat)*length(lon),12*lyrs)';

% 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)';
u10=reshape(permute(flip(ncread([varpth 'wind10_mon_1871_2011.nc'],'wind10'),2),[2 1 3]),length(lat)*length(lon),12*lyrs)';

% precipitation, will be converted from kg*m^-2*s^-1 to (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)';
pr=reshape(permute(flip(ncread([varpth 'precip_mon_1871_2011.nc'],'precip'),2),[2 1 3]),length(lat)*length(lon),12*lyrs)';

%An example of using opendap to read in pr from BCC
%url='http://strega.ldeo.columbia.edu:81/expert/CMIP5/.byScenario/.past1000/.atmos/.mon/.pr/.bcc-csm1-1/.r1i1p1/.pr/dods';
%pr=double(ncread(url,'pr'));

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

disp('Pre-computations...')

% 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
lndidx=find((sftlf(:)>50)); % get grid cells with < 50% land
%lndidx=find((sftlf(:)>50)&(nlat>-60)); % & not Antarctica => causes
%problems with the calculation

% 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(:);

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

% Surface net radiation (W/M^2):
%netrad=rsds-rsus+rlds-rlus;
disp('Net radiation calculation for ECHAM5-wiso')
netrad=srads+trads; 
%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

% Vapor pressure
vp=huss.*ps./(.622+.378*huss);
%e=huss*ps/(.622+.378*huss)Not generally inherent model output so must be calculated

% Estimate u2m from u10m using wind profile power law (empirical est.); the
% mean wind at 2m based on global observations is ~2 m/s
%u2=u10*(2/10)^0.143;
u2=u10*4.87/log(67.8*10-5.42);
%u2=u10;

% make sure there are no NaN precip values (not really an issue at monthly
% resolution)
pr(isnan(pr))=0;
% convert precipitation
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])
% LDEO code assumes that the lower layer of soil has AWC = 5, whereas mean
% of US climate locations is 7, with a range of 4-11 in
% awc_c=7;
% AWC=awc_c*ones(length(lat)*length(lon),1);

% % VARIABLE AWC
% % compute AWC by finding the maximum mrso then mapping this distribution
% % onto the scale 4-11
% varpth='/d1/nsteiger/climate-data/gfdl-esm2m/';
% mrsoi=permute(ncread([varpth 'mrsos_day_GFDL-ESM2M_historical_r1i1p1_18660101-20051231.nc'],'mrsos'),[2 1 3]);
% mrsoi(mrsoi==0)=NaN;
% mrso=max(mrso,[],3,'omitnan');
% xmax=max(mrso(:));xmin=min(mrso(:));
% % linear mapping to ymin/ymax
% % y = (ymax-ymin)*(x-xmin)/(xmax-xmin) + ymin;
% ymin=0;ymax=20;
% awc= (ymax-ymin)*(mrso-xmin)/(xmax-xmin) + ymin;
% AWC=reshape(awc,length(lat)*length(lon),1);

% USE MODEL'S INHERENT FIELD CAPACITY
load([varpth 'wsmx_fx_ECHAM5-wiso.mat'])% in meters
mrso=(wsmx.*sftlf/100)*(39.37-1); % just take same land portion as above, convert to inches minus top inch of water for top layer
AWC=reshape(mrso,length(lat)*length(lon),1);


% lat_col is a column of latitude 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
%=================

% Thornwaite method
%t_f = t_c*33.8; % C to  F
%PET = Thornthwaite_PET(t_f(:),lat_list,count_loc,beg_row_col,lat_col);

% Penman-Monteith method (mm/day)
u2o='gcm'; % 2 m wind from the climate model
%u2o='1ms'; % assume uniform 1 m/s wind at 2 m
if strcmp(u2o,'gcm')
    PET_mm=penmont_PET(t_c,netrad,vp,ps,u2);
elseif strcmp(u2o,'1ms')
    PET_mm=penmont_PET(t_c,netrad,vp,ps,[]);
end

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)
    j=1;k=12;
    for l=1:lyrs
        PET(n,:)=PET0(i,j:k);
        j=j+12;k=k+12;n=n+1;
    end
end

%=================
% WATER BALANCE
%=================

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

%==================
% Z-INDEX AND PDSI
%==================

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);


% NEED TO PARALLELIZE THIS PORTION OF THE CODE... SPLIT INTO MULTIPLE CHUNKS AND COMPUTE DIFFERENT GRID POINTS SEPARATELY

disp(' ')
disp('Computing PDSI...');tic
% table = [1='Latitude',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);
tmr=toc;
disp(['Time to compute PDSI: ' num2str(tmr/60) ' mins']);

% PDSI VALUES RE
pdsi_locs=table(:,10);
pdsi_lndonly=reshape(pdsi_locs,lyrs*12,length(lndidx))';
pdsi_f=NaN(length(lat)*length(lon),lyrs*12);
pdsi_f(lndidx,:)=pdsi_lndonly;
pdsi_f=reshape(pdsi_f,length(lat),length(lon),lyrs*12);

% SAVE THE DATA
tmstmp=datestr(now);tmstmp=strrep(tmstmp,' ','_'); % remove space
%save([varpth 'pdsi_output_AWC_' num2str(awc_c) '_u2_' u2o '_' tmstmp '.mat'],...
%    'pdsi_f','pdsi_lndonly','lndidx','lyrs','syr','eyr','lat','lon','-v7.3')
save([varpth 'pdsi_output_AWC_mdl_u2_' u2o '_' tmstmp '.mat'],...
    'pdsi_f','pdsi_lndonly','lndidx','lyrs','syr','eyr','lat','lon','-v7.3')
% save everything...
%save([varpth 'pdsi_output_AWC_' num2str(awc_c) '_u2_' u2o '_' tmstmp '.mat'],'-v7.3')























