%% Trying to explain the change in climatological U
%  PV equation:
%  U*Q_x + V*Q_y = -div(E)_y;
%  dU*Q_x + U*dQ_x + dV*Q_y + V*dQ_y = -div(dE)_y
%  Whether the wave drag forcing div(dE)_y is mainly due to MJO waves?
%  E = (vv-uu, -uv, f*vth/TH_p)


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;
tau=40*86400;
kappa=2./7.;

%% 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'));
         
        lat=squeeze(double(ncread(file,'lat')));
        lon=squeeze(double(ncread(file,'lon')));
        lev=squeeze(double(ncread(file,'lev'))); 
        nlon=length(lon);
        nlat=length(lat);
        nlev=length(lev);
        iy1=[2:nlat,nlat];
        iy0=[1,1:nlat-1];
        dlat=(lat(2)-lat(1))*d2r*2;
        ix1=[2:nlon,1];
        ix0=[nlon,1:nlon-1]; 
        dlon=(lon(2)-lon(1))*d2r*2;
        coslat_=repmat(cos(lat.*d2r)',[nlon,1,nlev]);
        seclat=1./cos(lat.*d2r)';
        seclat(1)=0;
        seclat(end)=0;
        seclat_=repmat(seclat,[nlon,1,nlev]);
        ip1=[2:nlev,nlev];
        ip0=[1,1:nlev-1];
        pres=permute(repmat(lev.*100,[1,nlon,nlat]),[2,3,1]);
        
        U=squeeze(double(ncread(file,'U')));
        V=squeeze(double(ncread(file,'V')));
        T=squeeze(double(ncread(file,'T')));
        VU=squeeze(double(ncread(file,'VU')));
        VV=squeeze(double(ncread(file,'VV')));
        VT=squeeze(double(ncread(file,'VT')));
        UU=squeeze(double(ncread(file,'UU')));
        Z=squeeze(double(ncread(file,'Z3')));    
        vu=VU-V.*U;
        uu=UU-U.*U;
        vv=VV-V.*V;
        vt=VT-V.*T;
        %TH_p_equivalent=(T(:,:,ip1)-T(:,:,ip0))./(pres(:,:,ip1)-pres(:,:,ip0))+T*(-kappa)./pres;
        Tbar=repmat(mean(T),[nlon,1,1]);
        TH_p_equivalent=(Tbar(:,:,ip1)-Tbar(:,:,ip0))./(pres(:,:,ip1)-pres(:,:,ip0))+Tbar*(-kappa)./pres;
        
        
        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')));
        vt_40=squeeze(double(ncread(file,'v_t_3d')));
        
        vu_40_k2=squeeze(double(ncread(file,'u_v_3d_k2')));
        uu_40_k2=squeeze(double(ncread(file,'u_u_3d_k2')));
        vv_40_k2=squeeze(double(ncread(file,'v_v_3d_k2')));
        vt_40_k2=squeeze(double(ncread(file,'v_t_3d_k2')));
        
        %EPx=vv-uu;
        %EPx_40=vv_40-uu_40;
        EPx=(vv-uu);
        EPx_40=(vv_40-uu_40);
        EPx_40_k2=(vv_40_k2-uu_40_k2);
        EPx=(-uu);
        EPx_40=(-uu_40);
        EPx_40_k2=(-uu_40_k2);
        EPy=-vu;
        EPy_40=-vu_40;
        EPy_40_k2=-vu_40_k2;
        EPz=f0.*coslat_.*vt./TH_p_equivalent;
        EPz_40=f0.*coslat_.*vt_40./TH_p_equivalent;
        EPz_40_k2=f0.*coslat_.*vt_40_k2./TH_p_equivalent;
        DIVEP=1/a.*seclat_.*(EPx(ix1,:,:)-EPx(ix0,:,:))./(dlon)...
            + 1/a.*seclat_.*(EPy(:,iy1,:).*coslat_(:,iy1,:)-EPy(:,iy0,:).*coslat_(:,iy0,:))./(dlat);...
            - (EPz(:,:,ip1)-EPz(:,:,ip0))./(pres(:,:,ip1)-pres(:,:,ip0));
        DIVEP_40=1/a.*seclat_.*(EPx_40(ix1,:,:)-EPx_40(ix0,:,:))./(dlon)...
            + 1/a.*seclat_.*(EPy_40(:,iy1,:).*coslat_(:,iy1,:)-EPy_40(:,iy0,:).*coslat_(:,iy0,:))./(dlat);...
            - (EPz_40(:,:,ip1)-EPz_40(:,:,ip0))./(pres(:,:,ip1)-pres(:,:,ip0));
        DIVEP_40_k2=1/a.*seclat_.*(EPx_40_k2(ix1,:,:)-EPx_40_k2(ix0,:,:))./(dlon)...
            + 1/a.*seclat_.*(EPy_40_k2(:,iy1,:).*coslat_(:,iy1,:)-EPy_40_k2(:,iy0,:).*coslat_(:,iy0,:))./(dlat);...
            - (EPz_40_k2(:,:,ip1)-EPz_40_k2(:,:,ip0))./(pres(:,:,ip1)-pres(:,:,ip0));
         
        fldname=char(strcat(Exps(iExp),cA(icase)));
        EPxs.(fldname)=EPx;
        EPys.(fldname)=EPy;
        EPzs.(fldname)=EPz;
        DIVEPs.(fldname)=DIVEP;
        EPx_40s.(fldname)=EPx_40;
        EPy_40s.(fldname)=EPy_40;
        EPz_40s.(fldname)=EPz_40;
        DIVEP_40s.(fldname)=DIVEP_40;
        EPx_40_k2s.(fldname)=EPx_40_k2;
        EPy_40_k2s.(fldname)=EPy_40_k2;
        EPz_40_k2s.(fldname)=EPz_40_k2;
        DIVEP_40_k2s.(fldname)=DIVEP_40_k2;
        Us.(fldname)=U;
        
    end
end
save('./E_vector.mat','EPxs','EPys','EPzs','DIVEPs','EPx_40s','EPy_40s','EPz_40s','DIVEP_40s','EPx_40_k2s','EPy_40_k2s','EPz_40_k2s','DIVEP_40_k2s','Us','lat','lon','lev');
else
    load('./E_vector.mat')
end


%% How well is this equation balanced
for iExp=1:NExps
    for icase=1:N
        figure(370+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(DIVEPs.(fldname)(:,:,ilev))'-squeeze(DIVEPs.(fldname0)(:,:,ilev))',[-0.0002:0.00002:0.0002],'LineColor','none');
        caxis([-0.0002,0.0002])
        
        hold on
        %colormap(jet(21))
        contour(lon,lat,-(squeeze(DIVEP_40s.(fldname)(:,:,ilev))'-squeeze(DIVEP_40s.(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)));
        DIVEP_temp=DIVEPs.(fldname);
        [~,ilat]=min(abs(lat-latshow));
        [~,ilev]=min(abs(lev-levshow));
        divep=100.*squeeze(mean(mean(DIVEP_temp(:,ilat-4:ilat+4,ilev-3:ilev+3),3),2));
        if icase==1
            divep0=divep;
        end
        
        DIVEP_40_temp=DIVEP_40s.(fldname);
        [~,ilat]=min(abs(lat-latshow));
        [~,ilev]=min(abs(lev-levshow));
        divep_40=100.*squeeze(mean(mean(DIVEP_40_temp(:,ilat-4:ilat+4,ilev-3:ilev+3),3),2));
        if icase==1
            divep_400=divep_40;
        end
     
        
        yyaxis left
        plot(lon,divep-divep0,'LineWidth',2,'Color',cs(icase,:),'LineStyle','-','Marker','none')
        ylim([-5e-3,5e-3])
        ylabel(strcat('\Delta DIV(EP)'))
        yyaxis right
        plot(lon,divep_40-divep_400,'LineWidth',2,'Color',cs(icase,:),'LineStyle','--','Marker','none')
        ylim([-5e-3,5e-3])
        ylabel(strcat('\Delta DIV(EP) MJO'))
    end
    xlabel('longitude (^oE)')
    
    
    title(strcat(char(Exps(iExp)),'Amps'),'Interpreter','none')
    set(gca,'FontSize',18)
   
end
        
        
        
        