%This script is designed to calculate the AB omega over the N. Hemi. from
%ERA interim data

close all
clear all

%add a bunch of paths to make sure netcdf handling is easier
addpath('/uufs/chpc.utah.edu/common/home/u0559900/MATLAB_SCRIPTS/');
addpath('/uufs/chpc.utah.edu/common/home/u0559900/MATLAB_SCRIPTS/mexcdf/mexnc/')
addpath('/uufs/chpc.utah.edu/common/home/u0559900/MATLAB_SCRIPTS/mexcdf/snctools');
addpath('/uufs/chpc.utah.edu/common/home/u0559900/MATLAB_SCRIPTS/netcdfAll-4.2.jar');

%MAP STUFF
load conus;
S=shaperead('landareas','UseGeoCoords',true);


%initialize a location for the final data
tracksum=zeros(61,280,7);

mns=['oct'; 'nov'; 'dec' ;'jan' ;'feb'; 'mar'; 'apr'];
yr1=2010; %year for the ending month
for mn=1:7
    if mn<4
        yr=yr1-1;
    else
        yr=yr1;
    end
    yrs=yr-1989;
    
    %file='output_feb_2011_alllevs.nc';
    file=strcat('output_',mns(mn,:),'_',num2str(yr),'.nc')

%%

%{
dimensions:
	longitude = 240 ;
	latitude = 40 ;
	levelist = 7 ;
	time = UNLIMITED ; (12 currently)

variables:
	// Preference 'PRESERVE_FVD':  false,
	// dimensions consistent with ncBrowse, not with native MATLAB netcdf package.
	single longitude(longitude), shape = [240]
		longitude:units = "degrees_east"
		longitude:long_name = "longitude"
	single latitude(latitude), shape = [40]
		latitude:units = "degrees_north"
		latitude:long_name = "latitude"
	int32 levelist(levelist), shape = [7]
		levelist:units = "millibars"
		levelist:long_name = "pressure_level"
	int32 time(time), shape = [12]
		time:units = "hours since 1900-01-01 00:00:0.0"
		time:long_name = "time"
	int16 r(time,levelist,latitude,longitude), shape = [12 7 40 240]
		r:scale_factor = 0.00234828
		r:add_offset = 61.7188
		r:_FillValue = -32767 s
		r:missing_value = -32767 s
		r:units = "%"
		r:long_name = "Relative humidity"
	int16 u(time,levelist,latitude,longitude), shape = [12 7 40 240]
		u:scale_factor = 0.00181751
		u:add_offset = 19.9043
		u:_FillValue = -32767 s
		u:missing_value = -32767 s
		u:units = "m s**-1"
		u:long_name = "U velocity"
	int16 v(time,levelist,latitude,longitude), shape = [12 7 40 240]
		v:scale_factor = 0.00187424
		v:add_offset = -1.34212
		v:_FillValue = -32767 s
		v:missing_value = -32767 s
		v:units = "m s**-1"
		v:long_name = "V velocity"
	int16 vo(time,levelist,latitude,longitude), shape = [12 7 40 240]
		vo:scale_factor = 8.7408e-09
		vo:add_offset = 9.34047e-05
		vo:_FillValue = -32767 s
		vo:missing_value = -32767 s
		vo:units = "s**-1"
		vo:long_name = "Vorticity (relative)"
	int16 z(time,levelist,latitude,longitude), shape = [12 7 40 240]
		z:scale_factor = 1.68581
		z:add_offset = 65945.4
		z:_FillValue = -32767 s
		z:missing_value = -32767 s
		z:units = "m**2 s**-2"
		z:long_name = "Geopotential"

//global Attributes:
		:Conventions = "CF-1.0"
		:history = "2011-08-26 17:56:59 GMT by mars2netcdf-0.92"


%}

%% Pull in the Data from the NETCDF files...
    
time=double(nc_varget(file,'time'))./24 + datenum(1900,01,01,00,00,00); %convert to matlab serial time
lons=double(nc_varget(file,'longitude')); %rereference the lons;
tidx=find(time>datenum(2051,01,01,00,00,00)); %old stuff to deal with a problem
time(tidx)=[];


