%load up the precip from the LGM/4XCO_2/ Jolocene and Hosing experiments -- all on the common grid runs and look at the annual mean shifts as
%a function of latitude
clear all; close all;
IT=10; %number of interations of the optimization to perform
LW=2.1; FSz=24; FN='Arial';
 LWA=2.5; LWM = 3.5; trans=.2; LWC=3.7; LW2=2.3; LW3 =2; LWV=1.5;
gray=[0.4 0.4 0.4];
green=[0.1 0.4 0.1];
lgreen= [127 255 0]/255;
dblue=[ 0 0 .5];
 cmap=colormap('jet'); NC=length(cmap);
 
cols = cmap(roundup( NC*(1:IT)/IT),:); %discrete color for each iteration
cbar_ytic=(1:IT)/IT;
flat=-20:.01:20;
%load the data and compose a single matrix of all the experiments --
%P_perturb and P_pi for the contro run
load('data\PRECIP_LGM_common', 'lat', 'lon', 'P_PI', 'P_LGM','MODS')
[NM NT NN N]=size(P_LGM);
do_lgm=1:NM;
P_perturb=P_LGM;
P_pi=P_PI;
load('data\PRECIP_4X_common', 'lat', 'lon', 'P_PI', 'P_4X','MODS')
[NM NT NN N]=size(P_4X);
start=7+1; %start the model count with the end of the LGM
do_4x=start:(NM+start-1);
P_perturb(start:NM+start-1,1:12,1:NN,1:N)=P_4X;
P_pi(start:NM+start-1,1:12,1:NN,1:N)=P_PI;

load('data\PRECIP_HOS_common', 'lat', 'lon', 'P_PI', 'P_HOS','MODS')
[NM NT NN N]=size(P_HOS);
start=23+1; %start the model count with the end of the LGM
do_hos=start:(NM+start-1);
P_perturb(start:NM+start-1,1:12,1:NN,1:N)=P_HOS;
P_pi(start:NM+start-1,1:12,1:NN,1:N)=P_PI;


%%%%%%%%%%%%%%%%%%%%%%% and the holocene runs
load('data\P_HOL', 'P_PI', 'P_HOL', 'lat', 'lon', 'MODS')
[NM NT NN N]=size(P_HOL);
start=30+1; %start the model count with the end of the LGM
do_hol=start:(NM+start-1);
P_perturb(start:NM+start-1,1:12,1:NN,1:N)=P_HOL;
P_pi(start:NM+start-1,1:12,1:NN,1:N)=P_PI;

[NM NT NN N]=size(P_pi);
mlat=lat; %just cause the coast.mat maps have a variable called lat

%now optimize the combined shift/intensification/contraction for each model
%simulaiton
for nm=1:1:NM;%loop over each simulations
lat=mlat; %in case the coast.mat lat is loaded
    
PC=squeeze(mean(squeeze(P_pi(nm,:,:,:)))); %for easy access - control precip we are dealing with
PN=squeeze(mean(squeeze(P_perturb(nm,:,:,:))));

