%% Trying to explain the change in climatological U
%  -f*dV = -g*Z_x - (du'u')_x - (du'v')_y - 1/tau*dU
%    |        |        |           |           |
%  dTERM1   dTERM2   dTERM3      dTERM4      dTERM5
%  TERM6 = (vv)_x [~ f*u_a (ageostrophic)]
%  TERM7 = U*U_x 
%  TERM6 ~ TERM7 (dominant balance shown by Hoskins 1983)

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=false;

latshow=40.;
levshow=250.;

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

%% 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')));
        VU=squeeze(double(ncread(file,'VU')));
        UU=squeeze(double(ncread(file,'UU')));
        VV=squeeze(double(ncread(file,'VV')));
        Z=squeeze(double(ncread(file,'Z3')));    
        vu=VU-V.*U;
        uu=UU-U.*U;
        vv=VV-V.*V;
        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'));
        vu_40=squeeze(double(ncread(file,'u_v_3d')));
        uu_40=squeeze(double(ncread(file,'u_u_3d')));
        vv_40=squeeze(double(ncread(file,'v_v_3d')));
        
        
%         vu=vu_40;
%         uu=uu_40;
%         vv=vv_40;
        
        nlon=length(lon);
        nlat=length(lat);
        nlev=length(lev);
        iy1=[2:nlat,1];
        iy0=[1:nlat];
        dlat=(lat(2)-lat(1))*d2r;
        ix1=[2:nlon,1];
        ix0=[1:nlon]; 
        dlon=(lon(2)-lon(1))*d2r;
        seclat=1./cos(lat.*d2r)';
        seclat(1)=0;
        seclat(end)=0;
        
        TERM1=-repmat(f0*sin(lat*d2r)',[nlon,1,nlev]).*V;
        
        Zx=repmat(seclat,[length(lon),1,length(lev)]).*(Z(ix1,:,:)-Z(ix0,:,:))./(a*dlon);
        TERM2=-g.*Zx;
        
        uux=repmat(seclat,[length(lon),1,length(lev)]).*(uu(ix1,:,:)-uu(ix0,:,:))./(a*dlon);
        TERM3=-uux;
        
        vuy=(vu(:,iy1,:)-vu(:,iy0,:))./(a*dlat);
        vuy(:,1,:)=0;
        vuy(:,end,:)=0;
        TERM4=-vuy;
        
        TERM5=-(1/tau).*U;
        
        vvx=repmat(seclat,[length(lon),1,length(lev)]).*(vv(ix1,:,:)-vv(ix0,:,:))./(a*dlon);
        TERM6=vvx;
        
        UUx=0.5.*repmat(seclat,[length(lon),1,length(lev)]).*(U(ix1,:,:).^2-U(ix0,:,:).^2)./(a*dlon);
        TERM7=UUx;
               
        fldname=char(strcat(Exps(iExp),cA(icase)));
        TERM1s.(fldname)=TERM1;
        TERM1s.('long_name')='-f*V';
        TERM2s.(fldname)=TERM2;
        TERM2s.('long_name')='-g*Z_x';
        TERM3s.(fldname)=TERM3;
        TERM3s.('long_name')='-\bar{uu}_x';
        TERM4s.(fldname)=TERM4;
        TERM4s.('long_name')='-\bar{vu}_y';
        TERM5s.(fldname)=TERM5;
        TERM5s.('long_name')='-U/tau';
        TERM6s.(fldname)=TERM6;
        TERM6s.('long_name')='\bar{vv}_x';
        TERM7s.(fldname)=TERM7;
        TERM7s.('long_name')='\bar{UU}_x';
        Us.(fldname)=U;
        
    end
end
save('./U_momentum.mat','TERM1s','TERM2s','TERM3s','TERM4s','TERM5s','TERM6s','TERM7s','Us','lat','lon','lev');
else
    load('./U_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(TERM4s.(fldname)(:,:,ilev))'+squeeze(TERM5s.(fldname)(:,:,ilev))')-(squeeze(TERM2s.(fldname0)(:,:,ilev))'+squeeze(TERM3s.(fldname0)(:,:,ilev))'+squeeze(TERM4s.(fldname0)(:,:,ilev))'+squeeze(TERM5s.(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(TERM5s.(fldname)(:,:,ilev))'-squeeze(TERM5s.(fldname0)(:,:,ilev))').*100,[-0.001:0.0001:0.001],'LineColor','none');
        caxis([-0.001,0.001])
        
        hold on
        %colormap(jet(21))
        contour(lon,lat,((squeeze(TERM3s.(fldname)(:,:,ilev))'+squeeze(TERM4s.(fldname)(:,:,ilev))')-(squeeze(TERM3s.(fldname0)(:,:,ilev))'+squeeze(TERM4s.(fldname0)(:,:,ilev))')).*10,[-0.001:0.0001:0.001],'LineWidth',1);
        caxis([-0.001,0.001])
        title(fldname,'Interpreter','none')
        xlabel('lon')
        ylabel('lat')
        set(gca,'FontSize',18)
        pause
    end
end

%% U*U_x = f*va = (vv)x
for iExp=1:NExps
    for icase=1:N
        figure(270+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(TERM7s.(fldname)(:,:,ilev))'-squeeze(TERM7s.(fldname0)(:,:,ilev))'),[-0.0002:0.00002:0.0002],'LineColor','none');
        caxis([-0.0002,0.0002])
        
        hold on
        %colormap(jet(21))
        contour(lon,lat,-((squeeze(TERM6s.(fldname)(:,:,ilev))')-(squeeze(TERM6s.(fldname0)(:,:,ilev))')),[-0.0002:0.00002:0.0002],'LineWidth',1);
        caxis([-0.0002,0.0002])
        title(fldname,'Interpreter','none')
        xlabel('lon')
        ylabel('lat')
        set(gca,'FontSize',18)
        pause
    end
end

%% plot line plot of TERM5 and TERM3+TERM4
cs=jet(N);
for iExp=1:NExps
    figure(70+iExp)
    clf
    hold on
    for icase=1:N
        fldname=char(strcat(Exps(iExp),cA(icase)));
        TERM5_temp=TERM5s.(fldname);
        [~,ilat]=min(abs(lat-latshow));
        [~,ilev]=min(abs(lev-levshow));
        term5=100.*squeeze(mean(mean(TERM5_temp(:,ilat-4:ilat+4,ilev-3:ilev+3),3),2));
        if icase==1
            term50=term5;
        end
        
        TERM34_temp=TERM1s.(fldname)+TERM4s.(fldname);
        [~,ilat]=min(abs(lat-latshow));
        [~,ilev]=min(abs(lev-levshow));
        term34=squeeze(mean(mean(TERM34_temp(:,ilat-4:ilat+4,ilev-3:ilev+3),3),2));
        if icase==1
            term340=term34;
        end
        
        yyaxis left
        plot(lon,term5-term50,'LineWidth',2,'Color',cs(icase,:),'LineStyle','-','Marker','none')
        ylim([-2e-4,2e-4])
        ylabel(strcat('U @',sprintf('%d',latshow),'N'))
        yyaxis right
        plot(lon,term34-term340,'LineWidth',2,'Color',cs(icase,:),'LineStyle','--','Marker','none')
        ylim([-2e-4,2e-4])
        ylabel(strcat('$uu_x+uv_y$'),'Interpreter','latex')
    end
    xlabel('longitude (^oE)')
    
    
    title(strcat(char(Exps(iExp)),'Amps'),'Interpreter','none')
    set(gca,'FontSize',18)
   
end
        
        
        
        