lonloc1=find(lons==178.5);lonloc2=find(lons==-180);
lons1=lons(lonloc2:lonloc1);
lons1(241)=180;
lons=(-180-30):1.5:(178.5+30);
px=length(lons); %number of points in the zonal direction

lats=double(nc_varget(file,'latitude'));
latloc1=find(lats==0);latloc2=find(lats==90)
py=length(lats); %number of points meridionally
lats=lats(latloc2:latloc1);
rvort=double(nc_varget(file,'vo'));
rvort(tidx,:,:,:)=[];
leveler1=size(rvort,2);

rvort1=rvort(:,:,latloc2:latloc1,lonloc2:lonloc1);
rvort=zeros(length(time),leveler1,61,280);
rvort(:,:,:,21:260)=rvort1;
rvort(:,:,:,1:20)=rvort1(:,:,:,221:240);
rvort(:,:,:,261:280)=rvort1(:,:,:,1:20);

hgt=double(nc_varget(file,'z'))./9.8;
hgt(tidx,:,:,:)=[];
hgt1=hgt(:,:,latloc2:latloc1,lonloc2:lonloc1);
hgt=zeros(length(time),leveler1,61,280);
hgt(:,:,:,21:260)=hgt1;
hgt(:,:,:,1:20)=hgt1(:,:,:,221:240);
hgt(:,:,:,261:280)=hgt1(:,:,:,1:20);




temps=double(nc_varget(file,'t'));
temps(tidx,:,:,:)=[];
temps1=temps(:,:,latloc2:latloc1,lonloc2:lonloc1);
temps=zeros(length(time),leveler1,61,280);
temps(:,:,:,21:260)=temps1;
temps(:,:,:,1:20)=temps1(:,:,:,221:240);
temps(:,:,:,261:280)=temps1(:,:,:,1:20);


levels=double(nc_varget(file,'levelist'))*10^2;
[tsteps,leveler,ltdim,lndim]=size(rvort);
R=297;

ompmat=zeros(py,px,leveler,tsteps);
[latmat,lonmat]=meshgrat(lats,lons);

ommonth=zeros(py,px,tsteps);
%% set up the x,y grids for gradients
hgtfield=squeeze(temps(1,4,:,:));

%stuff from gradientm function
geoid = almanac('earth','geoid','m');
mstruct = defaultm('pcarree');
mstruct.geoid = geoid;
mstruct = defaultm(mstruct);

[x,y] = mfwdtran(mstruct,latmat,lonmat); %convert to cartesian points

%handle the wrap around lat lon data (goes past -180/180)
for jk=1:20;
   x(:,21-jk)=x(:,21)-jk*(1.6698*10^5); 
   x(:,260+jk)=x(:,260)+jk*(1.6698*10^5);
end


%convergence of the grid
convfactor = departure(zeros(size(latmat)),ones(size(latmat)),latmat,geoid)/ departure(0,1,0,geoid);

convfactor(convfactor<10^-16) = NaN; % avoid divide by zeros



%%

rearth=6.37*10^6; %earth radius
phi=1.5*(pi/180); %radian grid spacing
dxmat=rearth.*cosd(latmat).*phi; %matrix of x spacing
dy=rearth.*phi; %yspacing
dymat=ones(size(lonmat)).*dy; %matrix of yspacing

latmat2=latmat(2:end-1,2:end-1); %remove the boundary points
dxmat2=dxmat(2:end-1,2:end-1);
dymat2=dymat(2:end-1,2:end-1);
amat=(dxmat2.*dymat2);

amatflat=amat(:);



latmat2=latmat(2:end-1,2:end-1);
dxmat2=dxmat(2:end-1,2:end-1);
dx2mat2=dxmat2.^2;
dx2mat2=dx2mat2(:);
dymat2=dymat(2:end-1,2:end-1);
dy2mat2=dymat2.^2;
dy2mat2=dy2mat2(:);


