%compare the seasonal heating from surface fluxes and direct solar
%absorption
clear all; close all;
%load the climatologiclay data
dr =    {'C:\Users\aaron\Desktop\aqua_depth\2.4\';  %directory listing dfor each run
         'C:\Users\aaron\Desktop\aqua_depth\6\';
         'C:\Users\aaron\Desktop\aqua_depth\12\';
         'C:\Users\aaron\Desktop\aqua_depth\24\';
         'C:\Users\aaron\Desktop\aqua_depth\50\'};
     
NR=5;
colormap('jet');cmap=colormap; %save the RGB colormap vector
CN=length(cmap);
figure(1); clf;
depth={'2.4';'6';'12';'24';'50'}
FSz=22; LW=2; FN='Arial'; %figure plottinf options
LW2=1.5; %for the phasing info
LW=2;
g=9.81;%gravity
cp=1007; %specific heat at constant pressurre (J/(kg *K(\)
depths=[2.4 6 12 24 50];
col=[         0         0    0.5625
         0    0.5000    1.0000
    0.4375    1.0000    0.5625
    1.0000    0.6250         0
    0.5         0         0]; % discrete colormap (stolen from jet)

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)) 'MSE_TEN.mat']);
     load([char(dr(nr)) 'T.mat']);
     
     %zonal average the temperature
     T=squeeze(mean(shiftdim(T,3)));
     NL=length(Pfull); NN=length(lat);
     %then take the fundamental amplitude and phase
     cross=42; %crossover latitude poleward of which to average over
     doz=find(Pfull>200); %troposphere
     doz2=find(Pfull > 10 & Pfull <300); %upper troposphere
     do=find(lat>cross); %domain of latitudes
     [coef, phase] = fund2_3d(T(:,:,do));
     PHASE=((2*pi-phase)*365/(2*pi)+15)-175; %day of maximum relative to solstice
     
     Phalf=mean([Pfull(1:end-1) Pfull(2:end)]')';
     Pdiff=[Phalf;1000]-[0;Phalf]; %pressure differenc by layer
     W=zeros(NL,1);W(doz)=Pdiff(doz); W=W/sum(W); %normalized weighting vector for mass integrals over the troposphere
     W2=zeros(NL,1);W2(doz2)=Pdiff(doz2); W2=W2/sum(W2); %normalized weighting vector for mass integrals over the troposphere
     
     for mn=1:12;
         field=squeeze(T(mn,:,:));
         TV(mn,1:NN)= W'*field; %mass weighted vertical average of T for each month and latitude 
         TV2(mn,1:NN)= W2'*field; %mass weighted vertical average of T for each month and latitude 
     end
     
     color_trans;
     
     %make the vertical structure of the temperature
     avg=ones(length(doz),1)*mean(coef(doz,:)); %column average seasonal cycle 
     
     
%      MAX=1.3*max(max(coef(doz,:)));
    [LAT,LEV]=meshgrid(lat(do),Pfull(doz)); %mesh for vertical cross sections
%      figure(10+nr)%;subplot(1,2,1);
%      pcolor(LAT, -LEV, coef(doz,:)./avg); colormap(cmap);
%      caxis([.8 1.2]);
%      colorbar;
     
          %color_dhoe3_fine;
     figure(20+nr)
     colormap('jet');
     pcolor(LAT, -LEV, coef(doz,:)); colorbar 
     ytic=-1000:100:-300;
     yticl=-ytic;
     xlabel('Latitude', 'FontSize', FSz, 'FontName', FN)
     ylabel('Pressure level (hPa)', 'FontSize', FSz, 'FontName', FN)
     shading interp
     set(gca,'FontSize', FSz, 'FontName', FN, 'LineWidth', 2, 'YTick', ytic, 'YTickLabel', yticl); 
     title([char(depth(nr,:)) ' Meters'], 'FontSize', FSz,'FontName', FN) 
     
     
     figure(30+nr)
     %ytic=-1000:100:-300;
     %yticl=-ytic;
     colormap('jet');
     pcolor(LAT, -LEV, PHASE(doz,:)); caxis([45 150]); colorbar;%caxis([45 115]);colorbar 
     xlabel('Latitude', 'FontSize', FSz, 'FontName', FN)
     ylabel('Pressure level (hPa)', 'FontSize', FSz, 'FontName', FN)
     shading interp
     set(gca,'FontSize', FSz, 'FontName', FN, 'LineWidth', 2, 'YTick', ytic, 'YTickLabel', yticl); 
     
     %load([char(dr(n)) 'T_clim.mat'], 'TS','T','lev');
     %nl = find2(lev, 700); %index for the vertical level you want
     
     SHF = squeeze(mean(shiftdim(FS -US + DS,2))); %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
     SWABS = squeeze(mean(shiftdim(DT - UT -DS + US,2))); %SW radiation absorbed within the atmosphere
     ZDT=squeeze(mean(shiftdim(DT,2)));
     ZOLR=squeeze(mean(shiftdim(OLR,2)));
     ZTETEN3=squeeze(mean(shiftdim(TTEN,2)));
     ZTETEN=MSE_TEN/2.03; %the 2 month difference was divided by thirty days when it sohould have been divided by 365/6 (this corrects that)
     ZDSETEN=DSE_TEN/2.03;
     ZTETEN2=(sum(Pdiff(doz))*100*cp/g) * ( TV([2:12 1],:) - TV([12 1:11],:)) / (24*3600*365/6);
     %ZTEDIV=squeeze(mean(shiftdim(TEDIV,2)));
     ZTEDIV=ZTETEN+ZOLR-SWABS-SHF; %as a residual
     figure(1); hold on; 
     plot(lat, mean(SWABS), '--','Color', col(nr, 1:3), 'LineWidth', LW);
     hold on; plot(lat, mean(SHF), '-','Color', col(nr, 1:3), 'LineWidth', LW);
     

     %now define the seasonal amplitude
     SF=2^.5;
     for nn=1:NN;
        ind =(sqrt(12/11))*squeeze(ZDT(:,nn)-squeeze(mean(ZDT(:,nn))))/std(squeeze(ZDT(:,nn))); %normalized index to give regression values (accounts for small sample size of std
         S_SWABS(nn) = SF*mean(ind.*squeeze(SWABS(:,nn)));
         S_SHF(nn) = SF*mean(ind.*squeeze(SHF(:,nn)));
         S_TEDIV(nn) = SF*mean(ind.*squeeze(ZTEDIV(:,nn)));
        S_OLR(nn) = -SF*mean(ind.*squeeze(ZOLR(:,nn)));
        S_TETEN(nn) = -SF*mean(ind.*squeeze(ZTETEN(:,nn)));
        %S_TETEN2(nn) = -SF*mean(ind.*squeeze(ZTETEN2(:,nn)));
        %S_TETEN(nn) = -SF*mean(ind.*squeeze(MSE_TEN(:,nn)));
     end
     
        figure(2); hold on; 
        plot(lat, S_SWABS, '--','Color', col(nr, 1:3), 'LineWidth', LW);
        hold on; plot(lat, S_SHF, '-','Color', col(nr, 1:3), 'LineWidth', LW);
     
        
        %average poleward of a certaon latitude
        x=sind(lat); xdo=x(do); %sine of latitude (for area weighting over the polar cap
        norm=trapz(xdo,ones(size(xdo))); %normaliztion for averages over this region
        for mn=1:12;
        RM=1; %remove annual mean? 1 for yes and 0 for no    
        R_SWABS(mn)=trapz(xdo, SWABS(mn,do)-RM*mean(SWABS(:,do)))/norm; %monthly average over the polar cap
        R_SHF(mn)=trapz(xdo, SHF(mn,do)-RM*mean(SHF(:,do)))/norm; %monthly average over the polar cap
        R_TEDIV(mn)=trapz(xdo, ZTEDIV(mn,do)-RM*mean(ZTEDIV(:,do)))/norm; %monthly average over the polar cap
        R_OLR(mn)=trapz(xdo, ZOLR(mn,do)-RM*mean(ZOLR(:,do)))/norm; %monthly average over the polar cap
        R_DT(mn)=trapz(xdo, ZDT(mn,do)-RM*mean(ZDT(:,do)))/norm; %monthly average over the polar cap
        R_TETEN(mn)=trapz(xdo, MSE_TEN(mn,do)-RM*mean(MSE_TEN(:,do)))/norm; %monthly average over the polar cap
        R_TV(mn)=trapz(xdo, TV(mn,do)-RM*mean(TV(:,do)))/norm;
        R_TV2(mn)=trapz(xdo, TV2(mn,do)-RM*mean(TV2(:,do)))/norm;
        end
        
        %now crank out the phases of the terms
        SWABS_ph(nr)=fund_phase(R_SWABS); %day of the peak
        DT_ph(nr)=fund_phase(R_DT); %day of the peak
        SHF_ph(nr)=fund_phase(R_SHF); %day of the peak
        TETEN_ph(nr)=fund_phase(R_TETEN); %day of the peak
        TV_ph(nr)=fund_phase(R_TV); %day of the peak
        TV2_ph(nr)=fund_phase(R_TV2); %day of the peak
        OLR_ph(nr)=fund_phase(R_OLR); %day of the peak
        TEDIV_ph(nr)=fund_phase(-R_TEDIV); %day of the peak
        HEAT_ph(nr)=fund_phase(R_SHF+R_SWABS);
        
        figure(3); hold on; 
        plot(monthday(1:12)+15, R_SWABS, '--','Color', col(nr, 1:3), 'LineWidth', LW);
        hold on; plot(monthday(1:12)+15, R_SHF, '-','Color', col(nr, 1:3), 'LineWidth', LW);
        
        figure(40+nr); hold on; 
        plot(monthday(1:12)+15, R_OLR, 'g', 'LineWidth', LW);
        hold on; plot(monthday(1:12)+15, R_TETEN, 'c', 'LineWidth', LW);
        hold on; plot(monthday(1:12)+15, -R_TEDIV, 'm', 'LineWidth', LW);
        hold on; plot(monthday(1:12)+15, 2*R_TV, 'k', 'LineWidth', LW)
        AX=axis;
        ymn=AX(3); ymx=AX(4);
        hold on; plot(OLR_ph(nr)*[1 1], [ymn ymx], 'g--');
        hold on; plot(TEDIV_ph(nr)*[1 1], [ymn ymx], 'm--');
        hold on; plot(TV_ph(nr)*[1 1], [ymn ymx], 'k--');
        hold on; plot(TETEN_ph(nr)*[1 1], [ymn ymx], 'c--');
        hold on; plot( (TETEN_ph(nr)+91.2)*[1 1], [ymn ymx], 'c:');
         
        rescale=5;
        figure(4); hold on; 
        plot(monthday(1:12)+15, R_SWABS+R_SHF, '-','Color', col(nr, 1:3), 'LineWidth', LW);
        hold on; plot(monthday(1:12)+15, rescale*R_TV, '--','Color', col(nr, 1:3), 'LineWidth', LW);
        
        

end
LWph=1.1; %linewdith for the vertical lines showing the phase of the peak
for nr=1:5;
    figure(3); hold on; 
    AX=axis; ymn=AX(3); ymx=AX(4);
    plot(SWABS_ph(nr)*[1 1], [ymn ymx], '--', 'Color', col(nr,1:3), 'LineWidth', LWph); %plot the SWABS PHASE as a vertical line
    plot(SHF_ph(nr)*[1 1], [ymn ymx], '-', 'Color', col(nr,1:3), 'LineWidth', LWph); %plot the SHF PHASE as a vertical line
    
    figure(4); hold on; 
    AX=axis; ymn=AX(3); ymx=AX(4);
    plot(HEAT_ph(nr)*[1 1], [ymn ymx], '-', 'Color', col(nr,1:3), 'LineWidth', LWph); %plot the total heating PHASE as a vertical line
    plot(TV_ph(nr)*[1 1], [ymn ymx], '--', 'Color', col(nr,1:3), 'LineWidth', LWph); %plot the temperature PHASE as a vertical line
end

%the phasing of the terms in each run
figure(7);clf;
loc=1:5; %location for each run on the phase plot
REF=monthday(6)+21; %reference date
MS=18;
for nr=1:5;
    hold on;plot(HEAT_ph(nr)-REF, loc(nr), 'rd', 'MarkerSize', MS, 'MarkerFaceColor', 'r');
    hold on;plot(TV_ph(nr)-REF, loc(nr), 'kd', 'MarkerSize', MS, 'MarkerFaceColor', 'k');
    hold on;plot(TETEN_ph(nr)-REF, loc(nr), 'cd', 'MarkerSize', MS, 'MarkerFaceColor', 'c');
    hold on;plot(OLR_ph(nr)-REF, loc(nr), 'gd', 'MarkerSize', MS, 'MarkerFaceColor', 'g');
    hold on;plot(TEDIV_ph(nr)-REF, loc(nr), 'md', 'MarkerSize', MS, 'MarkerFaceColor', 'm');
    hold on;plot(TV2_ph(nr)-REF, loc(nr), 'kd', 'MarkerSize', MS, 'MarkerFaceColor', 'w');
end

figure(7); hold on;
legend('Atmospheric Heating', 'Tropospheric Temperature', 'Column Tendency', 'OLR', 'Heat Transport div.','Stratospheric temp.')
ymn=0.5; ymx=7.1;
plot([0 0], [ymn ymx], 'r--', 'LineWidth', 1.5); hold on;
plot(91.2*[1 1], [ymn ymx], 'k--', 'LineWidth', 1.5); hold on;
xlabel('Days past the summer solstice', 'FontSize', FSz, 'FontName', FN);
ylabel('Mixed Layer Depth (meters)', 'FontSize', FSz, 'FontName', FN);
title({'Seasonal Phasing of energy fluxes and'; 'temperature as a function of mixed layer depth'}, 'FontSize', FSz, 'FontName', FN);
set(gca, 'FontSize', FSz, 'FontName', FN, 'LineWidth', 2, 'YTick', 1:5,'YTickLabel', {2.4; 6; 12; 24; 50})
axis([-45 100 ymn ymx])
box on;

break
figure(1); hold on;
legend('2.4m SWABS', '2.4m SHF', '6m SWABS', '6m SHF','12m SWABS', '12m SHF','24m SWABS', '24m SHF','50m SWABS', '50m SHF');
xlabel(' Latitude', 'FontSize', FSz, 'FontName', FN);
ylabel('Annual Mean Energy Flux to the atmosphere (W m^-^2)', 'FontSize', FSz, 'FontName', FN)
title('Annual mean atmospheric heating', 'FontSize', FSz, 'FontName', FN)
set(gca, 'FontSize', FSz, 'FontName', FN, 'LineWidth', 2)
axis([-90 90 0 250])
Box on

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
figure(2); hold on;
legend('2.4m SWABS', '2.4m SHF', '6m SWABS', '6m SHF','12m SWABS', '12m SHF','24m SWABS', '24m SHF','50m SWABS', '50m SHF');
xlabel(' Latitude', 'FontSize', FSz, 'FontName', FN);
ylabel({'Seasonal Amplitude Energy Flux'; 'to the atmosphere (W m^-^2)'}, 'FontSize', FSz, 'FontName', FN)
title('Seasonal atmospheric heating', 'FontSize', FSz, 'FontName', FN)
set(gca, 'FontSize', FSz, 'FontName', FN, 'LineWidth', 2)
axis([-90 90 -50 100])
hold on; plot([-90 90], [0 0], 'k', 'LineWidth', 2.5)
Box on

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
dtmn=0; dtmx=365; %date minnimum and maximum
figure(3); hold on;
legend('2.4m SWABS', '2.4m SHF', '6m SWABS', '6m SHF','12m SWABS', '12m SHF','24m SWABS', '24m SHF','50m SWABS', '50m SHF');
xlabel(' Month', 'FontSize', FSz, 'FontName', FN);
ylabel({'Seasonal Anomaly of extratropical'; 'energy flux  (W m^-^2)'}, 'FontSize', FSz, 'FontName', FN)
title('Seasonal phasing of extratropical atmospheric heating', 'FontSize', FSz, 'FontName', FN)
set(gca, 'FontSize', FSz, 'FontName', FN, 'LineWidth', 2)
axis([dtmn dtmx -75 120])
hold on; plot([dtmn dtmx], [0 0], 'k', 'LineWidth', 2.5)
Box on

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
FN='Helvetica';
figure(4); hold on;
%[AX,H1a,H2a] = plotyy(NaN,NaN,NaN,NaN);

%set(AX(1), 'YLim', [-200  200]);
%set(AX(2), 'YLim', [-200  200]);

%set(get(AX(1),'Ylabel'),'string',{'Seasonal Anomaly of extratropical'; 'energy flux  (W m^-^2)'}, 'FontSize', FSz, 'FontName', FN) %sets the labels
%set(AX(1),'XColor','k','YColor','k', 'FontSize', 20)
%set(get(AX(1),'XTick'), [])

%set(get(AX(2),'Ylabel'),'String',{'Tropospheric temperature'; '(anomaly from annual mean -- K)'}, 'Color', [0 0 0], 'FontSize', FSz, 'FontName', FN)
%set(AX(2),'XColor','k','YColor','k', 'FontSize', 20) 


legend('2.4m Heating ', '2.4 Temp.', '6m Heating', '6m Temp.','12m Heating', '12m Temp.','24 m Heating','24m Temp.','50m Heating', '50m Temp.' );
xlabel(' Day of year', 'FontSize', FSz, 'FontName', FN);
ylabel({'Seasonal Anomaly of extratropical'; 'energy flux  (W m^-^2)'}, 'FontSize', FSz, 'FontName', FN) 
title('Seasonal phasing of total extratropical atmospheric heating', 'FontSize', FSz, 'FontName', FN)
set(gca, 'FontSize', FSz, 'FontName', FN, 'LineWidth', 2)
axis([dtmn dtmx -100 180])
hold on; plot([dtmn dtmx], [0 0], 'k', 'LineWidth', 2.5)
Box on
