%% Trying to explain the change in climatological U
%  f*Ua = -(\bar{v'v'})_y; (proved to be dominant balance in Hoskins 1983) 
%  f*Ug = -g*Z_y;
%  f*dU = -g*Z_y - (dv'v')_y - 1/tau*dV
%  1/tau << f
%  f*dU = -g*Z_y - (dv'v')_y
%    |        |        |
%  dTERM1   dTERM2   dTERM3

Exps={'WA1Jan35_EQ1M40_','WA1Jan35_EQ1M40lon60180_','WA1Jan35_EQ1M40lon180300_','WA1Jan35_EQ1M40lon30060_'};
Amps=[0,1,2,3,4,5,6,7,10];
cA={'0','1','2','3','4','5','6','7','10'};
NExps=length(Exps);
colors={'k','r','b','g'};
N=length(Amps);
yijing=true;

latshow=40.;
levshow=250.;

d2r=pi/180.;
a=6371000.;
f0=4*pi/86400;
g=9.8;

%% calculate 3D EP, Eddy Energy Gain, and dU
if ~yijing
for iExp=1:NExps
    for icase=1:N
        file=char(strcat('vinth_Alls/vinth_All_',Exps(iExp),cA(icase),'.nc'));
        U=squeeze(double(ncread(file,'U')));
        V=squeeze(double(ncread(file,'V')));
        VV=squeeze(double(ncread(file,'VV')));
        Z=squeeze(double(ncread(file,'Z3')));    
        vv=VV-V.^2;
        lat=squeeze(double(ncread(file,'lat')));
        lon=squeeze(double(ncread(file,'lon')));
        lev=squeeze(double(ncread(file,'lev'))); 
        file=char(strcat('proj_eddys/IDEAL_',Exps(iExp),cA(icase),'_filter_4040DJFsingle.nc'));
        vv_40=squeeze(double(ncread(file,'v_v_3d')));
        
        nlon=length(lon);
        nlat=length(lat);
        nlev=length(lev);
        iy1=[2:nlat,1];
        iy0=[1:nlat];
        dlat=(lat(2)-lat(1))*d2r;
        
        TERM1=repmat(f0*sin(lat*d2r)',[nlon,1,nlev]).*U;
        
        Zy=(Z(:,iy1,:)-Z(:,iy0,:))./(a*dlat);
        Zy(:,1,:)=0;
        Zy(:,end,:)=0;
        TERM2=-g.*Zy;
        
        vvy=(vv(:,iy1,:)-vv(:,iy0,:))./(a*dlat);
        vvy(:,1,:)=0;
        vvy(:,end,:)=0;
        TERM3=-vvy;
        
        fldname=char(strcat(Exps(iExp),cA(icase)));
        TERM1s.(fldname)=TERM1;
        TERM1s.('long_name')='f*U';
        TERM2s.(fldname)=TERM2;
        TERM2s.('long_name')='-g*Z_y';
        TERM3s.(fldname)=TERM3;
        TERM3s.('long_name')='-\bar{vv}_y';
        Us.(fldname)=U;
        
    end
end
save('./V_momentum.mat','TERM1s','TERM2s','TERM3s','Us','lat','lon','lev');
else
    load('./V_momentum.mat')
end


%% How well is this equation balanced
for iExp=1:NExps
    for icase=1:N
        figure(70+iExp*10+icase)
        clf
        [~,ilev]=min(abs(lev-levshow));
        fldname=char(strcat(Exps(iExp),cA(icase)));
        fldname0=char(strcat(Exps(iExp),cA(1)));
        colormap(polarmap(21))
        contourf(lon,lat,squeeze(TERM1s.(fldname)(:,:,ilev))'-squeeze(TERM1s.(fldname0)(:,:,ilev))',[-0.001:0.0001:0.001],'LineColor','none');
        caxis([-0.001,0.001])
        
        hold on
        %colormap(jet(21))
        contour(lon,lat,(squeeze(TERM2s.(fldname)(:,:,ilev))'+squeeze(TERM3s.(fldname)(:,:,ilev))')-(squeeze(TERM2s.(fldname0)(:,:,ilev))'+squeeze(TERM3s.(fldname0)(:,:,ilev))'),[-0.001:0.0001:0.001],'LineWidth',2);
        caxis([-0.001,0.001])
        title(fldname,'Interpreter','none')
        xlabel('lon')
        ylabel('lat')
        set(gca,'FontSize',18)
        pause
    end
end

%% Is eddy drag dominant the U change?
for iExp=1:NExps
    for icase=1:N
        figure(170+iExp*10+icase)
        clf
        [~,ilev]=min(abs(lev-levshow));
        fldname=char(strcat(Exps(iExp),cA(icase)));
        fldname0=char(strcat(Exps(iExp),cA(1)));
        colormap(polarmap(21))
        contourf(lon,lat,squeeze(TERM1s.(fldname)(:,:,ilev))'-squeeze(TERM1s.(fldname0)(:,:,ilev))',[-0.001:0.0001:0.001],'LineColor','none');
        caxis([-0.001,0.001])
        
        hold on
        %colormap(jet(21))
        contour(lon,lat,(squeeze(squeeze(TERM3s.(fldname)(:,:,ilev))')-(squeeze(TERM3s.(fldname0)(:,:,ilev)))').*3,[-0.001:0.0001:0.001],'LineWidth',2);
        caxis([-0.001,0.001])
        title(fldname,'Interpreter','none')
        xlabel('lon')
        ylabel('lat')
        set(gca,'FontSize',18)
        pause
    end
end

%% plot line plot of TERM1 and TERM3
cs=jet(N);
for iExp=1:NExps
    figure(70+iExp)
    clf
    hold on
    for icase=1:N
        fldname=char(strcat(Exps(iExp),cA(icase)));
        TERM1_temp=TERM1s.(fldname);
        [~,ilat]=min(abs(lat-latshow));
        [~,ilev]=min(abs(lev-levshow));
        term1=squeeze(mean(mean(TERM1_temp(:,ilat-4:ilat+4,ilev-3:ilev+3),3),2));
        if icase==1
            term10=term1;
        end
        
        TERM3_temp=TERM3s.(fldname);
        [~,ilat]=min(abs(lat-latshow));
        [~,ilev]=min(abs(lev-levshow));
        term3=squeeze(mean(mean(TERM3_temp(:,ilat-4:ilat+4,ilev-3:ilev+3),3),2));
        if icase==1
            term30=term3;
        end
        
        yyaxis left
        plot(lon,term1-term10,'LineWidth',2,'Color',cs(icase,:),'LineStyle','-','Marker','none')
        ylim([-1e-3,1e-3])
        ylabel(strcat('U @',sprintf('%d',latshow),'N'))
        yyaxis right
        plot(lon,term3-term30,'LineWidth',2,'Color',cs(icase,:),'LineStyle','--','Marker','none')
        ylim([-1e-3,1e-3])
        ylabel(strcat('$vv_y$'),'Interpreter','latex')
    end
    xlabel('longitude (^oE)')
    
    
    title(strcat(char(Exps(iExp)),'Amps'),'Interpreter','none')
    set(gca,'FontSize',18)
   
end
        
        
        
        