%set up masks from the annual mean pressure
% to be used for all tiume steps thereafter
%load the mean fields
clear all; close all;
load('C:\Users\Aaron Donohoe\OneDrive - UW\Desktop\research\ERA5_AHT\ERA5_scripts\CLIM_PS_all_years.mat', 'lat', 'lon', "CPS", 'lev')

g=9.81;

[LON,LAT]=meshgrid(lon,lat); %mesh for calculations

%make normalization and weighting vectors (for vertical integrals)
MPS=CPS/100;
plev=lev; % already in hPa % /100
NL=length(plev);
NN=length(lat);
N=length(lon); 

%set up the topography mask
tmask=ones(NL,NN,N); %pre-allocate for a tmask of ones and zeros
TMASK=ones(NL,NN,N); %pre-allocate for a tmask of ones and NaNs

tic
for nn=1:NN
    for n=1:N
        mlev = [0; mean([plev(1:end-1,1) plev(2:end,1)]')'; max([squeeze(MPS(nn,n)) 1000 ]) ]; %vertical boundaries of levels

        
        %WEIGHT(1:NL,nn,n)=max([(min([Plev (MPS(nn,n)/100)*ones(NL,1)]')-[Plev(2:end)' 0]); zeros(1,NL)])*100/g; %the local weights for each discrete level based on the surface pressure
        WEIGHT(1:NL,nn,n)= max([(min([mlev(2:end) (MPS(nn,n))*ones(NL,1)]')' - mlev(1:end-1)) zeros(NL,1)]')*100/g;
        
    end
end

return

for nl=1:NL
    for nn=1:NN;
        do=find(squeeze(MPS(nn,:)) < (plev(nl)) ); %an error in the original script had a greater than sign
        tmask(nl,nn,do)=0;
        TMASK(nl,nn,do)=NaN;
        clear do
    end
end

WEIGHT2=WEIGHT.*tmask; %blokcs out partial levels

for nn=1:NN;    
    for nl=1:NL
        ZON_NORM(nl,nn,1:N)=WEIGHT(nl,nn,:)/(max([sum(WEIGHT(nl,nn,:)) 1*10^-5])); % allows weighted zonal averages at each level to be computed with a matric multiplication
        ZON_NORM2(nl,nn,1:N)=WEIGHT2(nl,nn,:)/(max([sum(WEIGHT2(nl,nn,:)) 1*10^-5])); % allows weighted zonal averages at each level to be computed with a matric multiplication
    end
end

%make a plot of the zonally averaged weight meridional cross section
for nl=1:NL
    ZON_WEIGHT(nl,1:NN)=mean(squeeze(WEIGHT(nl,:,:))'); % / weight(nl) if you want it as a fraction
    ZON_WEIGHT2(nl,1:NN)=mean(squeeze((WEIGHT(nl,:,:)).*tmask(nl,:,:))'); %weighting that eliminates partial levels
end

save('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', 'CPS')

%% check that sum of weights at ach grid point equals PS
sphcolor2jet(squeeze(sum(WEIGHT))*g - 100*MPS  ,lat,lon,3);
title("sum of weights test")

% look at cross sections of zonal weights
[LAT,LEV] = meshgrid(lat, plev);
figure(2); clf;
pcolor(LAT,-LEV, ZON_WEIGHT); colorbar;
