%% Trying to explain the change in climatological U
%  ([<u(T)>]-[<u(0)>])/T = -2dx([<u(0)\bar{\tilde{u}}>]) - dx([<\bar{\tilde{u}^2}>]) ...
%            |                           |                            |       
%          TERM1                        TERM2                       TERM3  
%  - dy([<u(0)\bar{\tilde{v}}>]) - dy([<v(0)\bar{\tilde{u}}>]) - dy([<\bar{\tilde{u}*\tilde{v}}>])
%            |                                |                                   |       
%          TERM4                            TERM5                              TERM6
%  + f[<\bar{\tilde{v}}>] - gdx([<\bar{\tilde{\Phi}}>])
%            |                          |  
%           TERM7                     TERM8

%addpath('/glade/scratch/wykang/analysis/matlabcode/')
addpath('/Users/wanyingkang/Desktop/research/matlab/code/basic/')
MJONAME='MJOB';
yijing=true;
append='_lasttwomean';
append='_timeavg';

latshow=40.;
levshowtop=100.;
levshowbot=1000.;

T=20.;
T1day=86400.;
d2r=pi/180.;
a=6371000.;
f0=4*pi/86400;
g=9.8;
tau=40*86400;
kappa=2./7.;

%% calculate each term
%  capital: <> ensemble average; lower case: each ensemble 
%  b: bar;      t: tilde;     E: ensemble average
if ~yijing
    lat=squeeze(double(ncread('vinth_Last_CTRL_ENSavg.nc','lat')));
    lon=squeeze(double(ncread('vinth_Last_CTRL_ENSavg.nc','lon')));
    lev=squeeze(double(ncread('vinth_Last_CTRL_ENSavg.nc','lev')));
    ilev=squeeze(double(ncread('vinth_Last_CTRL_ENSavg.nc','ilev')));
    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]);
    dp=ilev(2:end)-ilev(1:end-1);
    
    
    % TERM1
    UT=squeeze(double(ncread(strcat('vinth_Last_',MJONAME,'_ENSavg.nc'),'U')));
    U0=squeeze(double(ncread('vinth_Last_CTRL_ENSavg.nc','U')));
    TERM1 = (UT-U0)./T;
    
    % TERM2
    u0=squeeze(double(ncread('vinth_U_Last_CTRL.nc','U')));
    bt_u=squeeze(double(ncread(strcat('vinth_U_timeavg_',MJONAME,'.nc'),'U'))) - u0; % all 20 day
    E_u0_bt_u=squeeze(mean(u0.*bt_u,4));
    TERM2 = - 2.* seclat_.*(E_u0_bt_u(ix1,:,:)-E_u0_bt_u(ix0,:,:))./(a*dlon).*T1day;
    
    % C1 = -dx(<u(0)*u(0)>)
    E_u0_u0 = squeeze(mean(u0.*u0,4));
    C1 = - seclat_.*(E_u0_u0(ix1,:,:)-E_u0_u0(ix0,:,:))./(a*dlon);
    
    % TERM3
    Eb_t_u_t_u=squeeze(double(ncread(strcat('vinth',append,'_',MJONAME,'_ENSavg.nc'),'UU'))) - 2.*E_u0_bt_u - E_u0_u0; 
    TERM3 = - seclat_.*(Eb_t_u_t_u(ix1,:,:)-Eb_t_u_t_u(ix0,:,:))./(a*dlon).*T1day;
    
    % TERM4
    v0=squeeze(double(ncread('vinth_V_Last_CTRL.nc','V')));
    bt_v=squeeze(double(ncread(strcat('vinth_V_timeavg_',MJONAME,'.nc'),'V'))) - v0;  % all 20 day
    E_u0_bt_v=squeeze(mean(u0.*bt_v,4));
    TERM4 = -(E_u0_bt_v(:,iy1,:)-E_u0_bt_v(:,iy0,:))./(a*dlat).*T1day;
    
    % TERM5
    E_v0_bt_u=squeeze(mean(v0.*bt_u,4));
    TERM5 = -(E_v0_bt_u(:,iy1,:)-E_v0_bt_u(:,iy0,:))./(a*dlat).*T1day;
    
    % C2 = -dy(<u(o)*v(0)>)
    E_u0_v0 = squeeze(mean(u0.*v0,4));
    C2 = -(E_u0_v0(:,iy1,:)-E_u0_v0(:,iy0,:))./(a*dlat).*T1day;
    
    % TERM6
    Eb_t_u_t_v=squeeze(double(ncread(strcat('vinth',append,'_',MJONAME,'_ENSavg.nc'),'VU'))) - E_u0_bt_v - E_v0_bt_u - E_u0_v0; 
    TERM6 = -(Eb_t_u_t_v(:,iy1,:)-Eb_t_u_t_v(:,iy0,:))./(a*dlat).*T1day;
    
    % TERM7
	E_bt_v=squeeze(double(ncread(strcat('vinth',append,'_',MJONAME,'_ENSavg.nc'),'V')))-squeeze(double(ncread('vinth_Last_CTRL_ENSavg.nc','V'))); 
    TERM7 = f0.*sinlat_.*E_bt_v.*T1day;
    
    % TERM8
    E_bt_phi=squeeze(double(ncread(strcat('vinth',append,'_',MJONAME,'_ENSavg.nc'),'Z3')))-squeeze(double(ncread('vinth_Last_CTRL_ENSavg.nc','Z3'))); 
    TERM8 = -g.* seclat_.*(E_bt_phi(ix1,:,:)-E_bt_phi(ix0,:,:))./(a*dlon).*T1day;
        

