% convert E3SM output into standard SCM intercomparison output format
% note that there is time shift for some output variables to be consistent
% with the forcing.


% clear all;
% infile = ['../SCM_data/mc3e_nonudging.e3sm.h0.2011-04-24-00000.nc'];
% casename='mc3e';
function outfile=convert_E3SM_format(infile,casename,modeloption)

%addpath 'C:\Users\tang32\OneDrive - LLNL\MATLAB\library';


daystr = strsplit(infile,'.');
daystr = daystr{end-1};
daystr2=strsplit(daystr,'-');
year=daystr2{1};
mon=daystr2{2};
day=daystr2{3};

%outfile = ['../SCM_data/' casename '_e3sm_' modeloption '.' year mon day '.000000.nc'];
outfile = ['/global/cscratch1/sd/sqtang/SCM_runs/SCM_data_final/' casename '_e3sm_' modeloption '.' year mon day '.000000.nc'];
disp(outfile);

levnew = [25:25:1000]';
levlen = length(levnew);

VARNAME_in = {'PRECT','PRECC','TMQ','TREFHT','TS',...
    'FSNTOA','FSUTOA','FSUTOAC','FLUT','FLUTC',...
    'FSNS','FSNSC','FSDS','FSDSC','FLNS','FLDS','FLDSC',...
    'SHFLX','LHFLX','PS','U10','N/A','TGCLDCWP','TGCLDIWP','AQSNOW',...
    'CLDTOT','QREFHT','PBLH','N/A','N/A',...
    'U','V','T','Q','RELHUM','Z3','OMEGA',...
    'CLOUDFRAC_CLUBB','CONCLD','CLOUD','N/A','N/A','N/A','N/A','CLDLIQ','CLDICE',...
    'N/A','RHO_CLUBB','RCM_CLUBB','N/A','N/A','N/A',...
    'ZMMU','ZMMD','PTTEND','PTEQ','QRS','QRL','TTEND_TOT',...
    'DIVT','ZMDT','CMFDT','STEND_CLUBB','MPDT','TTGW',...
    'N/A','DIVQ','ZMDQ','CMFDQ','RVMTEND_CLUBB','MPDQ','EVAPQZM',...
    'EVAPTZM','QVADV','TVADV'}; % additional terms for calculation
VARNAME_out = {'pr','prc','prw','tas','ts',...
    'rsdt','rsut','rsutcs','rlut','rlutcs',...
    'rsus','rsuscs','rsds','rsdscs','rlus','rlds','rldscs',...
    'hfss','hfls','ps','uas','vas','clwvi','clivi','cliviall',...
    'clt','qs','pblh','mrsos','orog',...
    'ua','va','ta','hus','hur','zg','wap',...
    'cls','clc','cl','clws','clwc','clis','clic','clw','cli',...
    'pfull','rho','sclwc','dclwc','sclic','dclic',...
    'mcu','mcd','q1','q2','tntrsw','tntrlw','tnt',...
    'tnta','tntdp','tntsh','tntpbl','tntlscp','tntd',...
    'tnhus','tnhusa','tnhusdp','tnhussh','tnhuspbl','tnhuslscp','tnhusd'};
varlen_in = length(VARNAME_in);
varlen_out = length(VARNAME_out);
VARall=struct('var',{},'long_name',{},'units',{});

%% read in data

ncid = netcdf.open(infile, 'nc_nowrite');
varid = netcdf.inqVarID(ncid,'lat');
lat = netcdf.getVar(ncid,varid,'double');
varid = netcdf.inqVarID(ncid,'lon');
lon = netcdf.getVar(ncid,varid,'double');
varid = netcdf.inqVarID(ncid,'lev');
lev = netcdf.getVar(ncid,varid,'double');
varid = netcdf.inqVarID(ncid,'P0');
p0 = netcdf.getVar(ncid,varid,'double');
varid = netcdf.inqVarID(ncid,'PS');
ps = netcdf.getVar(ncid,varid,'double');
varid = netcdf.inqVarID(ncid,'time');
time = netcdf.getVar(ncid,varid,'double');
unit_time = netcdf.getAtt(ncid,varid,'units');
varid = netcdf.inqVarID(ncid,'hyam');
hyam = netcdf.getVar(ncid,varid,'double');
varid = netcdf.inqVarID(ncid,'hybm');
hybm = netcdf.getVar(ncid,varid,'double');
varid = netcdf.inqVarID(ncid,'hyai');
hyai = netcdf.getVar(ncid,varid,'double');
varid = netcdf.inqVarID(ncid,'hybi');
hybi = netcdf.getVar(ncid,varid,'double');
latlen = length(lat);
lonlen = length(lon);

