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

% LOAD THE PARAMETERS
%spei_params;

disp(' ')
disp('Loading data...')

syr=1901;eyr=2014;lyrs=length(syr:eyr);
lat=ncread('/d1/nsteiger/climate-data/cru_ts3/cru_ts3.23.1901.2014.pre.dat.nc','lat');
lon=ncread('/d1/nsteiger/climate-data/cru_ts3/cru_ts3.23.1901.2014.pre.dat.nc','lon');



% mm/month
pr_mm_mon=reshape(permute(ncread('/d1/nsteiger/climate-data/cru_ts3/cru_ts3.23.1901.2014.pre.dat.nc','pre'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
% convert mm/month to inches/month
pr_inmon=pr_mm_mon*0.03937;%Convert mm/month to inches/month.


% make sure there are no NaN precip values (not really an issue at monthly
% resolution)
%pr(isnan(pr))=0;


% constant available water content
awc_o='c';
awc_c=7;awc_o=[awc_o '_' num2str(awc_c)];
AWC=awc_c*ones(length(lat)*length(lon),1);


% mm/day
PET_mm=reshape(permute(ncread('/d1/nsteiger/climate-data/cru_ts3/cru_ts3.23.1901.2014.pet.dat.nc','pet'),[2 1 3]),length(lat)*length(lon),12*lyrs)';
% convert pet
k1 =[31 28 31 30 31 30 31 31 30 31 30 31]; % days in month, nonleap
F=repmat(k1(:),[lyrs 1]);


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.



% get land area (PET has missing values along some coastal regions and
% islands)
field=reshape(permute(ncread('/d1/nsteiger/climate-data/cru_ts3/cru_ts3.23.1901.2014.pet.dat.nc','pet'),[2 1 3]),length(lat)*length(lon),12*lyrs);
nan_f=sum(~isnan(field),2);
lndidx=find((nan_f>0)); % get grid cells with any land

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

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

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

% 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


%==================
% PDSI
%==================

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

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

disp('Reshaping PET...')
% reshape PET [locations*years,months]
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');]
%parpool(2)
%[table] = PDSI_parallel(Z_all,count_loc,beg_row_col,end_row_col,lat_col,yr_col,PET);
[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']);


save(['/d1/nsteiger/climate-data/pdsi/pdsi_cruts3_prelim_out.mat'],'-v7.3')

% PULL OUT PDSI VALUES
pdsi_locs=zeros(lyrs*12*length(lndidx),1);
jj=1;kk=lyrs*12;
for i=1:length(table)
    tbl=table{i};
    pdsi_locs(jj:kk)=tbl(:,10);
    jj=jj+lyrs*12;kk=kk+lyrs*12;
end

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

% Set outliers to +-20
pdsi_lndonly_orig=pdsi_lndonly;
pdsi_lndonly(pdsi_lndonly>20)=20;
pdsi_lndonly(pdsi_lndonly<-20)=-20;

% SAVE THE DATA
tmstmp=datestr(now);tmstmp=strrep(tmstmp,' ','_'); % remove space


save(['/d1/nsteiger/climate-data/pdsi/pdsi_cruts3_AWC_' awc_o '_' tmstmp '.mat'],...
        'pdsi_f','pdsi_lndonly','pdsi_lndonly_orig','lndidx','lyrs','syr','eyr','lat','lon','-v7.3')
% save everything...
%save([varpth 'pdsi_output_AWC_' num2str(awc_c) '_u2_' u2o '_' tmstmp '.mat'],'-v7.3')
