%% Now we need to get the stream function and non-divergent winds
 
%do it in parallel
matlabpool open 5

tic

parfor i=1:tsteps
    
    for ii=1:leveler
        hgtnow=squeeze(hgt(i,ii,:,:));
        rnow=squeeze(rvort(i,ii,2:end-1,2:end-1));
        rnow(:,1)=rnow(:,1)./hgtnow(2:end-1,2);
        rnow(1,2:end-1)=rnow(1,2:end-1)./hgtnow(1,3:end-2);
        rnow(:,end)=rnow(:,end-1)./hgtnow(2:end-1,end-1);
        rnow(end,2:end-1)=rnow(end,2:end-1)./hgtnow(end-1,3:end-2);
        
        
        %rnow=squeeze(rvort(i,ii,2:end-1,2:end-1));
        rnow=rnow(:);
        nx = py-2; % number of grid points in the x-direction;
        ny = px-2; % number of grid points in the y-direction;
        ex = ones(nx,1);
        Dxx = -spdiags([ex -2*ex ex], [-1 0 1], nx, nx); %1D discrete Laplacian in the x-direction ;
        [Bxx,dxx]=spdiags(Dxx);
        
        
        
        ey = ones(ny,1);
        Dyy = -spdiags([ey, -2*ey ey], [-1 0 1], ny, ny); %1D discrete Laplacian in the y-direction ;
        [Byy,dyy]=spdiags(Dyy);
        for kk=1:length(Bxx)
            Bxx(kk,:)=Bxx(kk,:)./dx2mat2(kk);
            
        end
        for kk=1:length(Byy);
            Byy(kk,:)=Byy(kk,:)./ dy2mat2(1);
        end
        
        Dxx=spdiags(Bxx,dxx,Dxx);
        Dyy=spdiags(Byy,dyy,Dyy);
        L = kron(speye(ny, ny), Dxx) + kron(Dyy, speye(nx, nx));
        
        
        U=L\rnow;
        
        psii=reshape(U,py-2,px-2);
        psii=padarray(psii,[1 1]);
        
        %non-divergent wind:
        [dpsiidx,dpsiidy]=mtxgradient_neil(psii,x,y);
        upsii=dpsiidy;
        vpsii=-dpsiidx./convfactor;
        
        if ii==-3
            figure(5)
            clf
            usamap([21 75],[-180 -10]);
            hold on;
            %contourm(lats,lons,10^-7*psii,-10:.5:10,'linewidth',2);
            %caxis([-10 10]);
            contourm(lats,lons,hgtnow,0:100:50000,'-k','linewidth',2)
            quiverm(latmat,lonmat,vpsii,upsii*2);
            geoshow(statelat,statelon,'Displaytype','line','color','k');
            geoshow(uslat,uslon,'Displaytype','line','color','k');
            geoshow([S.Lat],[S.Lon]);
            %caxis([-10 10])
            
        end
        [dtdx,dtdy] = mtxgradient_neil(squeeze(temps(i,ii,:,:)),x,y);
        %dtdx=dtdx./convfactor;
        %[dudx,dudy]=gradient(u,166.77*10^3);
        [dudx,dudy] = mtxgradient_neil(upsii,x,y);
        %dudx=dudx./convfactor;
        %[dvdx,dvdy]=gradient(v,166.77*10^3);
        [dvdx,dvdy] = mtxgradient_neil(vpsii,x,y);
        %dvdx=dvdx./convfactor;
        qx=(-R./levels(ii)).*((dudx.*dtdx)+(dvdx.*dtdy));
        qy=(-R./levels(ii)).*((dudy.*dtdx)+(dvdy.*dtdy));
        %[dqxdx,~]=gradient(qx,166.77*10^3);
        [dqxdx,~] = mtxgradient_neil(qx,x,y);
        %dqxdx=dqxdx./convfactor;
        %[~,dqydy]=gradient(qy,166.77*10^3);
        [~,dqydy] = mtxgradient_neil(qy,x,y);
        omproxy=-2*(dqxdx+dqydy);
        if ii==-3
            figure(5);
            hold on
            pcolor(omproxy);shading flat;
            %caxis([-.7*10^-16 .7*10^-16]);
            %
            %         colorbar
        end
        ompmat(:,:,ii,i)=omproxy(:,:);
        
    end
    %end