% timenew = [0:120]./24;  % days from day 0 to day 5
timenew = [time(1)*24:time(end)*24]./24;  % days for all
timelen = length(timenew);
 
 % vars in SE dyncore has dimesion of (1,p,t). shift to (1,1,p,t)
 if ndims(ps)==2
     ps=shiftdim(ps,-1);
 end

for kk=1:varlen_in
    try
        varid = netcdf.inqVarID(ncid,VARNAME_in{kk});
    catch
   %     disp(VARNAME_in{kk});
        VARall{kk}.var=NaN;
        VARall{kk}.units = 'N/A';
        VARall{kk}.long_name = 'Not Available';
        continue;
    end
    data_sigma = netcdf.getVar(ncid,varid,'double');
 % vars in SE dyncore has dimesion of (1,p,t). shift to (1,1,p,t)
    data_sigma = shiftdim(data_sigma,-1); 
    t_dim = ndims(data_sigma);
    % change vertical dimension
    if t_dim==4    % have lev dimension
        if size(data_sigma,3)==size(hyam,1)   % vertical coordinate is lev
         	data_pressure = convert_sigma_to_pres(data_sigma, ps, levnew, hyam, hybm);
      	elseif size(data_sigma,3)==size(hyai,1)   % vertical coordinate is ilev
         	data_pressure = convert_sigma_to_pres(data_sigma, ps, levnew, hyai, hybi);
        else
         	error('vertical dimension does not match');
        end
	elseif t_dim==3
        data_pressure=data_sigma;
    end

    % some variables have 1 or 2 timesteps lagging. shift back
    % the index bug has been fixed in Aug 2019. The timesteps lagging is only 0 or 1
    % variables with 1 timestep lagging is because of their output position is after time integration
    switch VARNAME_in{kk}
%        case {'TS','SHFLX','LHFLX'} % 1-d data with 1 timestep shift
%            data_pressure = cat(3,data_pressure(:,:,2:end),NaN*data_pressure(:,:,end));
%        case {'PS','U10','FLUT','FLUTC','FLNS','FLDS','FLDSC'} % 1-d data with 2 timesteps shift
%            data_pressure = cat(3,data_pressure(:,:,3:end),NaN*data_pressure(:,:,1:2));
%        case {'U','V','T','Q','RELHUM','Z3','OMEGA'}    %2d data with 2 timesteps shift
%            data_pressure = cat(4,data_pressure(:,:,:,3:end),NaN*data_pressure(:,:,:,1:2));

        case {'PS','U10','FLUT','FLUTC','FLNS','FLDS','FLDSC'} % 1-d data
            data_pressure = cat(3,data_pressure(:,:,2:end),NaN*data_pressure(:,:,1));
        case {'U','V','T','Q','RELHUM','Z3','OMEGA'}    %2d data
            data_pressure = cat(4,data_pressure(:,:,:,2:end),NaN*data_pressure(:,:,:,1));

    end

    % time average
    % model output is accumulated, adjust to centered at the time stamp
    switch t_dim
        case 3
            tmp_data = NaN(1,1,timelen);
            for tt=1:timelen
                t1=timenew(tt)-0.51/24;     % increase time range a little to ensure including data at the time
                t2=timenew(tt)+0.51/24;
                idx = time>t1 & time<t2;
                tmp_data(1,1,tt) = mean(data_pressure(1,1,idx),3,'omitnan');
            end
        case 4
            tmp_data = NaN(1,1,levlen,timelen);
            for tt=1:timelen
                t1=timenew(tt)-0.51/24;
                t2=timenew(tt)+0.51/24;
                idx = time>t1 & time<t2;
                tmp_data(1,1,:,tt) = mean(data_pressure(1,1,:,idx),4,'omitnan');
            end
        otherwise
            error('dimension number is neither 3 or 4');
    end
    VARall{kk}.var=tmp_data;
    VARall{kk}.units = netcdf.getAtt(ncid,varid,'units');
    VARall{kk}.long_name = netcdf.getAtt(ncid,varid,'long_name');
