% compute the toa atmospheric energy budegt for the 24 mm slab ocean
% aquaplanet
clear all; close all;

%constants
a=6.36*10^6;
g=9.81;
cp=1004;
L=2.51*10^6; %latent heat of evaporation in J/kg

%model data
load ~aaron/research/aqua_gcm/comp_co2/years  %has the start and end years
for yr=startyear:5:endyear;
    tic
    yr
dir = '/home/disk/eos10/dargan/aquafv/output/seas50mml/history/'; 
fil =[ num2str(yr), '0101.atmos_month.nc'];
FIL=[dir fil];
FIL=[dir fil];

nc=netcdf(FIL);
%ncdump(FIL);

lon = nc{'lon'}(:,:,:); N=length(lon);
lat = nc{'lat'}(:,:,:); NN=length(lat);
time = nc{'time'}(:); NT=length(time);

daysperyear=365;
yr=1982+rounddown(time/daysperyear);
day=time-(yr-1982)*daysperyear;
month=round((day+15)/(daysperyear/12));

plev = nc{'pfull'}(:); NL=length(plev);
pb=[0 mean([plev(2:end) plev(1:end-1)]') 1000]; %boundaries of each pressure level
w = (pb(2:end)-pb(1:end-1))/(max(pb)-min(pb))'; %mass weighting vector for vertical integrals

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%  TOA %%%%%%%%%%%%%%%%%%%%%%%%%%
data(1:12, 1:NL,1:NN) = squeeze(mean(shiftdim(nc{'mse'}(:,:,:),3))); %zonal mean mse
data2(1:12, 1:NL,1:NN) = squeeze(mean(shiftdim(nc{'dse'}(:,:,:),3))); %zonal mean dse

 for mn=1:12
     d=squeeze(data(mn,:,:));
     d2=squeeze(data2(mn,:,:));
     MSE(yr-startyear+1,mn,1:NN) = (w * d) * (10^5/9.81); %vertically integrated MSE in J/m^2;
     DSE(yr-startyear+1,mn,1:NN) = (w * d2) * (10^5/9.81); %vertically integrated MSE in J/m^2;
     %GMSE(yr-startyear+1,mn)=2*pi*(a^2)*trapz(sin(lat*pi/180), squeeze(MSE(mn,:))); %global mean MSE in J
     %GDSE(yr-startyear+1,mn)=2*pi*(a^2)*trapz(sin(lat*pi/180), squeeze(DSE(mn,:))); %global mean MSE in J
 end

 
 
%  figure(3); clf;
%  plot(lat, MSE);
%  legend('1','2','3','4','5','6','7','8','9','10','11','12')

end
MSE=squeeze(mean(GSE));
DSE=squeeze(mean(DSE));
 

MSE_TEN=(MSE([2:12 1],:) - MSE([12 1:11],:))/ ( 30 * 24 * 3600); %central difference
DSE_TEN=(DSE([2:12 1],:) - DSE([12 1:11],:))/ ( 30 * 24 * 3600); %central difference

save('MSE_TEN', 'lat','MSE_TEN', 'DSE_TEN')
 
    for yrr=yr:1:yr+4; %loop over individual years
        doy=find(yr == yrr) %indexes of the months in the given year
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%  TOA %%%%%%%%%%%%%%%%%%%%%%%%%%
        OLR(yrr-startyear+1,1:12,1:NN,1:N) = nc{'olr'}(doy,:,:); %TOA OLR
        SWD(yrr-startyear+1,1:12,1:NN,1:N) = nc{'swdn_toa'}(doy,:,:); %TOA SW down
        SWU(yrr-startyear+1,1:12,1:NN,1:N) = nc{'swup_toa'}(doy,:,:); %TOA SW up
        ASR(yrr-startyear+1,1:12,1:NN, 1:N) = SWD(yrr-startyear+1,1:12,1:NN,1:N) -SWU(yrr-startyear+1,1:12,1:NN,1:N); %seem right? right!
        NET(yrr-startyear+1,1:12,1:NN,1:N) = nc{'netrad_toa'} (doy,:,:) ; %net radition, reality check says this should equal ASR - OLR - CHECKS OUT KOSHER

        %surface radiation %%%%%%%%%%%%%%%%%%%%%%%
        LW_S_down(yrr-startyear+1,1:12,1:NN,1:N) = nc{'lwdn_sfc'}(doy,:,:); %downward longwave flux at surface
        LW_S_up(yrr-startyear+1,1:12, 1:NN, 1:N) = nc{'lwup_sfc'}(doy,:,:); %upward longwave flux at surface

        ALB=nc{'alb_sfc'}(:,:,:); %surface albedo 
        SW_S_down(yrr-startyear+1,1:12,1:NN,1:N) = nc{'swdn_sfc'}(doy,:,:); %downward SW at surface
        SW_S_up(yrr-startyear+1,1:12,1:NN,1:N) = nc{'swup_sfc'}(doy,:,:); %upward SW at surface

        LW_S_NET(yrr-startyear+1,1:12,1:NN,1:N) = LW_S_up(yrr-startyear+1,1:12,1:NN,1:N) - LW_S_down(yrr-startyear+1,1:12,1:NN,1:N); %to the atmosphere 
        LW_FLX(yrr-startyear+1,1:12,1:NN,1:N) = -nc{'lwflx'}(doy,:,:); %should equal the above - the define to the surface, negative sign makes it to the atmosphere
        SW_S_NET(yrr-startyear+1,1:12,1:NN,1:N) = SW_S_up(yrr-startyear+1,1:12,1:NN,1:N) - SW_S_down(yrr-startyear+1,1:12,1:NN,1:N); %SW surface flux to the atmosphere

        %surface sensible and latent
        SENS(yrr-startyear+1,1:12,1:NN,1:N) = nc{'shflx'}(doy,:,:); %sensible heat flux % to the atmosphere W/m^2
        P = nc{'precip'}(doy,:,:); %precipitation in kg/(m^2*s)
        E = nc{'evap'}(doy,:,:); %evaporation in kg/(m^2*s)
        LH_FLX(yrr-startyear+1,1:12,1:NN,1:N)=L*(E); %latent heat flux to the atmosphere in W/m^2

        FS(yrr-startyear+1,1:12,1:NN,1:N) = LH_FLX(yrr-startyear+1,1:12,1:NN,1:N) + SW_S_NET(yrr-startyear+1,1:12,1:NN,1:N) + LW_S_NET(yrr-startyear+1,1:12,1:NN,1:N) + SENS(yrr-startyear+1,1:12,1:NN,1:N);  % surface heat flux, to the atmosphere in W/m^2 

        load CTEN.mat;
        %diagnose the heat transport divergence as a residual
        TEDIV(yrr-startyear+1,1:12,1:NN,1:N) = CTEN - squeeze(NET(yrr-startyear+1,1:12,1:NN,1:N)) - squeeze(FS(yrr-startyear+1,1:12,1:NN,1:N)); %convergence of total energy transport
        toc
    end
end


% m_LH_FLX=squeeze(mean(LH_FLX));
% m_SH_FLX=squeeze(mean(SENS));
% m_FS=squeeze(mean(FS));
% m_ASR=squeeze(mean(ASR));
% m_OLR=squeeze(mean(OLR));
% m_TEDIV =squeeze(mean(TEDIV));
% 
% m_LW_S_down=squeeze(mean(LW_S_down));
% m_LW_S_up=squeeze(mean(LW_S_up));
% m_SW_S_down=squeeze(mean(SW_S_down));
% m_SW_S_up=squeeze(mean(SW_S_up));
% m_SW_up=squeeze(mean(SWU)); 
% m_SW_down=squeeze(mean(SWD));


LH_FLX=squeeze(mean(LH_FLX));
SH_FLX=squeeze(mean(SENS));
FS=squeeze(mean(FS));
ASR=squeeze(mean(ASR));
OLR=squeeze(mean(OLR));
TEDIV =squeeze(mean(TEDIV));

LW_S_down=squeeze(mean(LW_S_down));
LW_S_up=squeeze(mean(LW_S_up));
%SW_S_down=squeeze(mean(SW_S_down));
%SW_S_up=squeeze(mean(SW_S_up));
SW_up=squeeze(mean(SWU)); 
SW_down=squeeze(mean(SWD));

DT =squeeze(mean(SWD));
UT =squeeze(mean(SWU));

load('TTEN.mat')
TTEN=TTEN_bar;
DS=squeeze(mean(SW_S_down)); %downward shortwave at the surface
US=squeeze(mean(SW_S_up)); %downward shortwave at the surface

save('energy_climatology', 'DT', 'UT','DS', 'US','LH_FLX', 'SH_FLX', 'FS', 'ASR', 'OLR', 'LW_S_down', 'LW_S_up', 'TTEN','TEDIV', 'lat', 'lon')