%% Climatology change analysis: |V|, |\nabla Q|, cos(V,\nabla Q)
%% 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
%  The wave drag forcing div(dE)_y is PV source. can it explain the change
%  in PV advection?
%  Q = f + \nabla^2 \Psi + f^2 (1/S*\Psi_p)_p,  S = \bar(alpha)*(ln(\bar(\Theta)))_p
%    = f + curl(Vh) + f (-T'/\bar(T)/(ln\Theta)_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=false;

latshow=40.;
levshow=150.;
levshow2=600;

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]);
        sinlat_=repmat(sin(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')));
        Z=squeeze(double(ncread(file,'Z3')));    

        %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;
        
        Q = f0.*sinlat_ + (1/a).*seclat_.*((V(:,iy1,:).*coslat_(:,iy1,:)-V(:,iy0,:).*coslat_(:,iy0,:))./(dlat)...
                                            -(U(ix1,:,:)-U(ix0,:,:))./(dlon))...
            + f0.*sinlat_.* ((T(:,:,ip1)./TH_p_equivalent(:,:,ip1))-(T(:,:,ip0)./TH_p_equivalent(:,:,ip0)))./(pres(:,:,ip1)-pres(:,:,ip0));
        Qx = repmat(seclat,[length(lon),1,length(lev)]).*(Q(ix1,:,:)-Q(ix0,:,:))./(a*dlon);
        Qy = (Q(:,iy1,:)-Q(:,iy0,:))./(a*dlat);
        AdvQx = U.*Qx;
        AdvQy = V.*Qy;
        
        Ux = repmat(seclat,[length(lon),1,length(lev)]).*(U(ix1,:,:)-U(ix0,:,:))./(a*dlon);
        Uy = (U(:,iy1,:)-U(:,iy0,:))./(a*dlat);
        AdvUx = U.*Ux;
        AdvUy = V.*Uy;
        
        fldname=char(strcat(Exps(iExp),cA(icase)));
        Qs.(fldname)=Q;
        Qxs.(fldname)=Qx;
        Qys.(fldname)=Qy;
        Us.(fldname)=U;
        Vs.(fldname)=V;
        AdvQxs.(fldname)=AdvQx;
        AdvQys.(fldname)=AdvQy;
        AdvUxs.(fldname)=AdvUx;
        AdvUys.(fldname)=AdvUy;
        
    end
end
save('./PV.mat','Qs','Qxs','Qys','AdvQxs','AdvQys','AdvUxs','AdvUys','Us','Vs','lat','lon','lev');
else
    load('./PV.mat')
    load('./E_vector.mat')
    %load('./E_vector_noEPz40k2.mat')
end


%% How well is this equation balanced Delta(V\nabla Q)=-dy(div(E))
        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;
for iExp=1:NExps
    for icase=1:N
        figure(370+iExp*10+icase)
        clf
        [~,ilev]=min(abs(lev-levshow));
        [~,ilev2]=min(abs(lev-levshow2));
        fldname=char(strcat(Exps(iExp),cA(icase)));
        fldname0=char(strcat(Exps(iExp),cA(1)));
        colormap(polarmap(21))
        
        %contourf(lon,lat,(squeeze(AdvQxs.(fldname)(:,:,ilev))'+squeeze(AdvQys.(fldname)(:,:,ilev))')-(squeeze(AdvQxs.(fldname0)(:,:,ilev))'+squeeze(AdvQys.(fldname0)(:,:,ilev))'),[-1e-8:1e-9:1e-8],'LineColor','none');
        %caxis([-1e-8,1e-8]) 
        
        %contourf(lon,lat,(squeeze(Us.(fldname)(:,:,ilev))')-(squeeze(Us.(fldname0)(:,:,ilev))'),[-10:1:10],'LineColor','none');
        %caxis([-10,10]) 
        
        contourf(lon,lat,((squeeze(AdvUxs.(fldname)(:,:,ilev))'+squeeze(AdvUys.(fldname)(:,:,ilev))')-(squeeze(AdvUxs.(fldname0)(:,:,ilev))'+squeeze(AdvUys.(fldname0)(:,:,ilev))')).*1e5,[-10:1:10],'LineColor','none');
        caxis([-10,10])  
        
        hold on
        %colormap(jet(21))
        %contour(lon,lat,((squeeze(DIVEPs.(fldname)(:,iy1,ilev))'-squeeze(DIVEPs.(fldname0)(:,iy1,ilev))')-(squeeze(DIVEPs.(fldname)(:,iy0,ilev))'-squeeze(DIVEPs.(fldname0)(:,iy0,ilev))'))./(a*dlat).*30,[-1e-8:1e-9:1e-8],'LineWidth',1);
        %contour(lon,lat,((squeeze(DIVEP_40s.(fldname)(:,:,ilev))'-squeeze(DIVEP_40s.(fldname0)(:,:,ilev))'))./10000,[-1e-8:1e-9:1e-8],'LineWidth',1);
        %caxis([-1e-8,1e-8])
        contour(lon,lat,((squeeze(DIVEP_40_k2s.(fldname)(:,:,ilev))'-squeeze(DIVEP_40_k2s.(fldname0)(:,:,ilev))')).*1e5,[-10:1:10],'LineWidth',1);
        %contour(lon,lat,((squeeze(DIVEP_40_k2s.(fldname)(:,iy1,ilev))'-squeeze(DIVEP_40_k2s.(fldname0)(:,iy1,ilev))')-(squeeze(DIVEP_40_k2s.(fldname)(:,iy0,ilev))'-squeeze(DIVEP_40_k2s.(fldname0)(:,iy0,ilev))'))./(a*dlat).*1e11,[-10:1:10],'LineWidth',1);
        %contour(lon,lat,((squeeze(Us.(fldname)(:,:,ilev2))')-(squeeze(Us.(fldname0)(:,:,ilev2))')).*10,[-10:1:10]);
        caxis([-10,10]) 
        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)));
        [~,ilat]=min(abs(lat-latshow));
        [~,ilev]=min(abs(lev-levshow));
        
        DIVEP_temp=DIVEPs.(fldname);
        divep=100.*squeeze(mean(mean(DIVEP_temp(:,ilat-4:ilat+4,ilev-3:ilev+3),3),2));
        divepy=divep-100.*squeeze(mean(mean(DIVEP_temp(:,ilat-6:ilat+2,ilev-3:ilev+3),3),2));
        
        U_temp=Us.(fldname);
        u=squeeze(mean(mean(U_temp(:,ilat-4:ilat+4,ilev-3:ilev+3),3),2));
        
        AdvQ_temp=AdvQxs.(fldname)+AdvQys.(fldname);
        advq=squeeze(mean(mean(AdvQ_temp(:,ilat-4:ilat+4,ilev-3:ilev+3),3),2));
        
        DIVEP_40_temp=DIVEP_40_k2s.(fldname);
        divep_40=100.*squeeze(mean(mean(DIVEP_40_temp(:,ilat-4:ilat+4,ilev-3:ilev+3),3),2));
        divep_40y=divep_40-100.*squeeze(mean(mean(DIVEP_40_temp(:,ilat-6:ilat+2,ilev-3:ilev+3),3),2));
        
        if icase==1
            divep0=divep;
            divepy0=divepy;
            u0=u;
            advq0=advq;
            divep_400=divep_40;
            divep_40y0=divep_40y;
        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)'))
        
        plot(lon,u-u0,'LineWidth',2,'Color',cs(icase,:),'LineStyle','-','Marker','none')
        ylim([-10,10])
        ylabel(strcat('\Delta U'))
        
%         plot(lon,advq-advq0,'LineWidth',2,'Color',cs(icase,:),'LineStyle','-','Marker','none')
%         ylim([-1e-9,1e-9])
%         ylabel(strcat('\Delta Adv Q'))
        
        yyaxis right
        plot(lon,divep_40-divep_400,'LineWidth',2,'Color',cs(icase,:),'LineStyle','--','Marker','none')
        ylim([-2e-3,2e-3])
        ylabel(strcat('\Delta DIV(EP) MJO'))
        
%         plot(lon,divep_40y-divep_40y0,'LineWidth',2,'Color',cs(icase,:),'LineStyle','--','Marker','none')
%         ylim([-2e-3,2e-3])
%         ylabel(strcat('\Delta dY( DIV(EP) MJO )'))
    end
    xlabel('longitude (^oE)')
    
    
    title(strcat(char(Exps(iExp)),'Amps'),'Interpreter','none')
    set(gca,'FontSize',18)
   
end
        
        
        
        