end
netcdf.close(ncid);

% figure;plot(squeeze(VARall{11}.var))
% warning('Check FSNS and FSNSC!');

%% change some variables from input to output

idx = (strcmp(VARNAME_out,'pr'));
VARall{idx}.var = VARall{idx}.var*1000*86400; % m/s to mm/day
VARall{idx}.units = 'mm/day';
VARall{idx}.long_name = 'Total Precipitation';

idx = (strcmp(VARNAME_out,'prc'));
VARall{idx}.var = VARall{idx}.var*1000*86400; % m/s to mm/day
VARall{idx}.units = 'mm/day';
VARall{idx}.long_name = 'Convective Precipitation';

idx = (strcmp(VARNAME_out,'rsdt'));
swtoa_net = VARall{strcmp(VARNAME_in,'FSNTOA')}.var;
swtoa_up = VARall{strcmp(VARNAME_in,'FSUTOA')}.var;
VARall{idx}.var = (swtoa_net+swtoa_up); % sw_dn = sw_net+sw_up
VARall{idx}.units = 'W/m2';
VARall{idx}.long_name = 'TOA insident shortwave radiation';

idx = (strcmp(VARNAME_out,'rsus'));
swsfc_net = VARall{strcmp(VARNAME_in,'FSNS')}.var;
swsfc_dn = VARall{strcmp(VARNAME_in,'FSDS')}.var;
VARall{idx}.var = (swsfc_dn-swsfc_net); % sw_up = sw_dn-sw_net
VARall{idx}.units = 'W/m2';
VARall{idx}.long_name = 'Surface upward shortwave radiation';

idx = (strcmp(VARNAME_out,'rsuscs'));
swsfc_net_cs = VARall{strcmp(VARNAME_in,'FSNSC')}.var;
swsfc_dn_cs = VARall{strcmp(VARNAME_in,'FSDSC')}.var;
VARall{idx}.var = (swsfc_dn_cs-swsfc_net_cs); % sw_up = sw_dn-sw_net
VARall{idx}.units = 'W/m2';
VARall{idx}.long_name = 'Surface upward shortwave radiation in clearsky';

idx = (strcmp(VARNAME_out,'rlus'));
lwsfc_net = VARall{strcmp(VARNAME_in,'FLNS')}.var;
lwsfc_dn = VARall{strcmp(VARNAME_in,'FLDS')}.var;
VARall{idx}.var = (lwsfc_dn+lwsfc_net); % lw_up = lw_dn+lw_net
VARall{idx}.units = 'W/m2';
VARall{idx}.long_name = 'Surface upward longwave radiation';

idx = (strcmp(VARNAME_out,'rldscs'));
VARall{idx}.units = 'W/m2';
VARall{idx}.long_name = 'Surface downward longwave radiation in clearsky';

%----------------- two options to calculate us and vs ----------------
% calculate us and vs using U10 and wind direction of the lowest level of u and v
u = VARall{strcmp(VARNAME_in,'U')}.var;
v = VARall{strcmp(VARNAME_in,'V')}.var;
wspd = VARall{strcmp(VARNAME_in,'U10')}.var;
us = NaN(1,1,timelen);
vs = NaN(1,1,timelen);
for tt=1:timelen
    zz=find(~isnan(squeeze(u(1,1,:,tt))),1,'last');
    if isempty(zz)
        continue;
    end    
    us(1,1,tt) = u(1,1,zz,tt);
    vs(1,1,tt) = v(1,1,zz,tt);
end
idx = (strcmp(VARNAME_out,'uas'));
VARall{idx}.var = wspd./sqrt(us.^2+vs.^2).*us;
VARall{idx}.units = 'm/s';
VARall{idx}.long_name = 'Surface u-wind';
idx = (strcmp(VARNAME_out,'vas'));
VARall{idx}.var = wspd./sqrt(us.^2+vs.^2).*vs;
VARall{idx}.units = 'm/s';
VARall{idx}.long_name = 'Surface v-wind';

