% convert E3SM output into standard SCM intercomparison output format

% clear all;
%infile=['/global/cscratch1/sd/sqtang/SCM_runs/scam5out/pecan_nonudging.cam.h0.2015-07-13-00000.nc'];
%casename='pecan';
%modeloption='nonudging';
function outfile=convert_SCAM5_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 '_scam5_' modeloption '.' year mon day '.000000.nc'];
outfile = ['/global/cscratch1/sd/sqtang/SCM_runs/SCM_data_final/' casename '_scam5_' modeloption '.' year mon day '.000000.nc'];

disp(outfile);

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

VARNAME_in = {'PRECT','PRECC','TMQ','TREFHT','TS',...
    'SOLIN','FSUTOA','FSNTOAC','FLUT','FLUTC',...
    'FSNS','FSNSC','FSDS','FSDSC','FLNS','FLDS','FLNSC',...
    'SHFLX','LHFLX','PS','U10','N/A','TGCLDCWP','TGCLDIWP','AQSNOW',...
    'CLDTOT','QREFHT','PBLH','N/A','N/A',...
    'U','V','T','Q','RELHUM','Z3','OMEGA',...
    'N/A','CONCLD','CLOUD','CLDLIQSTR','CLDICESTR','CLDLIQCON','CLDICECON','CLDLIQ','CLDICE',...
    'N/A','N/A','N/A','N/A','N/A','N/A',...
    'ZMMU','ZMMD','PTTEND','PTEQ','QRS','QRL','TTEND',...
    'DIVT','ZMDT','CMFDT','DTV','MPDT','EVAPTZM',...
    'TEQ','DIVQ','ZMDQ','CMFDQ','VD01','MPDQ','EVAPQZM',...
    'MACPDT','QVADV','TVADV','MACPDQ','PTTEND'}; % 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);
    
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');
    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 timestep lagging. shift back
    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','FLUT','FLUTC','FLNS','FLDS','FLDSC','U10'} % 1-d data with 2 timesteps shift
            data_pressure = cat(3,data_pressure(:,:,3:end),NaN*data_pressure(:,:,1:2));
        case {'U','V'}    %2d data with 2 timesteps shift
            data_pressure = cat(4,data_pressure(:,:,:,3:end),NaN*data_pressure(:,:,:,1:2));
        case {'OMEGA','T','Q','RELHUM','Z3'}    %2d data with 1 timesteps shift. T/Q/RH/Z are not sure
            data_pressure = cat(4,data_pressure(:,:,:,2:end),NaN*data_pressure(:,:,:,end));
    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;
                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
    
% for radiative fluxes, SCAM5 only call radiation once per 3 time steps
% so radiation needs to multiply by 3
    if any(strcmp(VARNAME_in{kk},{'SOLIN','FSUTOA','FSNTOAC',...
    'FLUT','FLUTC','FSNS','FSNSC','FSDS','FSDSC','FLNS','FLDS','FLNSC'}))
        VARall{kk}.var=tmp_data*3;
    else
        VARall{kk}.var=tmp_data;
    end
    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,'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,'rsutcs'));
swtoa_net_cs = VARall{strcmp(VARNAME_in,'FSNTOAC')}.var;
swtoa_dn_cs = VARall{strcmp(VARNAME_in,'SOLIN')}.var;
VARall{idx}.var = (swtoa_dn_cs-swtoa_net_cs); % sw_up = sw_dn-sw_net
VARall{idx}.units = 'W/m2';
VARall{idx}.long_name = 'TOA upward shortwave radiation in clearsky';

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'));
lwsfc_up_cs = VARall{strcmp(VARNAME_in,'FLNS')}.var;    % already replaced with rlus above
lwsfc_net_cs = VARall{strcmp(VARNAME_in,'FLNSC')}.var;
VARall{idx}.var = (lwsfc_up_cs-lwsfc_net_cs); % lw_dn = lw_up-lw_net
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
idx = (strcmp(VARNAME_out,'cls'));
VARall{idx}.var = NaN(1,1,levlen,timelen);
VARall{idx}.units = '%';
VARall{idx}.long_name = 'Stratiform cloud fraction';

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

% 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,'mcd');
 VARall{idx}.var = -VARall{idx}.var;

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

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 advection';
idx = strcmp(VARNAME_out,'tntpbl');
% VARall{idx}.var = NaN(1,1,levlen,timelen);
% VARall{idx}.units = 'K/s';
VARall{idx}.long_name = 'change rate of temperature due to boundary layer';
idx = strcmp(VARNAME_out,'tntlscp');
T_micro = VARall{strcmp(VARNAME_in,'MPDT')}.var;
T_macro = VARall{strcmp(VARNAME_in,'MACPDT')}.var;
VARall{idx}.var = (T_micro+T_macro)/1005.7;
VARall{idx}.units = 'K/s';
VARall{idx}.long_name = 'change rate of temperature due to large-scale cloud and precipitation (micro+macro)';
idx = strcmp(VARNAME_out,'tntd');
T_evap = VARall{strcmp(VARNAME_in,'EVAPTZM')}.var;
VARall{idx}.var = T_evap; 
VARall{idx}.units = 'K/s';
VARall{idx}.long_name = 'change rate of temperature due to other diabatic processes';

idx = strcmp(VARNAME_out,'tnhusa');
Q_hori = VARall{strcmp(VARNAME_in,'DIVQ')}.var;
Q_vert = VARall{strcmp(VARNAME_in,'QVADV')}.var;
VARall{idx}.var = Q_hori+Q_vert;
VARall{idx}.units = 'kg/kg/s';
VARall{idx}.long_name = 'change rate of specific humidity due to advection';
% 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 boundary layer';
idx = strcmp(VARNAME_out,'tnhuslscp');
Q_micro = VARall{strcmp(VARNAME_in,'MPDQ')}.var;
Q_macro = VARall{strcmp(VARNAME_in,'MACPDQ')}.var;
VARall{idx}.var = Q_micro+Q_macro;
VARall{idx}.units = 'kg/kg/s';
VARall{idx}.long_name = 'change rate of specific humidity due to large-scale cloud and precipitation (micro+macro)';
idx = strcmp(VARNAME_out,'tnhusd');
Q_evap = VARall{strcmp(VARNAME_in,'EVAPQZM')}.var;
VARall{idx}.var = Q_evap;
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);