save(strcat('./',MJONAME,'_ENS_Ubudget',append,'.mat'),'TERM1','TERM2','TERM3','TERM4','TERM5','TERM6','TERM7','TERM8','lat','lon','lev','ilev','dp');
else
    load(strcat('./',MJONAME,'_ENS_Ubudget',append,'.mat'))
end

%% vertical average
dp_=permute(repmat(dp,[1,length(lon),length(lat)]),[2,3,1]);
indlev=find(lev<levshowbot&lev>levshowtop);
term1=squeeze(nansum(TERM1(:,:,indlev).*dp_(:,:,indlev),3))'./sum(dp(indlev)).*10;
term2=squeeze(nansum(TERM2(:,:,indlev).*dp_(:,:,indlev),3))'./sum(dp(indlev));
term3=squeeze(nansum(TERM3(:,:,indlev).*dp_(:,:,indlev),3))'./sum(dp(indlev));
term4=squeeze(nansum(TERM4(:,:,indlev).*dp_(:,:,indlev),3))'./sum(dp(indlev));
term5=squeeze(nansum(TERM5(:,:,indlev).*dp_(:,:,indlev),3))'./sum(dp(indlev));
term6=squeeze(nansum(TERM6(:,:,indlev).*dp_(:,:,indlev),3))'./sum(dp(indlev));
term7=squeeze(nansum(TERM7(:,:,indlev).*dp_(:,:,indlev),3))'./sum(dp(indlev));
term8=squeeze(nansum(TERM8(:,:,indlev).*dp_(:,:,indlev),3))'./sum(dp(indlev));

%% TERM2
figure(2)
clf
colormap(polarmap(21))
contourf(lon,lat,term1,[-5:0.5:5],'LineColor','none');
caxis([-5,5])

hold on
%colormap(jet(21))
contour(lon,lat,term2,[-5:0.5:5],'LineWidth',2);
caxis([-5,5])
title('TERM2','Interpreter','none')
xlabel('lon')
ylabel('lat')
set(gca,'FontSize',18)
pause

%% TERM3

figure(3)
clf
colormap(polarmap(21))
contourf(lon,lat,term1,[-5:0.5:5],'LineColor','none');
caxis([-5,5])

hold on
%colormap(jet(21))
contour(lon,lat,term3,[-5:0.5:5],'LineWidth',2);
caxis([-5,5])
title('TERM3','Interpreter','none')
xlabel('lon')
ylabel('lat')
set(gca,'FontSize',18)
pause

%% TERM4
figure(4)
clf
colormap(polarmap(21))
contourf(lon,lat,term1,[-5:0.5:5],'LineColor','none');
caxis([-5,5])

hold on
%colormap(jet(21))
contour(lon,lat,term4,[-5:0.5:5],'LineWidth',2);
caxis([-5,5])
title('TERM4','Interpreter','none')
xlabel('lon')
ylabel('lat')
set(gca,'FontSize',18)
pause

%% TERM5
figure(5)
clf
colormap(polarmap(21))
contourf(lon,lat,term1,[-5:0.5:5],'LineColor','none');
caxis([-5,5])

hold on
%colormap(jet(21))
contour(lon,lat,term5,[-5:0.5:5],'LineWidth',2);
caxis([-5,5])
title('TERM5','Interpreter','none')
xlabel('lon')
ylabel('lat')
set(gca,'FontSize',18)
pause

%% TERM6
figure(6)
clf
colormap(polarmap(21))
contourf(lon,lat,term1,[-5:0.5:5],'LineColor','none');
caxis([-5,5])

hold on
%colormap(jet(21))
contour(lon,lat,term6,[-5:0.5:5],'LineWidth',2);
caxis([-5,5])
title('TERM6','Interpreter','none')
xlabel('lon')
ylabel('lat')
set(gca,'FontSize',18)
pause