% % us and vs are using the lowest level of u and v
% idx = (strcmp(VARNAME_out,'uas'));
% u = VARall{strcmp(VARNAME_in,'U')}.var;
% us = NaN(1,1,timelen);
% for tt=1:timelen
%     zz=find(~isnan(squeeze(u(1,1,:,tt))),1,'last');
%     us(1,1,tt) = u(1,1,zz,tt);
% end
% VARall{idx}.var = us;
% VARall{idx}.units = 'm/s';
% VARall{idx}.long_name = 'Surface u-wind';
% idx = (strcmp(VARNAME_out,'vas'));
% v = VARall{strcmp(VARNAME_in,'V')}.var;
% vs = NaN(1,1,timelen);
% for tt=1:timelen
%     zz=find(~isnan(squeeze(v(1,1,:,tt))),1,'last');
%     vs(1,1,tt) = v(1,1,zz,tt);
% end
% VARall{idx}.var = vs;
% VARall{idx}.units = 'm/s';
% VARall{idx}.long_name = 'Surface v-wind';
%-------------------------------------------------------------------

idx = (strcmp(VARNAME_out,'mrsos'));
VARall{idx}.var = NaN(1,1,timelen);
VARall{idx}.units = 'kg/m2';
VARall{idx}.long_name = '0-10cm Soil Moisture';

idx = (strcmp(VARNAME_out,'orog'));
VARall{idx}.var = NaN(1,1,timelen);
VARall{idx}.units = 'm';
VARall{idx}.long_name = 'Surface altitude';

idx = (strcmp(VARNAME_out,'clt'));
VARall{idx}.var = VARall{idx}.var *100;
VARall{idx}.units = '%';

idx = (strcmp(VARNAME_out,'cl'));
VARall{idx}.var = VARall{idx}.var *100;
VARall{idx}.units = '%';

% % no stratiform cloud fraction; use cloud fraction in CLUBB
idx = (strcmp(VARNAME_out,'cls'));
VARall{idx}.long_name = 'Cloud fraction in CLUBB';

% no stratiform or convective cloud liquid and ice
idx = strcmp(VARNAME_out,'clws');
VARall{idx}.var = NaN(1,1,levlen,timelen);
VARall{idx}.units = 'kg/kg';
VARall{idx}.long_name = 'Stratiform cloud liquid water';
idx = strcmp(VARNAME_out,'clwc');
VARall{idx}.var = NaN(1,1,levlen,timelen);
VARall{idx}.units = 'kg/kg';
VARall{idx}.long_name = 'Convective cloud liquid water';
idx = strcmp(VARNAME_out,'clis');
VARall{idx}.var = NaN(1,1,levlen,timelen);
VARall{idx}.units = 'kg/kg';
VARall{idx}.long_name = 'Stratiform cloud ice water';
idx = strcmp(VARNAME_out,'clic');
VARall{idx}.var = NaN(1,1,levlen,timelen);
VARall{idx}.units = 'kg/kg';
VARall{idx}.long_name = 'Convective cloud ice water';

idx = strcmp(VARNAME_out,'pfull');
p = repmat(levnew,1,timelen)*100;   % hPa to Pa
VARall{idx}.var = reshape(p,1,1,levlen,timelen);
VARall{idx}.units = 'Pa';
VARall{idx}.long_name = 'pressure';

% idx = strcmp(VARNAME_out,'rho');
% VARall{idx}.var = NaN(1,1,levlen,timelen);
% VARall{idx}.units = 'kg/m3';
% VARall{idx}.long_name = 'air density';

% idx = strcmp(VARNAME_out,'sclwc');
% VARall{idx}.var = NaN(1,1,levlen,timelen);
% VARall{idx}.units = 'kg/kg';
% VARall{idx}.long_name = 'Mass fraction of shallow convective cloud liquid water';
idx = strcmp(VARNAME_out,'dclwc');
VARall{idx}.var = NaN(1,1,levlen,timelen);
VARall{idx}.units = 'kg/kg';
VARall{idx}.long_name = 'Mass fraction of deep convective cloud liquid water';
idx = strcmp(VARNAME_out,'sclic');
VARall{idx}.var = NaN(1,1,levlen,timelen);
VARall{idx}.units = 'kg/kg';
VARall{idx}.long_name = 'Mass fraction of shallow convective cloud ice water';
idx = strcmp(VARNAME_out,'dclic');
VARall{idx}.var = NaN(1,1,levlen,timelen);
VARall{idx}.units = 'kg/kg';
VARall{idx}.long_name = 'Mass fraction of deep convective cloud ice water';

