% 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
startyear=1987;
endyear=1991;
for yr=startyear:endyear;
dir = '/home/disk/eos10/dargan/aquafv/output/seas6mml/history/'; 
fil =[ num2str(yr), '0101.atmos_month.nc'];
FIL=[dir fil];

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

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

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%  TOA %%%%%%%%%%%%%%%%%%%%%%%%%%
OLR(yr-startyear+1,1:12,1:NN,1:N) = nc{'olr'}(:,:,:); %TOA OLR
SWD(yr-startyear+1,1:12,1:NN,1:N) = nc{'swdn_toa'}(:,:,:); %TOA SW down
SWU(yr-startyear+1,1:12,1:NN,1:N) = nc{'swup_toa'}(:,:,:); %TOA SW up
ASR(yr-startyear+1,1:12,1:NN, 1:N) = SWD(yr-startyear+1,1:12,1:NN,1:N) -SWU(yr-startyear+1,1:12,1:NN,1:N); %seem right? right!
NET(yr-startyear+1,1:12,1:NN,1:N) = nc{'netrad_toa'} (:,:,:) ; %net radition, reality check says this should equal ASR - OLR - CHECKS OUT KOSHER

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

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

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



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

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


%sphcolor2(mean(SW_S_NET), lat,lon,5); title('SW SURFACE HEAT FLUX _ ANNUAL MEAN')
%sphcolor2(mean(LW_S_NET), lat,lon,6); title('LW SURFACE HEAT FLUX - ANNUAL MEAN')
%sphcolor2(mean(SENS), lat,lon,7); title('SENSIBLE SURFACE HEAT FLUX - ANNUAL MEAN')
%sphcolor2(mean(LH_FLX), lat,lon,8); title('LATENT SURFACE HEAT FLUX - ANNUAL MEAN')
%sphcolor2(mean(FS), lat,lon,9); title('TOTAL SURFACE HEAT FLUX- ANNUAL MEAN')

%annual mean budgets
%SWG=trapz(sin(lat*pi/180),(mean(squeeze(mean(SW_S_NET(:,:,:)))'))')/2
%LWG=trapz(sin(lat*pi/180),(mean(squeeze(mean(LW_S_NET(:,:,:)))'))')/2
%SHG=trapz(sin(lat*pi/180),(mean(squeeze(mean(SENS(:,:,:)))'))')/2
%LHG=trapz(sin(lat*pi/180),(mean(squeeze(mean(LH_FLX(:,:,:)))'))')/2

%load the column tendency
load ('TTEN.mat')

%diagnose the heat transport divergence as a residual
TEDIV(yr-startyear+1,1:12,1:NN,1:N) = TTEN_bar - squeeze(NET(yr-startyear+1,1:12,1:NN,1:N)) - squeeze(FS(yr-startyear+1,1:12,1:NN,1:N)); %convergence of total energy transport
end




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));

%monthly global mean
for mn = 1:12; 
    G_FS(mn) = trapz(sin(lat*pi/180),squeeze(mean(squeeze(FS(mn,:,:))')'))/2;
    G_NET(mn) = trapz(sin(lat*pi/180),squeeze(mean(squeeze(ASR(mn,:,:) - OLR(mn,:,:))')'))/2;
    G_TEN(mn) = trapz(sin(lat*pi/180),squeeze(mean(squeeze(TTEN_bar(mn,:,:))')'))/2;
end


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));

TTEN=TTEN_bar;

DT =squeeze(mean(SWD));
UT =squeeze(mean(SWU));
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')

