%start reading in ERA5 data to calculate transients
clear all; close all;
%%% preamble -- load the time invariant masks and weights
%climatoligcal masks
load('C:\Users\Aaron Donohoe\Desktop\research\ERA5_AHT\ERA5_scripts\CLIM_MASKS.mat', 'lat', 'lon', 'WEIGHT','WEIGHT2', 'ZON_NORM', 'ZON_NORM2', 'ZON_WEIGHT', 'ZON_WEIGHT2','TMASK', 'lev')


%%%%%%%%%%%%%%%%%%%%%%%%%%%%% go level be level
a=6.36*10^6;
dlat=mean(lat(1:end-1)-lat(2:end));
dY=a*dlat*pi/180;
g=9.81;
cp=1004;
L=2.51*10^6; %latent heat of evaporation in J/kg
[LON,LAT]=meshgrid(lon,lat); %mesh for calculations

mlev = [0; mean([lev(1:end-1,1) lev(2:end,1)]')'; 1000]; %vertical boundaries of levels
norm=(mlev(1:end-1) - mlev(2:end))/(max(mlev)-min(mlev)); %normalized vertical averaging vector based on the pressure level boundaries
weight=-(mlev(1:end-1) - mlev(2:end))*100/g; % in kg/m^2 



%%%%%%%%%%%%%%%%%%%%%%%% load instantaneous data
dir = {'D:\ERA5_AHT\raw_data\00\';
        'D:\ERA5_AHT\raw_data\06\';
       'D:\ERA5_AHT\raw_data\12\';
       'D:\ERA5_AHT\raw_data\18\';};
 
       120*7/60
   


YEARS=    [ 2022];% 2019 2020 2021 2022]% 2003 2004 2005 2011 2012 2013 2014 2015]; 
st_month = [  3    1    1    1    1];
for ny  = 1:length(YEARS);
    yr = YEARS(ny) 