end
matlabpool close

toc

%taper the grid boundaries

for jj=1:20
    ompmat(:,281-jj,:,:)=ompmat(:,281-jj,:,:).*(jj./20);
    ompmat(:,jj,:,:)=ompmat(:,jj,:,:).*(jj./20);
end

%% Now calculate the whole shebang

% first extend spatial matrices to 3-D
latmat3=repmat(latmat(2:end-1,2:end-1),[1 1 (leveler-2)]);
dxmat3=repmat(dxmat(2:end-1,2:end-1),[1 1 (leveler-2)]);


dymat3=repmat(dymat(2:end-1,2:end-1),[1 1 (leveler-2)]);
dx2mat3=dxmat3.^2;
dx2mat3=dx2mat3(:);
dy2mat3=dymat3.^2;
dy2mat3=dy2mat3(:);
ffact=2*(7.292*10^-5).*sind(latmat3);
temps=permute(temps,[3,4,2,1]);
asqr=(160*10^3).^2; %~150 km grid
dpsqr=(100 *10^2).^2; % 100 hpa in the vertical
%%
fflat=ffact(:).^2;
fflat=(fflat*0)+(10^-4)^2;%comment this line out to use variable coriolis forcing

%% now compute the static stability

clear dthtdp;
[~,~,dthtdp,~]=gradient(temps,-10000);%returns vertical gradient of potential temperature
dthtdp(dthtdp>=0)=-.0000001; %should be negative... otherwise it is unstable

clear sigma;
sigma=zeros(size(dthtdp));

for j=1:leveler;
    sigma(:,:,j,:)=(-287./100000).*(100000./levels(j)).^(717/1004).*dthtdp(:,:,j,:);
end

%% set up the laplacian matrix to be inverted later
ny=px-2;
nx=py-2;
nz=leveler-2;
[~,~,A1]=laplacian([nx ny nz]); %this is a use generated function that uses kronecker sums of the individual 1-D laplacian matrices


    scrz=get(0,'screensize');
    f=figure(500);
    set(f,'color','w')
    
    
    
    cmap=[gray(16);flipud(gray(16))];
    cmap(1:13,1)=1;
    cmap(14,:)=[1 1 1];
    cmap(15,:)=[1 1 1];
    
    cmap(16,:)=[1 1 1];
    cmap(17,:)=[1 1 1];
    cmap(18,:)=[1 1 1];
    cmap(19,:)=[1 1 1];
    cmap(20:end,3)=1;
    colormap(cmap);

%% Now compute the AB omega 

matlabpool open 5
tic
parfor i=1:tsteps
    %if time(i)>datenum(2002,04,14,00,00,00) && time(i)<datenum(2002,04,18,00,00,00)
    A=A1; %reset A to its base value
    
    
    sigmamat=sigma(2:end-1,2:end-1,2:end-1,i);
    sgflat=(sigmamat(:)*1);
    
    %retrieve the diagnols of A
    [B,d]=spdiags(A);
    
    %now update the poins along the diagnols with appropriate grid spacings
    %and constants
    
    for k=1:length(B)
        
        B(k,1)=B(k,1).*(fflat(k)./dpsqr);
        B(k,7)=B(k,7).*(fflat(k)./dpsqr);
        B(k,2)=B(k,2).*(sgflat(k)./dx2mat3(k));
        B(k,6)=B(k,6).*(sgflat(k)./dx2mat3(k));
        B(k,4)=(2*sgflat(k)./dx2mat3(k))+(2*sgflat(k)./dy2mat3(1))+(2*fflat(k)./dpsqr);
        
    end
    B(:,3)=B(:,3).*(sgflat./dy2mat3(1));
    B(:,5)=B(:,5).*(sgflat./dy2mat3(1));
    
    %replace the diagnols of A
    A=spdiags(B,d,A);
    
    %this is the RHS term
    omnow=squeeze(ompmat(2:end-1,2:end-1,2:end-1,i));
    omnow=(omnow(:));
    omnow(isnan(omnow))=0; %check for nans... cause they'll screw it up
    
    
    %now solve with the sparse left matrix divide
    U=A\omnow;
    om=reshape(U,py-2,px-2,leveler-2);
    om=padarray(om,[1 1 1]);
    
    %optional plot for animations