%% TERM7
figure(7)
clf
colormap(polarmap(21))
contourf(lon,lat,term1,[-5:0.5:5],'LineColor','none');
caxis([-5,5])

hold on
%colormap(jet(21))
contour(lon,lat,term7,[-5:0.5:5],'LineWidth',2);
caxis([-5,5])
title('TERM7','Interpreter','none')
xlabel('lon')
ylabel('lat')
set(gca,'FontSize',18)
pause

%% TERM8
figure(8)
clf
colormap(polarmap(21))
contourf(lon,lat,term1,[-5:0.5:5],'LineColor','none');
caxis([-5,5])

hold on
%colormap(jet(21))
contour(lon,lat,term8,[-5:0.5:5],'LineWidth',2);
caxis([-5,5])
title('TERM8','Interpreter','none')
xlabel('lon')
ylabel('lat')
set(gca,'FontSize',18)
pause

%% TERM2+TERM3
figure(23)
clf
colormap(polarmap(21))
contourf(lon,lat,term1,[-5:0.5:5],'LineColor','none');
caxis([-5,5])

hold on
%colormap(jet(21))
contour(lon,lat,term2+term3,[-5:0.5:5],'LineWidth',2);
caxis([-5,5])
title('TERM2 + TERM3','Interpreter','none')
xlabel('lon')
ylabel('lat')
set(gca,'FontSize',18)
pause

%% TERM4+TERM5+TERM6
figure(1456)
clf
colormap(polarmap(21))
contourf(lon,lat,term1,[-5:0.5:5],'LineColor','none');
caxis([-5,5])

hold on
%colormap(jet(21))
contour(lon,lat,term4+term5+term6,[-5:0.5:5],'LineWidth',2);
caxis([-5,5])
title('TERM4 + TERM5 + TERM6','Interpreter','none')
xlabel('lon')
ylabel('lat')
set(gca,'FontSize',18)
pause

%% TERM7+TERM8
figure(78)
clf
colormap(polarmap(21))
contourf(lon,lat,term1,[-5:0.5:5],'LineColor','none');
caxis([-5,5])

hold on
%colormap(jet(21))
contour(lon,lat,term7+term8,[-5:0.5:5],'LineWidth',2);
caxis([-5,5])
title('TERM7 + TERM8','Interpreter','none')
xlabel('lon')
ylabel('lat')
set(gca,'FontSize',18)
pause

%% TERM2+TERM4+TERM5
figure(245)
clf
colormap(polarmap(21))
contourf(lon,lat,term1,[-5:0.5:5],'LineColor','none');
caxis([-5,5])

hold on
%colormap(jet(21))
contour(lon,lat,term4+term5+term2,[-5:0.5:5],'LineWidth',2);
caxis([-5,5])
title('TERM2 + TERM4 + TERM5','Interpreter','none')
xlabel('lon')
ylabel('lat')
set(gca,'FontSize',18)
pause

%% TERM3+TERM6
figure(36)
clf
colormap(polarmap(21))
contourf(lon,lat,term1,[-5:0.5:5],'LineColor','none');
caxis([-5,5])

hold on
%colormap(jet(21))
contour(lon,lat,term3+term6,[-5:0.5:5],'LineWidth',2);
caxis([-5,5])
title('TERM3 + TERM6','Interpreter','none')
xlabel('lon')
ylabel('lat')
set(gca,'FontSize',18)
pause
        
%% TERM2+TERM3+TERM4+TERM5+TERM6
figure(23456)
clf
colormap(polarmap(21))
contourf(lon,lat,term1,[-5:0.5:5],'LineColor','none');
caxis([-5,5])

hold on
%colormap(jet(21))
contour(lon,lat,term2+term3+term4+term5+term6,[-5:0.5:5],'LineWidth',2);
caxis([-5,5])
title('TERM2 + TERM3 + TERM4 + TERM5 + TERM6','Interpreter','none')
xlabel('lon')
ylabel('lat')
set(gca,'FontSize',18)
pause

%% TERM2+TERM3+TERM4+TERM5+TERM6++TERM7+TERM8
figure(2345678)
clf
colormap(polarmap(21))
contourf(lon,lat,term1,[-5:0.5:5],'LineColor','none');
caxis([-5,5])

hold on
%colormap(jet(21))
contour(lon,lat,term2+term3+term4+term5+term6+term7+term8,[-5:0.5:5],'LineWidth',2);
caxis([-5,5])
title('TERM2 + TERM3 + TERM4 + TERM5 + TERM6 + TERM7 + TERM8','Interpreter','none')
xlabel('lon')
ylabel('lat')
set(gca,'FontSize',18)
pause



       
        
