%% E vector DIV ~ dU
% need the output from drift_Ubudget.m, 'ENS_Ubudget.mat', which is 
% calculated on cluster and downloaded. 
levshowtop=100.;
levshowbot=800.;
MJONAME='MJOC';
append='_lasttwomean';
append='_timeavg';

T=20.;
T1day=86400.;
d2r=pi/180.;
a=6371000.;
f0=4*pi/86400;
g=9.8;
tau=40*86400;
kappa=2./7.;
load(strcat('./',MJONAME,'_ENS_Ubudget',append,'.mat'));

% prepare coordinate
file='vinth_timeavg_CTRL_reynold_term_ENSavg.nc';
nlon=length(lon);
nlat=length(lat);
nlev=length(lev);
ip1=[2:nlev,nlev];
ip0=[1,1:nlev-1];
pres=permute(repmat(lev.*100,[1,length(lon),length(lat)]),[2,3,1]);
dp=ilev(2:end)-ilev(1:end-1);
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)',[length(lon),1,length(lev)]);
seclat=1./cos(lat.*d2r)';
seclat(1)=0;
seclat(end)=0;
seclat_=repmat(seclat,[nlon,1,nlev]);

%read data
uu_CTRL=squeeze(double(ncread(file,'uu')));
vu_CTRL=squeeze(double(ncread(file,'vu')));
vv_CTRL=squeeze(double(ncread(file,'vv')));
vt_CTRL=squeeze(double(ncread(file,'vt')));
T_CTRL=squeeze(double(ncread(file,'T')));
Tbar_CTRL=repmat(mean(T_CTRL),[nlon,1,1]);
TH_p_CTRL=(Tbar_CTRL(:,:,ip1)-Tbar_CTRL(:,:,ip0))./(pres(:,:,ip1)-pres(:,:,ip0))+Tbar_CTRL*(-kappa)./pres;

file=strcat('vinth',append,'_',MJONAME,'_reynold_term_ENSavg.nc');

uu_MJO=squeeze(double(ncread(file,'uu')));
vu_MJO=squeeze(double(ncread(file,'vu')));
vv_MJO=squeeze(double(ncread(file,'vv')));
vt_MJO=squeeze(double(ncread(file,'vt')));
T_MJO=squeeze(double(ncread(file,'T')));
Tbar_MJO=repmat(mean(T_MJO),[nlon,1,1]);
TH_p_MJO=(Tbar_MJO(:,:,ip1)-Tbar_MJO(:,:,ip0))./(pres(:,:,ip1)-pres(:,:,ip0))+Tbar_MJO*(-kappa)./pres;

%EP flux and divergence
EPx_CTRL=(vv_CTRL-uu_CTRL);
EPx_MJO=(vv_MJO-uu_MJO);
EPy_CTRL=-vu_CTRL;
EPy_MJO=-vu_MJO;
EPz_CTRL=f0.*coslat_.*vt_CTRL./TH_p_CTRL;
EPz_MJO=f0.*coslat_.*vt_MJO./TH_p_MJO;

DIVEP_CTRL=1/a.*seclat_.*(EPx_CTRL(ix1,:,:)-EPx_CTRL(ix0,:,:))./(dlon)...
      + 1/a.*seclat_.*(EPy_CTRL(:,iy1,:).*coslat_(:,iy1,:)-EPy_CTRL(:,iy0,:).*coslat_(:,iy0,:))./(dlat)...
      - (EPz_CTRL(:,:,ip1)-EPz_CTRL(:,:,ip0))./(pres(:,:,ip1)-pres(:,:,ip0));
  
DIVEP_MJO=1/a.*seclat_.*(EPx_MJO(ix1,:,:)-EPx_MJO(ix0,:,:))./(dlon)...
      + 1/a.*seclat_.*(EPy_MJO(:,iy1,:).*coslat_(:,iy1,:)-EPy_MJO(:,iy0,:).*coslat_(:,iy0,:))./(dlat)...
      - (EPz_MJO(:,:,ip1)-EPz_MJO(:,:,ip0))./(pres(:,:,ip1)-pres(:,:,ip0));
  %DIVEP_CTRL=DIVEP_CTRL.*0;
  
%vertical integrate
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;
divep_ctrl=squeeze(nansum(DIVEP_CTRL(:,:,indlev).*dp_(:,:,indlev),3))'./sum(dp(indlev)).*T1day.*3;
divep_mjo=squeeze(nansum(DIVEP_MJO(:,:,indlev).*dp_(:,:,indlev),3))'./sum(dp(indlev)).*T1day.*3;
ex_ctrl=squeeze(nansum(EPx_CTRL(:,:,indlev).*dp_(:,:,indlev),3))'./sum(dp(indlev)).*T1day./1e7;
ey_ctrl=squeeze(nansum(EPy_CTRL(:,:,indlev).*dp_(:,:,indlev),3))'./sum(dp(indlev)).*T1day./1e6;
ex_mjo=squeeze(nansum(EPx_MJO(:,:,indlev).*dp_(:,:,indlev),3))'./sum(dp(indlev)).*T1day./1e6;
ey_mjo=squeeze(nansum(EPy_MJO(:,:,indlev).*dp_(:,:,indlev),3))'./sum(dp(indlev)).*T1day./3e5;

%plot
figure(1)
clf
colormap(polarmap(21))
%contourf(lon,lat,term1,[-5:0.5:5],'LineColor','none');
contourf(lon,lat,struct.du_h2_./5,[-5:0.5:5],'LineColor','none');
caxis([-5,5])

hold on
%colormap(jet(21))
d_divep=(divep_mjo-divep_ctrl);
%d_divep=divep_mjo./2;
d_divep=smoothn(d_divep,'robust');
contour(lon,lat,d_divep,[-5:0.5:5],'LineWidth',2);
caxis([-5,5])
title('DIV EP (forced-unforced)','Interpreter','none')
%title('DIV EP (forced)','Interpreter','none')
xlabel('lon')
ylabel('lat')
set(gca,'FontSize',18)

%plot
figure(10)
clf
colormap(polarmap(21))
%contourf(lon,lat,term1,[-5:0.5:5],'LineColor','none');
contourf(lon,lat,struct.du_h2_./5,[-5:0.5:5],'LineColor','none');
caxis([-5,5])

hold on
%colormap(jet(21))
d_ex=(ex_mjo-ex_ctrl);
%d_ex=ex_mjo./2;
d_ex=smoothn(d_ex,'robust');
contour(lon,lat,d_ex,[-5:0.5:5],'LineWidth',2);
caxis([-5,5])
title('Ex (forced-unforced)','Interpreter','none')
%title('Ex (forced)','Interpreter','none')
xlabel('lon')
ylabel('lat')
set(gca,'FontSize',18)

%plot
figure(100)
clf
colormap(polarmap(21))
%contourf(lon,lat,term1,[-5:0.5:5],'LineColor','none');
contourf(lon,lat,struct.du_h2_./5,[-5:0.5:5],'LineColor','none');
caxis([-5,5])

hold on
%colormap(jet(21))
d_ey=(ey_mjo-ey_ctrl);
%d_ey=ey_mjo./2;
d_ey=smoothn(d_ey,'robust');
contour(lon,lat,d_ey,[-5:0.5:5],'LineWidth',2);
caxis([-5,5])
title('Ey (forced-unforced)','Interpreter','none')
%title('Ey (forced)','Interpreter','none')
xlabel('lon')
ylabel('lat')
set(gca,'FontSize',18)