%  Trying to explain the change in climatological U
%% dEKE/dt = (EPx,EPy)*\nabla_H(U)
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.;

%% 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'));
        UU=squeeze(double(ncread(file,'UU')));
        U=squeeze(double(ncread(file,'U')));
        VV=squeeze(double(ncread(file,'VV')));
        V=squeeze(double(ncread(file,'V')));
        UV=squeeze(double(ncread(file,'VU')));
        lat=squeeze(double(ncread(file,'lat')));
        lon=squeeze(double(ncread(file,'lon')));
        lev=squeeze(double(ncread(file,'lev')));
        
        uu=UU-U.^2;
        vv=VV-V.^2;
        uv=UV-U.*V;
        EPx=vv-uu;
        EPy=-uv;
        
        ix1=[2:length(lon),1];
        ix0=[1:length(lon)]; 
        dlon=(lon(2)-lon(1))*d2r;
        seclat=1./cos(lat.*d2r)';
        seclat(1)=0;
        seclat(end)=0;
        Ux=repmat(seclat,[length(lon),1,length(lev)]).*(U(ix1,:,:)-U(ix0,:,:))./(a*dlon);
        iy1=[2:length(lat),1];
        iy0=[1:length(lat)];
        dlat=(lat(2)-lat(1))*d2r;
        Uy=(U(:,iy1,:)-U(:,iy0,:))./(a*dlat);
        Uy(:,1,:)=0;
        Uy(:,end,:)=0;
        fldname=char(strcat(Exps(iExp),cA(icase)));
        eddygainxs.(fldname)=EPx.*(Ux);
        eddygainys.(fldname)=EPy.*(Uy);
        EPxs.(fldname)=EPx;
        EPys.(fldname)=EPy;
        Us.(fldname)=U;
    end
end
save('./Eddygain.mat','eddygainxs','eddygainys','EPxs','EPys','Us','lat','lon','lev');
else
    load('./Eddygain.mat')
end

%% plot quiver plot of (EPx,EPy) and overlap contour plot of U
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)));
        %quiver(lon(1:3:end),lat(1:3:end),squeeze(EPxs.(fldname)(1:3:end,1:3:end,ilev))',squeeze(EPys.(fldname)(1:3:end,1:3:end,ilev))');
        quiver(lon(1:3:end),lat(1:3:end),squeeze(EPxs.(fldname)(1:3:end,1:3:end,ilev))'-squeeze(EPxs.(fldname0)(1:3:end,1:3:end,ilev))',squeeze(EPys.(fldname)(1:3:end,1:3:end,ilev))'-squeeze(EPys.(fldname0)(1:3:end,1:3:end,ilev))');
        hold on
        colormap(jet(21))
        contour(lon,lat,squeeze(Us.(fldname)(:,:,ilev))',[-10:3:50]);
        caxis([-10,50])
        title(fldname,'Interpreter','none')
        xlabel('lon')
        ylabel('lat')
        set(gca,'FontSize',18)
        pause
    end
end

%% plot contour plot of eddy energy gain and overlap contour plot of dU
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(Us.(fldname)(:,:,ilev))'-squeeze(Us.(fldname0)(:,:,ilev))',[-20:2:20],'LineColor','none');
        caxis([-20,20])
        
        hold on
        %colormap(jet(21))
        contour(lon,lat,(squeeze(eddygainxs.(fldname)(:,:,ilev))'+squeeze(eddygainys.(fldname)(:,:,ilev))').*2e4-(squeeze(eddygainxs.(fldname0)(:,:,ilev))'+squeeze(eddygainys.(fldname0)(:,:,ilev))').*2e4,[-20:4:20],'LineWidth',2);
        caxis([-20,20])
        title(fldname,'Interpreter','none')
        xlabel('lon')
        ylabel('lat')
        set(gca,'FontSize',18)
        pause
    end
end

%% plot line plot of dU and Eddy Energy Gain
cs=jet(N);
for iExp=1:NExps
    figure(70+iExp)
    clf
    hold on
    for icase=1:N
        fldname=char(strcat(Exps(iExp),cA(icase)));
        eddygain_temp=eddygainxs.(fldname);
        %eddygain_temp=eddygainxs.(fldname)+eddygainys.(fldname);
        [~,ilat]=min(abs(lat-latshow));
        [~,ilev]=min(abs(lev-levshow));
        eddygain1=squeeze(mean(mean(eddygain_temp(:,ilat-4:ilat+4,ilev-3:ilev+3),3),2));
        if icase==1
            eddygain0=eddygain1;
        end
        
        U_temp=Us.(fldname);
        U1=squeeze(mean(mean(U_temp(:,ilat-4:ilat+4,ilev-3:ilev+3),3),2));
        if icase==1
            U0=U1;
        end
        
        yyaxis left
        %plot(lon,eddygain1,'LineWidth',2,'Color',cs(icase,:),'LineStyle','-','Marker','none')
        %ylim([-1e-3,1e-3])
        %ylabel(strcat('Eddy Energy Gain @',sprintf('%d',latshow),'N, ',sprintf('%d',levshow),'mb'))
        plot(lon,U1-U0,'LineWidth',2,'Color',cs(icase,:),'LineStyle','-','Marker','none')
        ylim([-10,10])
        ylabel(strcat('U @',sprintf('%d',latshow),'N'))
        yyaxis right
        plot(lon,eddygain1-eddygain0,'LineWidth',2,'Color',cs(icase,:),'LineStyle','--','Marker','none')
        ylim([-1e-3,1e-3])
        ylabel(strcat('$\Delta\ Eddy\ Energy\ Gain\ @',sprintf('%d',latshow),'N, ',sprintf('%d',levshow),'mb$'),'Interpreter','latex')
    end
    xlabel('longitude (^oE)')
    
    
    title(strcat(char(Exps(iExp)),'Amps'),'Interpreter','none')
    set(gca,'FontSize',18)
   
end
        
        
        
        