% idx = strcmp(VARNAME_out,'mcu');
% VARall{idx}.var = NaN(1,1,levlen,timelen);
% VARall{idx}.units = 'kg/m2/s';
% VARall{idx}.long_name = 'convective updraft mass flux';
 idx = strcmp(VARNAME_out,'mcd');
 VARall{idx}.var = -VARall{idx}.var;
% VARall{idx}.units = 'kg/m2/s';
% VARall{idx}.long_name = 'convective downdraft mass flux';
% idx = strcmp(VARNAME_out,'q1');
% VARall{idx}.var = NaN(1,1,levlen,timelen);
% VARall{idx}.units = 'K/s';
% VARall{idx}.long_name = 'Apparent heat source';

% no snow water path, so set cliviall is missing
 idx = strcmp(VARNAME_out,'cliviall');
 ice = VARall{strcmp(VARNAME_in,'CLDICE')}.var;
 snow = VARall{strcmp(VARNAME_in,'AQSNOW')}.var;
 ice_all = ice+snow;
 ice_all(isnan(ice_all))=0;
 data_tmp = sum(ice_all*2500/9.8,3); %kg/kg to kg/m2
 data_tmp=reshape(data_tmp,1,1,timelen);
 VARall{idx}.var = data_tmp;
 VARall{idx}.units = 'kg/m2';
 VARall{idx}.long_name = 'Ice + Snow water path';


idx = strcmp(VARNAME_out,'q2');
VARall{idx}.var = - VARall{idx}.var;    % moisture sink is -dq/dt
% VARall{idx}.units = 'kg/kg/s';
% VARall{idx}.long_name = 'Apparent moisture sink';

idx = strcmp(VARNAME_out,'tnta');
T_hori = VARall{strcmp(VARNAME_in,'DIVT')}.var;
T_vert = VARall{strcmp(VARNAME_in,'TVADV')}.var;
VARall{idx}.var = T_hori+T_vert;
VARall{idx}.units = 'K/s';
VARall{idx}.long_name = 'change rate of temperature due to total advection';
% idx = strcmp(VARNAME_out,'tntc');
% VARall{idx}.var = NaN(1,1,levlen,timelen);
% VARall{idx}.units = 'K/s';
% VARall{idx}.long_name = 'change rate of temperature due to convection';
idx = strcmp(VARNAME_out,'tntpbl');
VARall{idx}.var = VARall{idx}.var/1000; % STEND_CLUBB need to divided by 1000
VARall{idx}.units = 'K/s';
VARall{idx}.long_name = 'change rate of temperature due to CLUBB scheme';
idx = strcmp(VARNAME_out,'tntlscp');
VARall{idx}.var = VARall{idx}.var/1005.7;
VARall{idx}.units = 'K/s';
VARall{idx}.long_name = 'change rate of temperature due to microphysics';
idx = strcmp(VARNAME_out,'tntd');
T_drag = VARall{strcmp(VARNAME_in,'TTGW')}.var;
T_evap = VARall{strcmp(VARNAME_in,'EVAPTZM')}.var;
VARall{idx}.var = T_drag+T_evap; 
VARall{idx}.units = 'K/s';
VARall{idx}.long_name = 'change rate of temperature due to other diabatic processes';

idx = strcmp(VARNAME_out,'tnt');
VARall{idx}.var = T_hori+T_vert+VARall{strcmp(VARNAME_in,'PTTEND')}.var;
VARall{idx}.units = 'K/s';
VARall{idx}.long_name = 'Total change rate of temperature';