%     figure(500);
%     clf;
%     axesm('eqaazim','MapLatLimit',[0 90])
%     axis off; framem on; gridm on; mlabel on; plabel on;
%     setm(gca,'MLabelParallel',-110)      
%     %usamap([20 75],[-180 180]);
%     hold on;
%     pcolorm(lats,lons1,sum(om(:,21:261,:),3));%shading flat;
%     caxis([-1.5 1.5])
%     contourm((lats),lons1,squeeze(hgt(i,5,:,21:261)),0:50:8000,'-k','linewidth',1.5);
%     %contourm((lats),lons,squeeze(hgt(i,end,:,:)),0:50:8000,'-m','linewidth',1);
%     geoshow(statelat,statelon,'Displaytype','line','color','k');
%     geoshow(uslat,uslon,'Displaytype','line','color','k');
%     geoshow([S.Lat],[S.Lon]);
      %end
    
      %aribitrary threshold value to use just strong ascent
      om(om<0.05)=0;
      om=sum(om,3);
      ommonth(:,:,i)=om;
      %tracksum(:,:,mn)=tracksum(:,:,mn)+om;
end
matlabpool close;
tracksum(:,:,mn)=sum(ommonth,3); %save it in monthly format;

toc
%% plot the monthly track
    
    f=figure(500+mn);
    set(f,'pos',[scrz(1) scrz(2) scrz(3)*.5 scrz(4)],'color','w')
    
    
    cmap=flipud(gray(12));

    colormap(cmap);

    clf;
    axesm('eqaazim','MapLatLimit',[0 90])
    axis off; framem on; gridm on; mlabel on; plabel off;
    setm(gca,'MLabelParallel',-110)      
    hold on;
    %pcolorm(lats,lons1,squeeze(tracksum(:,21:261,mn)));%shading flat;
    contourfm(lats,lons1,squeeze(tracksum(:,21:261,mn)),0:1:70,'linestyle','none');
    caxis([0 60]);
    geoshow([S.Lat],[S.Lon],'color','k');
    tstr=strcat(mns(mn,:),'  ',num2str(yr));
    title(tstr,'fontweight','bold','fontsize',18);
    %title('Feb 2011 all-levels ');
    export_fig(sprintf('%s%s.png', mns(mn,:),num2str(yr)),'-m2.5');
end

%% now plot the annual track
ansum=sum(tracksum,3);
 f=figure(500+10);
    set(f,'pos',[scrz(1) scrz(2) scrz(3)*.5 scrz(4)],'color','w')
    
    
    cmap=flipud(gray(12));

    colormap(cmap);

    clf;
    axesm('eqaazim','MapLatLimit',[0 90])
    axis off; framem on; gridm on; mlabel on; plabel off;
    setm(gca,'MLabelParallel',-110)      
    hold on;
    pcolorm(lats,lons1,ansum(:,21:261));%shading flat;
    %contourfm(lats,lons1,squeeze(tracksum(:,21:261,mn)),0:1:180,'linestyle','none');
    %caxis([0 140]);
    geoshow([S.Lat],[S.Lon],'color','k');
    tstr=strcat(mns(mn,:),'  ',num2str(yr));
    %title(tstr,'fontweight','bold','fontsize',18);
    title('2009-2010 Storm Track');
    export_fig(sprintf('%s.png',num2str(yr)),'-m2.5');