%compare the seasonal heating from surface fluxes and direct solar
%absorption
clear all; close all;
%load the climatologiclay data
dr =    {'C:/research/aqua_gcm/slab_2.4m/';  %directory listing dfor each run
         'C:/research/aqua_gcm/slab_6m/';
         'C:/research/aqua_gcm/slab_12m/';
         'C:/research/aqua_gcm/slab_24m/';
         'C:/research/aqua_gcm/slab_50m/'};
     
NR=5;
colormap('jet');cmap=colormap; %save the RGB colormap vector
cmap=[0 0 1;0 0 1; cmap];

CN=length(cmap);
FSz=22; LW=2.8; FN='Arial'; %figure plottinf options
LW2=1.5; %for the phasing info
depths=[2.4 6 12 24 50];
off=10; %opffset to chop off the end of the colobar
for n=1:3
cmap2(1:CN-off,n)=interp1((0:(CN-1-off))/(CN-1-off), cmap(1:end-off,n), ((0:(CN-1-off))/(CN-1-off)).^.4); %elongate the shallow end of the color scale to add contrast there
end


cut_lat=38; %cutoff latitude
for nr=1:5;  %loop over run (depths)
     
     %cn=round(CN*nr/NR); %index of the given run in the colomap matrix (could make it scale as depth if you want
     cn=round((CN-off)*depths(nr)/depths(NR)); %index of the given run in the colomap matrix (could make it scale as depth if you want
     
     load([char(dr(nr)) 'energy_climatology.mat']);

     FS = FS -US + DS; %the surface heat flux upward from the ocean (not counting the solar that passes through the atmosphere

     NN=length(lat); N=length(lon); %dimensionality
     
     %define the absorbed radiation
     absorb = DT - UT -DS + US; %SW radiation absorbed within the atmosphere
     Zabsorb=squeeze(mean(shiftdim(absorb,2))); %zonal mean absorption
     ZDT=squeeze(mean(shiftdim(DT,2)));
     ZFS=squeeze(mean(shiftdim(FS,2)));
     
     %now integrate over the extratropics
     do = find(lat>38);
     x=sind(lat);
     norm=trapz(x(do), ones(size(do))); %normalization
     LWPH=2; %LineWidth for phase lines
     for mn=1:12;
         EX_ABS(nr, mn)=trapz(x(do), Zabsorb(mn,do))/norm;
         EX_SHF(nr, mn)=trapz(x(do), ZFS(mn,do))/norm;
         EX_DT(nr, mn)=trapz(x(do), ZDT(mn,do))/norm;
     end
     
     ABS_phase(nr)=fund_phase(EX_ABS(nr,:));
     SHF_phase(nr)=fund_phase(EX_SHF(nr,:));
     TOT_phase(nr)=fund_phase(EX_ABS(nr,:)+EX_SHF(nr,:));
     S_phase(nr)=fund_phase(EX_DT(nr,:));
     
     X=monthday(1:12)+15;
     xmn=0; xmx=365;
     ymn=-70; ymx=120;
     EX=1.7;
     RM=1; %Remove the annual mean, 1 for yes and 0 for no
     figure(1); hold on;
     subplot(2,1,1);
     hold on; plot(X,EX_ABS(nr,:)-RM*mean(squeeze(EX_ABS(nr,:))),'-','Color', cmap2(cn,1:3), 'LineWidth', LW);
     hold on; plot(ABS_phase(nr)*[1 1], [ymn ymx], '--', 'Color', cmap2(cn,1:3), 'LineWidth', LWPH);
     figure(2); hold on; subplot(2,1,1);
     hold on; plot(X, EX_SHF(nr,:)-RM*mean(squeeze(EX_SHF(nr,:))),'-','Color', cmap2(cn,1:3), 'LineWidth', LW)
     hold on; plot(SHF_phase(nr)*[1 1], [ymn ymx], '--', 'Color', cmap2(cn,1:3), 'LineWidth', LWPH);
     figure(3); hold on; subplot(2,1,1)
     hold on; plot(X, EX_SHF(nr,:)+ EX_ABS(nr,:) -RM*(mean(squeeze(EX_SHF(nr,:)))+mean(squeeze(EX_ABS(nr,:)))) ,'-','Color', cmap2(cn,1:3), 'LineWidth', LW)
     hold on; plot(TOT_phase(nr)*[1 1], EX*[ymn ymx], '--', 'Color', cmap2(cn,1:3), 'LineWidth', LWPH);
end
xtic=monthday(1:12)+15;
xticl={'JAN';'FEB';'MAR';'APR';'MAY';'JUN';'JUL';'AUG';'SEP';'OCT';'NOV';'DEC'};

LWA=3; LWZ=2.5;

figure(1); hold on;
axis([xmn xmx ymn ymx]);
set(gca, 'FontSize', FSz, 'FontName', FN, 'LineWidth', LWA, 'XTick',xtic,'XTickLabel', xticl); box on; 
hold on; plot([xmn xmx], [0 0], 'k-','LineWidth', LWZ) 
hold on; plot(mean(S_phase)*[1 1], [ymn ymx], 'k--', 'LineWidth', LWPH)

figure(2); hold on;
axis([xmn xmx ymn ymx]);
set(gca, 'FontSize', FSz, 'FontName', FN, 'LineWidth', LWA,'XTick',xtic,'XTickLabel', xticl); box on; 
hold on; plot([xmn xmx], [0 0], 'k-', 'LineWidth', LWZ) 
hold on; plot(mean(S_phase)*[1 1], [ymn ymx], 'k--', 'LineWidth', LWPH)

figure(3); hold on;
axis([xmn xmx EX*ymn EX*ymx]);
set(gca, 'FontSize', FSz, 'FontName', FN, 'LineWidth', LWA, 'XTick',xtic,'XTickLabel', xticl); box on; 
hold on; plot([xmn xmx], [0 0], 'k-','LineWidth', LWZ) 
hold on; plot(mean(S_phase)*[1 1], EX*[ymn ymx], 'k--', 'LineWidth', LWPH)

break
%add in the observations
load('C:\Users\thedhoe\Desktop\aqua_depth\figures\atmos_heating\NH_EX_OBS_EENRGY.mat','NH_TETEN','NH_ABS', 'NH_SHF','NH_TEDIV','NH_OLR', 'min_lat');
load('C:\Users\thedhoe\Desktop\aqua_depth\figures\atmos_heating\SH_EX_OBS_EENRGY.mat','SH_TETEN','SH_ABS', 'SH_SHF','SH_TEDIV','SH_OLR', 'min_lat');

X_NH=monthday(1:12)+15;
NH_off=monthday(6)+22;
SH_off=NH_off-(365/2); %to convert to days past solstice
X_SH=X - (SH_off-NH_off); %put the seasonality onto the NH calendar day by adjusting relative to solstice 
X_SH(find(X_SH>=365))= X_SH(find(X_SH>=365))-365; %make positive throughout
ord=[7:12 1:6]; %order of indexes for the SH


     figure(1); hold on;
     subplot(2,1,1);
     phase_off=-182.5;
     ABS_PHASE_NH=fund_phase(NH_ABS);
     ABS_PHASE_SH=fund_phase(SH_ABS)+phase_off;
     hold on; plot(X_NH,NH_ABS-RM*mean(squeeze(NH_ABS)),'m-', 'LineWidth', LW);
     hold on; plot(X_SH(ord),SH_ABS(ord)-RM*mean(squeeze(SH_ABS)),'m--', 'LineWidth', LW);
     hold on; plot(ABS_PHASE_NH*[1 1], [ymn ymx], 'm-', 'LineWidth', LWPH);
     hold on; plot(ABS_PHASE_SH*[1 1], [ymn ymx], 'm--', 'LineWidth', LWPH);
     
     figure(2); hold on;
     subplot(2,1,1);
     SHF_PHASE_NH=fund_phase(NH_SHF);
     SHF_PHASE_SH=fund_phase(SH_SHF)+phase_off;
     hold on; plot(X_NH,NH_SHF-RM*mean(squeeze(NH_SHF)),'m-', 'LineWidth', LW);
     hold on; plot(X_SH(ord),SH_SHF(ord)-RM*mean(squeeze(SH_SHF)),'m--', 'LineWidth', LW);
     hold on; plot(SHF_PHASE_NH*[1 1], [ymn ymx], 'm-', 'LineWidth', LWPH);
     hold on; plot(SHF_PHASE_SH*[1 1], [ymn ymx], 'm--', 'LineWidth', LWPH);
     

     NH_TOT=NH_ABS+NH_SHF;
     SH_TOT=SH_ABS+SH_SHF;
       
     figure(3); hold on;
     subplot(2,1,1);
     TOT_PHASE_NH=fund_phase(NH_TOT);
     TOT_PHASE_SH=fund_phase(SH_TOT)+phase_off;
     hold on; plot(X_NH,NH_TOT-RM*mean(squeeze(NH_TOT)),'m-', 'LineWidth', LW);
     hold on; plot(X_SH(ord),SH_TOT(ord)-RM*mean(squeeze(SH_TOT)),'m--', 'LineWidth', LW);
     hold on; plot(TOT_PHASE_NH*[1 1], EX*[ymn ymx], 'm-', 'LineWidth', LWPH);
     hold on; plot(TOT_PHASE_SH*[1 1], EX*[ymn ymx], 'm--', 'LineWidth', LWPH);
break
     ANN_absorb(nr,1:NN) = mean(squeeze(mean(absorb))');
     ANN_FS(nr, 1:NN) = mean(squeeze(mean(FS))');
     
     %harmonic analysis of the solar (downweliing at TOA
     [Sc,Sph]=fund2(squeeze(mean((shiftdim(DT,2))))'); %find the magnitude and phase of the fundamental of the zonal mean downwelling
     
     %harmonic analysis of absorbed radiation in the atmosphere
     [Ac,Aph]=fund2(squeeze(mean((shiftdim(absorb,2))))'); %find the magnitude and phase of the fundamental of the zonal mean absorbed
     SEAS_absorb(nr,1:NN)=Ac;
     
     %harmonic analysis of surface heat flux
     [Fc,Fph]=fund2(squeeze(mean((shiftdim(FS,2))))'); %find the magnitude and phase of the fundamental of the zonal mean surface heat flux
     SEAS_FS(nr,1:NN)=Fc;
     elim=25; %latitude equatorward of which to eliminate from the phase analysis
     ph_SF = 360/(2*pi); %converts the phase to days
     for nn=1:NN %loop over latitude
         %define the phase difference between the solar and the atmospheric
         %absorption
         if abs(lat(nn)) <= elim
             absorb_phase(nr, nn) = NaN; 
         elseif (Aph(nn)-Sph(nn)) < pi
            absorb_phase(nr,nn)=ph_SF*(Aph(nn) -Sph(nn)); %phase is the difference between the solar phase and the absorbed
         else
            absorb_phase(nr,nn)=ph_SF*(Aph(nn) -Sph(nn) - 2*pi); %phase is the difference between the solar phase and the absorbed
         end
         
         %define the phase difference between the solar and FS
         if abs(lat(nn)) <= elim
             FS_phase(nr, nn) = NaN; 
         elseif (Fph(nn)-Fph(nn)) < pi
            FS_phase(nr,nn)=ph_SF*(Fph(nn) -Sph(nn)); %phase is the difference between the solar phase and the absorbed
         else
            F_phase(nr,nn)=ph_SF*(Fph(nn) -Sph(nn) - 2*pi); %phase is the difference between the solar phase and the absorbed
         end
         
         %solar phasing (day of maximum)
         if Sph(nn) <= 0;
            sol_max(nr,nn) =-Sph(nn)*365/(2*pi);
         else
             sol_max(nr,nn) = 365 -Sph(nn)*365/(2*pi);
         end
         
         %Surface Heat flux phasing (day of maximum)
         if Fph(nn) <= 0;
            FS_max(nr,nn) =-Fph(nn)*365/(2*pi);
         else
             FS_max(nr,nn) = 365 -Fph(nn)*365/(2*pi);
         end
         
         %Surface Heat flux phasing (day of maximum)
         if Aph(nn) <= 0;
            A_max(nr,nn) =-Aph(nn)*365/(2*pi);
         else
             A_max(nr,nn) = 365 -Aph(nn)*365/(2*pi);
         end
         
         
         
     %end
     
     FS_lag(nr,1:NN) = FS_max(nr,1:NN) - sol_max(nr,1:NN);
     FS_lag(nr, find(FS_lag(nr,:) < 0)) = FS_lag(nr,find(FS_lag(nr,:) < 0)) + 365; %makes is positive
     FS_lag(nr, find(abs(lat)<=elim))= NaN;
     
     A_lag(nr,1:NN) = A_max(nr,1:NN) - sol_max(nr,1:NN);
     A_lag(nr, find(A_lag(nr,:) < 0)) = A_lag(nr,find(A_lag(nr,:) < 0)) + 365; %makes is positive
     A_lag(nr, find(abs(lat)<=elim))= NaN;
     
     ymin=0; ymax=250;
     figure(19); hold on;
     ns=find2(lat,50); %location  to look at the seasonality
     hold on; plot(monthday(1:12),Zabsorb(:,ns), '-','Color', cmap(cn,1:3), 'LineWidth', LW)
     %hold on; plot(monthday(1:12), ZDT(:,ns), ':','Color', cmap(cn,1:3), 'LineWidth', LW)
     hold on; plot(monthday(1:12), ZFS(:,ns),'--', 'Color', cmap(cn,1:3), 'LineWidth', LW);
     hold on; plot(sol_max(nr,ns)*[1 1], [ymin ymax], 'k-', 'LineWidth', 3);
     hold on; plot(A_max(nr,ns)*[1 1], [ymin ymax], '-', 'Color', cmap(cn,1:3), 'LineWidth', LW2);
     hold on; plot(FS_max(nr,ns)*[1 1], [ymin ymax], '--', 'Color', cmap(cn,1:3), 'LineWidth', LW2);

     
     X=lat; %can make it sine of latitude
     figure(1); hold on;
     plot(X,ANN_absorb(nr,:),'-','Color', cmap(cn,1:3), 'LineWidth', LW)
     hold on; plot(X,ANN_FS(nr,:),'--','Color', cmap(cn,1:3), 'LineWidth', LW)
     legend('Atmospheric Absorption', 'Surface Heat Flux')
     
     figure(2); hold on;
     plot(X,SEAS_absorb(nr,:),'-','Color', cmap(cn,1:3), 'LineWidth', LW)
     hold on; plot(X,SEAS_FS(nr,:),'--','Color', cmap(cn,1:3), 'LineWidth', LW)
     legend('Atmospheric Absorption', 'Surface Heat Flux')
    
     figure(3); hold on; 
     plot(X, A_lag(nr,:), '-','Color', cmap(cn,1:3), 'LineWidth', LW)
     plot(X,FS_lag(nr,:), '--','Color', cmap(cn,1:3), 'LineWidth', LW)
     %plot(X, FS_phase(nr,:),'--','Color', cmap(cn,1:3), 'LineWidth', LW)
     
     Xtic(nr) = cn;
     %Xticl(nr) = nr;
     Xticl(nr)=depths(nr);
     
     
     %sphcolor2jet(std(absorb), lat,lon,3)
     %sphcolor2jet(std(FS),lat,lon,4)
end

figure(1); hold on;
set(gca, 'FontSize', FSz, 'FontName', FN)
 xlabel('Latitude', 'FontSize', FSz', 'FontName', FN)
 ylabel(' Annual Mean Heating (Wm^-^2)','FontSize', FSz, 'FontName', FN) 
 H=colorbar('location', 'SouthOutside', 'XTick',Xtic, 'XTickLabel', Xticl); 
 set(gca, 'FontSize', FSz, 'FontName', FN);
 title('Annual mean atmospheric heating', 'FontSize', FSz, 'FontName', FN)
 
 figure(2); hold on;
set(gca, 'FontSize', FSz, 'FontName', FN)
 xlabel('Latitude', 'FontSize', FSz', 'FontName', FN)
 ylabel('Seasonal Amplitude of Heating (Wm^-^2)','FontSize', FSz, 'FontName', FN) 
 H=colorbar('location', 'SouthOutside', 'XTick',Xtic, 'XTickLabel', Xticl); 
 set(gca, 'FontSize', FSz, 'FontName', FN);
 title('Seasonal Amplitude of Atmospheric Heating', 'FontSize', FSz, 'FontName', FN)
 
  
 figure(3); hold on;
set(gca, 'FontSize', FSz, 'FontName', FN)
 xlabel('Latitude', 'FontSize', FSz', 'FontName', FN)
 ylabel('Phase lag of maximum from incident solar maximum (days)','FontSize', FSz, 'FontName', FN) 
 H=colorbar('location', 'SouthOutside', 'XTick',Xtic, 'XTickLabel', Xticl); 
 set(gca, 'FontSize', FSz, 'FontName', FN);
 title('Seasonal Phase of Atmospheric Heating', 'FontSize', FSz, 'FontName', FN)
 
  figure(19); hold on;
set(gca, 'FontSize', FSz, 'FontName', FN)
 xlabel('Month', 'FontSize', FSz', 'FontName', FN)
 ylabel('Zonal mean heating (Wm^-^2)','FontSize', FSz, 'FontName', FN) 
 H=colorbar('location', 'SouthOutside', 'XTick',Xtic, 'XTickLabel', Xticl); 
 set(gca, 'FontSize', FSz, 'FontName', FN);
 title(['Seasonal Cyle of Heating at ' num2str(round(lat(ns))) ' ^o N'], 'FontSize', FSz, 'FontName', FN)
 legend(' Absorbed Directly', 'Surface Flux','Solar Phase', 'Absorption Phase', 'Surface Phase' )
 
 cd figures
 