daysperyear=365.2425;    
    for mn=1:12;%st_month(ny):1:12;
    off=0;
        for nt=1:4;%4; %hourly loop
        [yr mn nt]
        tic
        filt = [char(dir(nt,:)) num2str(yr) '.t.nc'];
        filv = [char(dir(nt,:)) num2str(yr) '.v.nc'];
        filq = [char(dir(nt,:)) num2str(yr) '.q.nc'];
        filz = [char(dir(nt,:)) num2str(yr) '.z.nc'];
        filu = [char(dir(nt,:)) num2str(yr) '.u.nc'];

        T1_V= ncread(filv, 'time',[1], [1]);
        T1_U= ncread(filu, 'time',[1], [1]);
        T1_T= ncread(filt, 'time',[1], [1]);
        T1_Q= ncread(filq, 'time',[1], [1]);
        T1_Z= ncread(filz, 'time',[1], [1]);

        time_differences = [T1_V T1_U T1_T T1_Q T1_Z] - T1_V

    
        lat = double(ncread(filt, 'latitude')); NN=length(lat);
        lon = double(ncread(filt, 'longitude')); N=length(lon);
        lev = double(ncread(filt, 'level')); NL=length(lev);
    
        time =  ncread(filt, 'time'); NT=length(time);
        year = 1900 + rounddown(time/(24*daysperyear));
        day = (time-((year-1900)*daysperyear*24))/24;

        %day_off=16;
        %cen = monthday(mn)+day_off;

        do = (monthday(mn)+1):1:monthday(mn+1);%max([1 cen-wid]):1:min([length(time) cen+wid]); %hopefullt this edit makes the total time series continous
        DO(mn, nt) = length(do);
    
        t(off+(1:length(do)), 1:NL, 1:NN,1:N) = permute(ncread(filt, 't', [ 1 1 1 min(do)], [Inf Inf Inf length(do)]), [4 3 2 1]);
        v(off+(1:length(do)), 1:NL, 1:NN,1:N)  = permute(ncread(filv, 'v', [ 1 1 1 min(do)], [Inf Inf Inf length(do)]), [4 3 2 1]);
        q(off+(1:length(do)), 1:NL, 1:NN,1:N) = permute(ncread(filq, 'q', [ 1 1 1 min(do)], [Inf Inf Inf length(do)]), [4 3 2 1]);
        z(off+(1:length(do)), 1:NL, 1:NN,1:N)  = permute(ncread(filz, 'z', [ 1 1 1 min(do)], [Inf Inf Inf length(do)]), [4 3 2 1]);
        u(off+(1:length(do)), 1:NL, 1:NN,1:N)  = permute(ncread(filu, 'u', [ 1 1 1 min(do)], [Inf Inf Inf length(do)]), [4 3 2 1]);
        
        %%%% now partition instant by instant
        [NT, NL,NN,N] = size(v); %deminsonality


 
        for it = 1:NT; %instants
            [it NT]
             tic
             SE = cp*squeeze(t(it,:,:,:));
             PE = squeeze(z(it,:,:,:)); %it's geopotential (no geopotential height)
             ME = L*squeeze(q(it,:,:,:));
             MSE = ME+SE+PE; %Moist static energy, used in the MMC heat transport calculation 
             V = squeeze(v(it,:,:,:)); %it's geopotential (no geopotential height)
             U = squeeze(u(it,:,:,:)); %it's geopotential (no geopotential height)
             
            %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%5            divergence %%
            
            for nl=1:NL;
                    divU(nl,1:NN,1:N) = sphdiv_per_stripped(squeeze(U(nl,:,:)).*squeeze(WEIGHT(nl,:,:)), squeeze(V(nl,:,:)).*squeeze(WEIGHT(nl,:,:)), lat,lon); %spherical diveregnce of mass flux at each gridpoint on the specified level kg/(s m^2)
            end
     
            %% LOCAL DEPARTURES FROM VERTICAL AVERAGE mse
            for nn=1:NN;
                  for n=1:N;
                        loc_E_bar(nn,n) = sum(squeeze(MSE(:,nn,n)).*squeeze(WEIGHT(:,nn,n)))/sum(squeeze(WEIGHT(:,nn,n))); %units of energy per mass (J/kg)
                        aE(1:NL,nn,n)= MSE(:,nn,n) - loc_E_bar(nn,n); %local departure of energy from the vertical average

                        %% and moisture
                        loc_ME_bar(nn,n) = sum(squeeze(ME(:,nn,n)).*squeeze(WEIGHT(:,nn,n)))/sum(squeeze(WEIGHT(:,nn,n))); %units of energy per mass (J/kg)
                        aME(1:NL,nn,n)= ME(:,nn,n) - loc_ME_bar(nn,n); %local departure of energy from the vertical average


                  end
            end

             
            divergent=squeeze(sum(divU.*aE)); %vertical integrate the full divergence term
            ME_divergent=squeeze(sum(divU.*aME)); %vertical integrate the full divergence term
                 
            %%%% advective
              dlat=lat(2)-lat(1); %latitude spacingZonal_div_
              dy=a*dlat*pi/180; %meridional grid spacing
              dlon=mean(diff(lon)); 
              e1=ones(N,1);
              DDX=spdiags([e1 -e1 e1 -e1], [-N+1 -1 1 N-1],N,N); %matrix operator for centered differencing x 
              e0=zeros(N,1);
              DDY=spdiags([-e1 e1], [0 2*N], N,3*N)/(2*dy);

              for nl=1:NL;
              for nn=2:NN-1
                ddx=DDX/(cosd(lat(nn))*(2*a*dlon*pi/180)); %divide the operator by twice the x spacing
                
               field=squeeze(MSE(nl,nn-1:nn+1,:)); %extended field for ddy ing
               Field=field(2,:); %middle row fo ddx
               E_x(nl,nn,1:N) = ddx * Field';
               E_y(nl,nn,1:N)= DDY * g2c(field);


                %moisture
               field=squeeze(ME(nl,nn-1:nn+1,:)); %extended field for ddy ing
               Field=field(2,:); %middle row fo ddx
               ME_x(nl,nn,1:N) = ddx * Field';
               ME_y(nl,nn,1:N)= DDY * g2c(field);

               %ZERO PAD POLES
               E_y(nl,NN,1:N)=zeros(1,N); E_x(nl,NN,1:N)=zeros(1,N);
               ME_y(nl,NN,1:N)=zeros(1,N); ME_x(nl,NN,1:N)=zeros(1,N);
              end
              end
               advect(1:NN,1:N) = squeeze(sum(WEIGHT.*U.*E_x + WEIGHT.*V.*E_y));
               advect_X(1:NN,1:N) = squeeze(sum(WEIGHT.*U.*E_x ));
               advect_Y(1:NN,1:N) = squeeze(sum(WEIGHT.*V.*E_y ));
               total(1:NN,1:N) = divergent(:,:) + advect(:,:);

               Madvect(1:NN,1:N) = squeeze(sum(WEIGHT.*U.*ME_x + WEIGHT.*V.*ME_y));
               Madvect_X(1:NN,1:N) = squeeze(sum(WEIGHT.*U.*ME_x ));
               Madvect_Y(1:NN,1:N) = squeeze(sum(WEIGHT.*V.*ME_y ));
               Mtotal(1:NN,1:N) = ME_divergent(:,:) + Madvect(:,:);
        
        
              %%% save for export
              iDIVERGENT(4*(it-1)+nt, 1:NN,1:N)=divergent;
              iADVECT_X(4*(it-1)+nt, 1:NN,1:N)=advect_X;
              iADVECT_Y(4*(it-1)+nt, 1:NN,1:N)=advect_Y;
              iTOTAL(4*(it-1)+nt, 1:NN,1:N)=total;
              iE_bar(4*(it-1)+nt, 1:NN,1:N) = loc_E_bar;

              %% moisture
              M_iDIVERGENT(4*(it-1)+nt, 1:NN,1:N)=ME_divergent;
              M_iADVECT_X(4*(it-1)+nt, 1:NN,1:N)=Madvect_X;
              M_iADVECT_Y(4*(it-1)+nt, 1:NN,1:N)=Madvect_Y;
              M_iTOTAL(4*(it-1)+nt, 1:NN,1:N)=Mtotal;
              M_iE_bar(4*(it-1)+nt, 1:NN,1:N) = loc_ME_bar;
              TIME(4*(it-1)+nt, 1:NN,1:N) = time(it);
              
              clear advect advect_X advect_Y total divergent Madvect Madvect_X Mavect_Y Mtotal ME_divergent V U ME LE MSE SE loc_E_bar loc_ME_bar;
              toc

            end %% instant loop -- it

            clear v t q z V T Q Z SE ME MSE PE
     

        end %6 hour Z
        save(['C:\Users\Aaron Donohoe\Desktop\research\ERA5_AHT\instant_AHT_2d\monthly\' num2str(YEARS(ny)) '_' num2str(mn) '_instant2d.mat'], 'lat','lon', 'lev','time', 'year', 'day', 'iE_bar', 'M_iE_bar', 'iDIVERGENT', 'iADVECT_X','iADVECT_Y', 'iTOTAL', 'M_iDIVERGENT', 'M_iADVECT_X','M_iADVECT_Y', 'M_iTOTAL' ) ;
         clear lat lon lev time year day iDIVERGENT iADVECT_Y iADVECT_X iTOTAL M_iDIVERGENT M_iADVECT_X M_iADVECT_Y M_iTOTAL TIME M_iE_bar iE_bar
         
    end %month

    %save

end %year