%% 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';

Day0=1;
Day1=20;
file0='vinth_timeavg_CTRL_ENSavg.nc';
file1=strcat('vinth_timeavg_',MJONAME,'_ENSavg.nc');
file_trans0='vinth_CTRL_H2_ENS800.nc';
file_trans1=strcat('vinth_',MJONAME,'_H2_ENS800.nc');
%file_trans1=strcat('vinth_','MJOB','_H2_ENS800.nc');

% Day0=18;
% Day1=20;
% file0='vinth_timeavg_CTRL_ENSavg.nc';
% file1=strcat('vinth_lasttwomean_',MJONAME,'_reynold_term_ENSavg.nc');
% file_trans0='vinth_CTRL_H2_ENS800.nc';
% %file_trans1=strcat('vinth_',MJONAME,'_H2_ENS800.nc');
% file_trans1=strcat('vinth_','MJOB','_H2_ENS800.nc');

%% calculate E vector
L=Day1-Day0+1;
T1day=86400.;
d2r=pi/180.;
a=6371000.;
f0=4*pi/86400;
g=9.8;
tau=40*86400;
kappa=2./7.;

% prepare coordinate
lat=squeeze(double(ncread(file0,'lat')));
lon=squeeze(double(ncread(file0,'lon')));
lev=squeeze(double(ncread(file0,'lev')));
ilev=squeeze(double(ncread(file0,'ilev')));
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=squeeze(double(ncread(file1,'UU')))-squeeze(double(ncread(file0,'UU')));
vu=squeeze(double(ncread(file1,'VU')))-squeeze(double(ncread(file0,'VU')));
vv=squeeze(double(ncread(file1,'VV')))-squeeze(double(ncread(file0,'VV')));
vt=squeeze(double(ncread(file1,'VT')))-squeeze(double(ncread(file0,'VT')));
T_CTRL=squeeze(double(ncread(file0,'T')));
U_CTRL=squeeze(double(ncread(file0,'U')));

Tbar=repmat(mean(T_CTRL),[nlon,1,1]);
TH_p=(Tbar(:,:,ip1)-Tbar(:,:,ip0))./(pres(:,:,ip1)-pres(:,:,ip0))+Tbar*(-kappa)./pres;

%EP flux and divergence
EPx=(vv-uu);
EPy=-vu;
EPz=f0.*coslat_.*vt./TH_p;

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));
  
% \Delta U
if Day0==1
    dU_h2=squeeze(double(ncread(file_trans1,'U',[1,1,1,Day1],[Inf,Inf,Inf,1])))-squeeze(double(ncread(file_trans0,'U',[1,1,1,L],[Inf,Inf,Inf,1])));
else
    dU_h2=squeeze(double(ncread(file_trans1,'U',[1,1,1,Day1],[Inf,Inf,Inf,1])))-squeeze(double(ncread(file_trans1,'U',[1,1,1,Day0-1],[Inf,Inf,Inf,1])));
end
dU_h2=dU_h2./L;

% eddy energy gain
Ux=1/a.*seclat_.*(U_CTRL(ix1,:,:)-U_CTRL(ix0,:,:))./(dlon);
Uy=1/a.*seclat_.*(U_CTRL(:,iy1,:).*coslat_(:,iy1,:)-U_CTRL(:,iy0,:).*coslat_(:,iy0,:))./(dlat);
EddyEnergyGain=EPx.*(Ux)+EPy.*(Uy);
  
%vertical integrate
dp_=permute(repmat(dp,[1,length(lon),length(lat)]),[2,3,1]);
indlev=find(lev<levshowbot&lev>levshowtop);
du_h2=squeeze(nansum(dU_h2(:,:,indlev).*dp_(:,:,indlev),3))'./sum(dp(indlev)).*10;
divep=squeeze(nansum(DIVEP(:,:,indlev).*dp_(:,:,indlev),3))'./sum(dp(indlev)).*T1day;
eddyenergygain=squeeze(nansum(EddyEnergyGain(:,:,indlev).*dp_(:,:,indlev),3))'./sum(dp(indlev)).*T1day./10;
ex=squeeze(nansum(EPx(:,:,indlev).*dp_(:,:,indlev),3))'./sum(dp(indlev)).*T1day/1e7;
ey=squeeze(nansum(EPy(:,:,indlev).*dp_(:,:,indlev),3))'./sum(dp(indlev)).*T1day/1e6;
ez=squeeze(nansum(EPz(:,:,indlev).*dp_(:,:,indlev),3))'./sum(dp(indlev)).*T1day/1e6;

%% plot \Delta U and \Delta div(E)
figure(111)
clf
colormap(polarmap(21))
contourf(lon,lat,du_h2,[-5:0.5:5],'LineColor','none');
caxis([-5,5])

hold on
%colormap(jet(21))
contour(lon,lat,divep,[-5:0.5:5],'LineWidth',2);
caxis([-5,5])
title('DIV EP (forced-unforced, same ref)','Interpreter','none')
xlabel('lon')
ylabel('lat')
set(gca,'FontSize',18)

%% plot \Delta U and \Delta eddy energy gain
figure(122)
clf
colormap(polarmap(21))
contourf(lon,lat,du_h2,[-5:0.5:5],'LineColor','none');
caxis([-5,5])

hold on
%colormap(jet(21))
contour(lon,lat,eddyenergygain,[-5:0.5:5],'LineWidth',2);
caxis([-5,5])
title('eddy energy gain (forced-unforced, same ref)','Interpreter','none')
xlabel('lon')
ylabel('lat')
set(gca,'FontSize',18)

%% plot  \Delta U and Ex
figure(133)
clf
colormap(polarmap(21))
contourf(lon,lat,du_h2,[-5:0.5:5],'LineColor','none');
caxis([-5,5])

hold on
%colormap(jet(21))
contour(lon,lat,ex,[-5:0.5:5],'LineWidth',2);
caxis([-5,5])
title('Ex (forced-unforced, same ref)','Interpreter','none')
xlabel('lon')
ylabel('lat')
set(gca,'FontSize',18)

% plot  \Delta U and Ey
figure(144)
clf
colormap(polarmap(21))
contourf(lon,lat,du_h2,[-5:0.5:5],'LineColor','none');
caxis([-5,5])

hold on
%colormap(jet(21))
contour(lon,lat,ey,[-5:0.5:5],'LineWidth',2);
caxis([-5,5])
title('Ey (forced-unforced, same ref)','Interpreter','none')
xlabel('lon')
ylabel('lat')
set(gca,'FontSize',18)

% plot  \Delta U and Ez
figure(55)
clf
colormap(polarmap(21))
contourf(lon,lat,du_h2,[-5:0.5:5],'LineColor','none');
caxis([-5,5])

hold on
%colormap(jet(21))
contour(lon,lat,ez,[-5:0.5:5],'LineWidth',2);
caxis([-5,5])
title('Ez (forced-unforced, same ref)','Interpreter','none')
xlabel('lon')
ylabel('lat')
set(gca,'FontSize',18)