%calculate the precipitation centroid
fP=interp1(lat, nanmean(PC'),flat);
cdf=cumsum(fP)/sum(fP);
PCENT(nm)=flat(find2(cdf,0.5));

fP=interp1(lat, nanmean(PN'),flat);
cdf=cumsum(fP)/sum(fP);
LPCENT(nm)=flat(find2(cdf,0.5));

sp=0.5; %interpolate the tropical precip to a finer grid over which the shoft will be optimized -- can change the spacing
ilat=-20:sp:20; iNN=length(ilat); %the domain we will be optimizing the shift over
ilat_ext=-30:sp:30; iNN2=length(ilat_ext); %you have to extend the domain so that the shifted and contracted fields extend over the optimized domain
ilat_ext2=-45:sp:45; iNN3=length(ilat_ext); % I think this is just so you can extend the fraction of variance explained by the shifted precip beyond the optimization domain for making spatial maps of the fraction of change explained 
[LON LAT]=meshgrid(lon,lat);
[ILON ILAT]=meshgrid(lon, ilat);
[ILON_ext ILAT_ext]=meshgrid(lon, ilat_ext);
[ILON_ext2 ILAT_ext2]=meshgrid(lon, ilat_ext2);

%interpolate to the domains -- perturbed precip first
iPN=interp2(LON, LAT, PN,ILON,ILAT);
iPN_ext=interp2(LON, LAT, PN,ILON_ext,ILAT_ext);
iPN_ext2=interp2(LON, LAT, PN,ILON_ext2,ILAT_ext2);
iZPN=nanmean(iPN'); % the perturbed precip on the interpolation grid
iZPN_ext=nanmean(iPN_ext'); % the perturbed precip on the interpolation grid
iZPN_ext2=nanmean(iPN_ext2'); % the perturbed precip on the interpolation grid

%and control precip
iPC=interp2(LON, LAT, PC,ILON,ILAT);
iPC_ext=interp2(LON, LAT, PC,ILON_ext,ILAT_ext);
iPC_ext2=interp2(LON, LAT, PC,ILON_ext2,ILAT_ext2);
iZPC=nanmean(iPC'); % the perturbed precip on the interpolation grid
iZPC_ext=nanmean(iPC_ext'); % the perturbed precip on the interpolation grid
iZPC_ext2=nanmean(iPC_ext2'); % the perturbed precip on the interpolation grid

 DELP(nm,1:iNN2, 1:N) = iPN_ext-iPC_ext; %precip change on the extended domain
 shift=-2:0.02:2; NS=length(shift); %vector of shifts to consider for the optimization -- if the shift gets greater than 2 degrees, you need to extend this 
 int=.80:.008:1.15; NI=length(int); %vector of intensification scalars
 con=.85:.002:1.1; NC=length(con); %vector of intensification scalars (the latitude will be multiplied by this prioir to interpolation)
 %start with zonal mean shifts, expansion and intesnification considered
 %collectively
 
 %start with initial guesses from performing each shift independently

 %%%%%%%%%%%%%%%%%%%%%%% SHIFTING MODE %%%%%%%%%%%%%%%%%%%%%%%
err_ref = (mean((iZPN-iZPC).^2))^.5; %the error in the unchanged precipitation
 sh_old=0; %initial guesses are unshifted
 con_old=1; %uncontracted
 int_old=1; %unintensified

 tic
 for i=1:IT; %iterations of the error minimization -- I set this to 10 in the first line
     
     
 
    for ns=1:NS; %loop over all shift values considered
         ZP_CI = int_old*interp1(con_old*(ilat_ext2), iZPC_ext2, ilat_ext); %the contracted and intensified precipitation from the previous iteration -- this will be optimally shifted below
         ZP_S = interp1((ilat_ext+shift(ns)), ZP_CI, ilat); %shift the pattern of optimally contracted and intensified precip from last iteration and record the resulting precipitation error below
        err_s(ns) = mean(squeeze(mean(((iZPN-ZP_S).^2))))^.5; % the precip error for the given shift value
    end
 
    sh_new=shift(find(err_s == min(err_s))); %set the new shift to the value that minnimizes the error
    SH_I(nm,i)=sh_new; %save for posteriority
    %visualize the initial and optimaly shifted precip fo each iteration
    ZP_S_new=interp1((ilat_ext+sh_new), ZP_CI, ilat); % the optimally shifted precip for this iteration

  %%%%%%%%%%%%%%%%%%%%%%% CONTRACTION MODE %%%%%%%%%%%%%%%%%%%%%%%
    for nc=1:NC; %loop over all considered contraction values
         ZP_SI = int_old*interp1((ilat_ext2+sh_old), iZPC_ext2, ilat_ext); %the shifted and intensified precipitation from the previous iteration
         ZP_CON = interp1((ilat_ext*con(nc)), ZP_SI, ilat); %now contract that for each contraction value
         err_c(nc) = mean(squeeze(mean(((iZPN-ZP_CON).^2))))^.5; %the resulting precipitation error
    end
 
    con_new=con(find(err_c == min(err_c))); %set the new contraction value to the error minninmizing value
    CON_I(nm,i)=con_new; %save for posteriority
    %visualize the initial and optimaly shifted precip fo each iteration
    ZP_C_new=interp1((con_new*ilat_ext), ZP_SI, ilat); % the optimally shifted precip for this iteration 
 %%%%%%%%%%%%%%%%%%%%%%% INTENSIFICATION MODE %%%%%%%%%%%%%%%%%%%%%%%
    for ni=1:NI; %loop over intesnification scalars considered
         ZP_SC = interp1(con_old*(ilat_ext2+sh_old), iZPC_ext2, ilat_ext);  %the shifted and contracted precipitation from the previous iteration 
         ZP_INT = int(ni)*interp1((ilat_ext), ZP_SC, ilat); %now intensifiy that
         err_i(ni) = mean(squeeze(mean(((iZPN-ZP_INT).^2))))^.5; %and record the error in precip
    end
 
    int_new=int(find(err_i== min(err_i))); %the new intensification scalar is that which minnimizes the error
    INT_I(nm,i)=int_new; %save for posteriority

    %visualize the initial and optimaly shifted precip fo each iteration
    ZP_I_new=int_new*interp1((ilat_ext), ZP_SC, ilat); % the optimally intensified precip for this iteration
 %%%%%%%%%%%%%%%%% prep for next interation
 inc=0.5; %the increment of the optimal change to step forward by -- you sometimes over shoot the best answer if you go to the optimum, so just nudget the answer in that direction
 sh_old=sh_old+inc*(sh_new-sh_old);
 con_old=con_old+inc*(con_new-con_old);
 int_old=int_old+inc*(int_new-int_old);    
        
end
             
%percent of change explained by all three, and excluding on mechanism
ZDP=iZPN-iZPC;
ZALL_err = int_old*interp1(con_old*(ilat_ext+sh_old), iZPC_ext , ilat) - iZPN;  %precip pattern using the shift/intensifiaction and contraction
ZSHIFT_err = int_old*interp1(con_old*(ilat_ext), iZPC_ext , ilat) - iZPN; %precip pattern with no shift - contraction and intensification only
ZCON_err = int_old*interp1((ilat_ext+sh_old), iZPC_ext , ilat) - iZPN; % precip pattern with no contraction -- intensification and shift only
ZINT_err = interp1(con_old*(ilat_ext+sh_old), iZPC_ext , ilat) - iZPN; % error with intensification

ZPER(nm,1:4) = [1- ( (mean((ZALL_err).^2))^.5) / ((mean((ZDP).^2))^.5); %percent of precip change explained with all modes
               1- ( (mean((ZSHIFT_err).^2))^.5) / ((mean((ZDP).^2))^.5); % percent of change explained with no shift -- the % explained by the shift will be the first entry minus this number -- see line 189
               1- ( (mean((ZCON_err).^2))^.5) / ((mean((ZDP).^2))^.5);
               1- ( (mean((ZINT_err).^2))^.5) / ((mean((ZDP).^2))^.5)];
           
           
%ALTERNATIVELY  -- error reduction resulting from a single operation from
%the means state
ZSHIFT_err2 = interp1((ilat_ext+sh_old), iZPC_ext , ilat) - iZPN;
ZCON_err2 = interp1(con_old*(ilat_ext), iZPC_ext , ilat) - iZPN;
ZINT_err2 = int_old*interp1((ilat_ext), iZPC_ext , ilat) - iZPN;

ZPER_ASCI2(nm,1:4) = [ 1-rms(ZALL_err)/ rms(ZDP);
                       1-rms(ZSHIFT_err2)/ rms(ZDP);
                       1-rms(ZCON_err2)/ rms(ZDP);
                       1-rms(ZINT_err2)/ rms(ZDP)];
            
 %lastly, the best answer from the first iteration
ZSHIFT_err3 = interp1((ilat_ext+SH_I(nm,1)), iZPC_ext , ilat) - iZPN;
ZCON_err3 = interp1(CON_I(nm,1)*(ilat_ext), iZPC_ext , ilat) - iZPN;
ZINT_err3 = INT_I(nm,1)*interp1((ilat_ext), iZPC_ext , ilat) - iZPN;

ZPER_ASCI3(nm,1:4) = [ 1-rms(ZALL_err)/ rms(ZDP);
                       1-rms(ZSHIFT_err3)/ rms(ZDP);
                       1-rms(ZCON_err3)/ rms(ZDP);
                       1-rms(ZINT_err3)/ rms(ZDP)];
toc

%% the fraction of change explained by each mode defined as how much of a reduction you get when excluding that mode from the optimized (shift/contract/intensified) precip pattern
ZPER_ASCI(nm, 1:4) = [ZPER(nm,1);
                      ZPER(nm,1)-ZPER(nm,2);
                      ZPER(nm,1)-ZPER(nm,3);
                      ZPER(nm,1)-ZPER(nm,4)];

zonal_percentages = [nm NM round(100*ZPER_ASCI(nm,1:4));
                     nm NM round(100*ZPER_ASCI2(nm,1:4));
                     nm NM round(100*ZPER_ASCI3(nm,1:4));]           
           %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
           %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% 2D
           %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
           %%%%%%%%%%%%%%%%%%% start from initial guesses above
           
%start with initial guesses from performing each shift independently

 
 %%%%%%%%%%%%%%%%%%%%%%% SHIFTING MODE %%%%%%%%%%%%%%%%%%%%%%%
 err_ref = mean(squeeze(mean((iPN-iPC).^2)))^.5; %the error in the unchanged precipitation
tic 

 for i=1:IT; %iterations of the error minimization

     
 
    for ns=1:NS;
         P_CI = int_old*interp2(ILON_ext2, con_old*(ILAT_ext2), iPC_ext2, ILON_ext, ILAT_ext); %the contracted and intensified precipitation from the previous iteration
         P_S = interp2(ILON_ext, (ILAT_ext+shift(ns)), P_CI, ILON, ILAT);
        err_s(ns) = mean(squeeze(mean(((iPN-P_S).^2))))^.5;
    end
 
    sh_new=shift(find(err_s == min(err_s)));
    SH_I_T(nm,i)=sh_new; %save for posteriority
 
  %%%%%%%%%%%%%%%%%%%%%%% CONTRACTION MODE %%%%%%%%%%%%%%%%%%%%%%%
    for nc=1:NC;
         P_SI = int_old*interp2(ILON_ext2,(ILAT_ext2+sh_old), iPC_ext2, ILON_ext, ILAT_ext); %the shifted and intensified precipitation from the previous iteration
         P_CON = interp2(ILON_ext, (ILAT_ext*con(nc)), P_SI, ILON, ILAT);
         err_c(nc) = mean(squeeze(mean(((iPN-P_CON).^2))))^.5;
    end
 
    con_new=con(find(err_c == min(err_c)));
    CON_I_T(nm,i)=con_new; %save for posteriority

 %%%%%%%%%%%%%%%%%%%%%%% INTENSIFICATION MODE %%%%%%%%%%%%%%%%%%%%%%%
    for ni=1:NI;
         P_SC = interp2(ILON_ext2, con_old*(ILAT_ext2+sh_old), iPC_ext2, ILON_ext, ILAT_ext);  %the shifted and contracted precipitation from the previous iteration 
         P_INT = int(ni)*interp2(ILON_ext,(ILAT_ext), P_SC, ILON, ILAT);
         err_i(ni) = mean(squeeze(mean(((iPN-P_INT).^2))))^.5;
    end
 
    int_new=int(find(err_i== min(err_i)));
    INT_I_T(nm,i)=int_new;

 %%%%%%%%%%%%%%%%% prep for next interation
 %the increment of the optimal change to step forward by
 sh_old=sh_old+inc*(sh_new-sh_old);
 con_old=con_old+inc*(con_new-con_old);
 int_old=int_old+inc*(int_new-int_old);   
        
       
        
end
             



%percent of change explained by all three, and excluding on mechanism
DP=iPN-iPC;
ALL_err = int_old*interp2(ILON_ext, con_old*(ILAT_ext+sh_old), iPC_ext , ILON, ILAT) - iPN; 
SHIFT_err = int_old*interp2(ILON_ext, con_old*(ILAT_ext), iPC_ext , ILON, ILAT) - iPN;
CON_err = int_old*interp2(ILON_ext, (ILAT_ext+sh_old), iPC_ext , ILON, ILAT) - iPN;
INT_err = interp2(ILON_ext, con_old*(ILAT_ext+sh_old), iPC_ext , ILON, ILAT) - iPN;

TPER(nm,1:4) = [1- ( (mean(squeeze(mean((ALL_err).^2))))^.5) /err_ref ;
               1- ( mean(squeeze(mean((SHIFT_err).^2) ))^.5) / err_ref;
               1- ( (mean(squeeze(mean((CON_err).^2))))^.5) / err_ref;
               1- ( mean(squeeze(mean((INT_err).^2)))^.5) / err_ref];
           
TPER_ASCI(nm, 1:4) = [TPER(nm,1);
                      TPER(nm,1)-TPER(nm,2);
                      TPER(nm,1)-TPER(nm,3);
                      TPER(nm,1)-TPER(nm,4)];
                  

                  
 %ALTERNATIVELY -- try just perturb the single mode off the mean state
 SHIFT_err2 = interp2(ILON_ext, (ILAT_ext+sh_old), iPC_ext , ILON, ILAT) - iPN;
CON_err2 = interp2(ILON_ext, con_old*(ILAT_ext), iPC_ext , ILON, ILAT) - iPN;
INT_err2 = int_old*interp2(ILON_ext, (ILAT_ext), iPC_ext , ILON, ILAT) - iPN;

TPER_ASCI2(nm,1:4) = [1- ( (mean(squeeze(mean((ALL_err).^2))))^.5) /err_ref ;
               1- ( mean(squeeze(mean((SHIFT_err2).^2) ))^.5) / err_ref;
               1- ( (mean(squeeze(mean((CON_err2).^2))))^.5) / err_ref;
               1- ( mean(squeeze(mean((INT_err2).^2)))^.5) / err_ref];
           
 %LASTLY -- the answer from the first iteration
 SHIFT_err3 = interp2(ILON_ext, (ILAT_ext+SH_I_T(nm,1)), iPC_ext , ILON, ILAT) - iPN;
CON_err3 = interp2(ILON_ext, CON_I_T(nm,1)*(ILAT_ext), iPC_ext , ILON, ILAT) - iPN;
INT_err3 = INT_I_T(nm,1)*interp2(ILON_ext, (ILAT_ext), iPC_ext , ILON, ILAT) - iPN;

TPER_ASCI3(nm,1:4) = [1- ( (mean(squeeze(mean((ALL_err).^2))))^.5) /err_ref ;
               1- ( mean(squeeze(mean((SHIFT_err3).^2) ))^.5) / err_ref;
               1- ( (mean(squeeze(mean((CON_err3).^2))))^.5) / err_ref;
               1- ( mean(squeeze(mean((INT_err3).^2)))^.5) / err_ref];
                  
DP = iPN - iPC;
DP_ALL = int_old*interp2(ILON_ext, con_old*(ILAT_ext+sh_old), iPC_ext , ILON, ILAT) - iPC;
DP_SHIFT = interp2(ILON_ext, (ILAT_ext+sh_old), iPC_ext , ILON, ILAT) - iPC;
DP_CON = interp2(ILON_ext, con_old*(ILAT_ext), iPC_ext , ILON, ILAT) - iPC;
DP_INT = int_old*interp2(ILON_ext, (ILAT_ext), iPC_ext , ILON, ILAT) - iPC;

DEL_P_EXT(nm, 1:length(ilat_ext), 1:N) = iPN_ext-iPC_ext;

DEL_P_ASCI(nm, 1:length(ilat_ext), 1:N) = int_old*interp2(ILON_ext2, con_old*(ILAT_ext2+sh_old), iPC_ext2 , ILON_ext, ILAT_ext) - iPC_ext;

%and for the individual modess
DEL_P_S(nm, 1:length(ilat_ext), 1:N) = interp2(ILON_ext2, (ILAT_ext2+sh_old), iPC_ext2 , ILON_ext, ILAT_ext) - iPC_ext;
DEL_P_C(nm, 1:length(ilat_ext), 1:N) = interp2(ILON_ext2, con_old*(ILAT_ext2), iPC_ext2 , ILON_ext, ILAT_ext) - iPC_ext;
DEL_P_I(nm, 1:length(ilat_ext), 1:N) = int_old*interp2(ILON_ext2, (ILAT_ext2), iPC_ext2 , ILON_ext, ILAT_ext) - iPC_ext;

twod_percentages = [nm NM round(100*TPER_ASCI(nm,1:4));
                    nm NM round(100*TPER_ASCI2(nm,1:4));
                    nm NM round(100*TPER_ASCI3(nm,1:4))] 
toc
end

%break
%save the results
save('C:\Users\Aaron\Desktop\spatial_ITCZ\ASCI_results.mat', 'DEL_P_EXT','DEL_P_ASCI', 'DEL_P_S', 'DEL_P_I', 'DEL_P_C', 'TPER_ASCI', 'ZPER_ASCI', 'TPER_ASCI2', 'ZPER_ASCI2','TPER_ASCI3', 'ZPER_ASCI3', 'do_lgm', 'do_4x', 'do_hol', 'do_hos', 'ilat', 'lon', 'SH_I', 'CON_I','INT_I', 'SH_I_T', 'CON_I_T', 'INT_I_T')