%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
CN=length(cmap);
figure(1); clf;
FSz=22; LW=2; FN='Arial'; %figure plottinf options
LW2=1.5; %for the phasing info
depths=[2.4 6 12 24 50];
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
     load([char(dr(nr)) 'energy_climatology.mat']);
     load([char(dr(nr)) 'TV_trop.mat']);
     %load([char(dr(n)) 'T_clim.mat'], 'TS','T','lev');
     %nl = find2(lev, 700); %index for the vertical level you want
     %TV=squeeze(T(:,nl,:,:));
     TV=TV_trop;
     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)));
     
          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
 