idx = strcmp(VARNAME_out,'tnhus');
Q_hori = VARall{strcmp(VARNAME_in,'DIVQ')}.var;
Q_vert = VARall{strcmp(VARNAME_in,'QVADV')}.var;
Q_phy = VARall{strcmp(VARNAME_in,'PTEQ')}.var;
VARall{idx}.var = Q_hori+Q_vert+Q_phy;
VARall{idx}.units = 'kg/kg/s';
VARall{idx}.long_name = 'Total change rate of specific humidity';
idx = strcmp(VARNAME_out,'tnhusa');
VARall{idx}.var = Q_hori+Q_vert;
VARall{idx}.units = 'kg/kg/s';
VARall{idx}.long_name = 'change rate of specific humidity due to total advection';
% idx = strcmp(VARNAME_out,'tnhusc');
% VARall{idx}.var = NaN(1,1,levlen,timelen);
% VARall{idx}.units = 'kg/kg/s';
% VARall{idx}.long_name = 'change rate of specific humidity due to convection';
idx = strcmp(VARNAME_out,'tnhuspbl');
% VARall{idx}.var = NaN(1,1,levlen,timelen);
VARall{idx}.units = 'kg/kg/s';
VARall{idx}.long_name = 'change rate of specific humidity due to CLUBB';
% idx = strcmp(VARNAME_out,'tnhuslscp');
% VARall{idx}.var = NaN(1,1,levlen,timelen);
% VARall{idx}.units = 'kg/kg/s';
% VARall{idx}.long_name = 'change rate of specific humidity due to large-scale cloud and precipitation';
idx = strcmp(VARNAME_out,'tnhusd');
% VARall{idx}.var = NaN(1,1,levlen,timelen);
% VARall{idx}.units = 'kg/kg/s';
VARall{idx}.long_name = 'change rate of specific humidity due to other diabatic processes';

%% write data
% ---------- write data ---------------
ncid = netcdf.create(outfile,'NC_WRITE');
% dimension
latdimID = netcdf.defDim(ncid,'lat',latlen);
londimID = netcdf.defDim(ncid,'lon',lonlen);
levdimID = netcdf.defDim(ncid,'lev',levlen);
tdimID  = netcdf.defDim(ncid,'time',timelen);
% variable
timeid = netcdf.defVar(ncid,'time','double',tdimID);
latid = netcdf.defVar(ncid,'lat','double',latdimID);
lonid = netcdf.defVar(ncid,'lon','double',londimID);
levid = netcdf.defVar(ncid,'lev','double',levdimID);
varid=NaN(varlen_out,1);
for kk=1:varlen_out
    t_dim = ndims(VARall{kk}.var);
    if t_dim==4
        varid(kk)=netcdf.defVar(ncid,VARNAME_out{kk},'float',[londimID,latdimID,levdimID,tdimID]);
    else
        varid(kk)=netcdf.defVar(ncid,VARNAME_out{kk},'float',[londimID,latdimID,tdimID]);
    end
end
netcdf.endDef(ncid);
netcdf.putVar(ncid,lonid,lon);
netcdf.putVar(ncid,latid,lat);
netcdf.putVar(ncid,levid,levnew);
netcdf.putVar(ncid,timeid,timenew);
for kk=1:varlen_out
    tmp_data=VARall{kk}.var;
    tmp_data(isnan(tmp_data))=-9999.0;
    netcdf.putVar(ncid,varid(kk),tmp_data);
end

% parameters
netcdf.reDef(ncid);
netcdf.putAtt(ncid,timeid,'long_name','time');
netcdf.putAtt(ncid,timeid,'units',unit_time);
netcdf.putAtt(ncid,latid,'long_name','latitude');
netcdf.putAtt(ncid,latid,'units','degrees_north');
netcdf.putAtt(ncid,lonid,'long_name','longitude');
netcdf.putAtt(ncid,lonid,'units','degrees_east');
netcdf.putAtt(ncid,levid,'long_name','pressure');
netcdf.putAtt(ncid,levid,'units','hPa');
netcdf.putAtt(ncid,levid,'positive','down');
for kk=1:varlen_out
    netcdf.putAtt(ncid,varid(kk),'units',VARall{kk}.units);
    netcdf.putAtt(ncid,varid(kk),'long_name',VARall{kk}.long_name);
    netcdf.putAtt(ncid,varid(kk),'missing_value',-9999.);
end
% global attributes
varid = netcdf.getConstant('GLOBAL');
netcdf.putAtt(ncid,varid,'Conventions','CF-1.0');
netcdf.putAtt(ncid,varid,'description','processed by Shuaiqi Tang (tang32@llnl.gov)');
netcdf.close(ncid);

