clear;close all;clc;

addpath('/qfs/people/xudo627/Setup-E3SM-Mac/matlab-scripts-for-mosart/');
addpath('/qfs/people/xudo627/Setup-E3SM-Mac/matlab-scripts-to-process-inputs/');
addpath('/qfs/people/xudo627/mylib/m');
re = 6.37122e6;% Earth radius

%[lon_gsim,lat_gsim,S,yr,mo,da,mu_gsim,sd,cv] = get_GSIM_discharge('BR_0002835',1);
load('Pantanal_shp.mat');



tag = 'Susquehanna_hydraulic';%'Susquehanna_run3';%'Pantanal_hydraulic2'; %'Pantanal_sur';%'Pantanal_geo'; %'Pantanal_all';
% run3       --> hydrology calibration
% hydraulic2 --> calibration hydraulic parameters after hydrology calibration 
if contains(tag,'Pantanal')
    load('Pantanal_shp.mat');
    aream2 = 5.7945e+11;%areaint(S.Y,S.X)*4*pi.*(re^2); %
elseif contains(tag,'Susquehanna')
    load('Susquehanna_shp.mat');
    aream2 = 6.5271e+10;
end
rng(1);

mosart_input_file = '/compyfs/inputdata/rof/mosart/MOSART_Global_half_20200720.nc';
elm_surface_file  = '/compyfs/inputdata/lnd/clm2/surfdata_map/surfdata_0.5x0.5_simyr2000_c190418.nc';
elm_domain_file   = '/compyfs/inputdata/share/domains/domain.lnd.r05_oEC60to30v3.190418.nc';

[ioutlet,contributing] = find_mosart_cell(mosart_input_file,lon_gsim,lat_gsim,aream2);
in = [ioutlet; contributing];
numc = length(in);

if strcmp(tag,'Pantanal_all')
    ntot = 1000;
    micro_sigma_default = ones(numc,1).*0.1;
    
    xtrain = NaN(ntot,6);
    xtrain(:,1)= 0.1 + (5-0.1)*rand(ntot,1);
    xtrain(:,2)= 0.1 + (5-0.1)*rand(ntot,1);
    xtrain(:,3)= 0.001 + (1-0.001)*rand(ntot,1); % micro_sigma
    xtrain(:,4)= 0.01 + (1-0.01)*rand(ntot,1); % fc   
    xtrain(:,5)= 2.6 + (20.2-2.6)*rand(ntot,1); % rwid [2.6 - 20.2]
    xtrain(:,6)= 0.12 + (0.63-0.12)*rand(ntot,1); % rdep  

elseif strcmp(tag,'Pantanal_geo')
    ntot = 100;
    micro_sigma_default = [];
    
    xtrain = NaN(ntot,6);
    xtrain(:,1)= ones(ntot,1).*2.5;
    xtrain(:,2)= ones(ntot,1).*0.5;
    xtrain(:,3)= ones(ntot,1).*0.1;
    xtrain(:,4)= ones(ntot,1).*0.4;  
    xtrain(:,5)= 2.6 + (20.2-2.6)*rand(ntot,1);   % rwid [2.6 - 20.2]
    xtrain(:,6)= 0.12 + (0.63-0.12)*rand(ntot,1); % rdep  
    
elseif strcmp(tag,'Pantanal_sub')
    ntot = 100;
    micro_sigma_default = [];
    
    xtrain = NaN(ntot,6);
    xtrain(:,1)= 0.1 + (5-0.1)*rand(ntot,1);
    xtrain(:,2)= ones(ntot,1).*0.5;
    xtrain(:,3)= ones(ntot,1).*0.1;
    xtrain(:,4)= ones(ntot,1).*0.4;  
    xtrain(:,5)= ones(ntot,1).*7.2;  
    xtrain(:,6)= ones(ntot,1).*0.26;   
    
elseif strcmp(tag,'Pantanal_sur')
    ntot = 500;
    micro_sigma_default = ones(numc,1).*0.1;
    
    xtrain = NaN(ntot,6);
    xtrain(:,1)= ones(ntot,1).*2.5;
    xtrain(:,2)= 0.1 + (5-0.1)*rand(ntot,1);
    xtrain(:,3)= 10.^(-3 + (0+3)*rand(1,ntot)); % micro_sigma
    xtrain(:,4)= 10.^(-3 + (-0.155+3)*rand(1,ntot)); %fc   
    xtrain(:,5)= ones(ntot,1).*7.2;  
    xtrain(:,6)= ones(ntot,1).*0.26; 

elseif strcmp(tag,'Pantanal_sur2')
    ntot = 100;
    micro_sigma_default = [];
    
    xtrain = NaN(ntot,6);
    xtrain(:,1)= ones(ntot,1).*2.5;
    xtrain(:,2)= 0.1 + (5-0.1)*rand(ntot,1);
    xtrain(:,3)= ones(ntot,1).*0.1;
    xtrain(:,4)= 10.^(-3 + (-0.155+3)*rand(1,ntot)); %fc   
    xtrain(:,5)= ones(ntot,1).*7.2;  
    xtrain(:,6)= ones(ntot,1).*0.26;  
    
elseif strcmp(tag,'Pantanal_sur3')
    ntot = 100;
    micro_sigma_default = [];
    
    xtrain = NaN(ntot,6);
    xtrain(:,1)= ones(ntot,1).*2.5;
    xtrain(:,2)= ones(ntot,1).*0.5;
    xtrain(:,3)=ones(ntot,1).*0.1;
    xtrain(:,4)= 0 + (0.5-0)*rand(ntot,1); % fc   
    xtrain(:,5)= ones(ntot,1).*7.2;  
    xtrain(:,6)= ones(ntot,1).*0.26; 
    
elseif strcmp(tag,'Pantanal_run')
    ntot = 500;
    micro_sigma_default = [];
    
    xtrain = NaN(ntot,6);
    xtrain(:,1)= 0.1 + (5-0.1)*rand(ntot,1); % fdrain
    xtrain(:,2)= 0.1 + (5-0.1)*rand(ntot,1);
    xtrain(:,3)= ones(ntot,1).*0.1;
    xtrain(:,4)= 0 + (0.5-0)*rand(ntot,1); % fc   
    xtrain(:,5)= ones(ntot,1).*7.2;  
    xtrain(:,6)= ones(ntot,1).*0.26;  
    
elseif strcmp(tag,'Pantanal_run2')
    ntot = 100;
    micro_sigma_default = [];
    
    xtrain = NaN(ntot,6);
    xtrain(:,1)= 0.1 + (5-0.1)*rand(ntot,1); % fdrain
    xtrain(:,2)=ones(ntot,1).*0.5; % fover
    xtrain(:,3)=ones(ntot,1).*0.1;
    xtrain(:,4)= 0 + (1-0)*rand(ntot,1); % fc   
    xtrain(:,5)= ones(ntot,1).*7.2;  
    xtrain(:,6)= ones(ntot,1).*0.26;  
elseif strcmp(tag,'Pantanal_run3') || strcmp(tag,'Susquehanna_run3')
    ntot = 1000;
    micro_sigma_default = ones(numc,1).*0.1;
    
    xtrain = NaN(ntot,6);
    xtrain(:,1)= 0.1 + (5-0.1)*rand(ntot,1);
    xtrain(:,2)= 0.1 + (5-0.1)*rand(ntot,1);
    xtrain(:,3)= 0.001 + (1-0.001)*rand(ntot,1); % micro_sigma
    xtrain(:,4)= 0.001 + (0.5-0.001)*rand(ntot,1); % fc   
    xtrain(:,5)= ones(ntot,1).*7.2;  
    xtrain(:,6)= ones(ntot,1).*0.26;  
elseif strcmp(tag,'Pantanal_run4')
    ntot = 1000;
    micro_sigma_default = ones(numc,1).*0.1;
    
    xtrain = NaN(ntot,6);
    xtrain(:,1)= 0.1 + (5-0.1)*rand(ntot,1);
    xtrain(:,2)= 0.1 + (5-0.1)*rand(ntot,1);
    xtrain(:,3)= 10.^(-3 + (0+3)*rand(1,ntot)); % micro_sigma
    xtrain(:,4)= 10.^(-3 + (-0.155+3)*rand(1,ntot)); %fc   
    xtrain(:,5)= ones(ntot,1).*7.2;  
    xtrain(:,6)= ones(ntot,1).*0.26;  
elseif strcmp(tag,'Pantanal_run5')
    ntot = 1000;
    micro_sigma_default = [];
    
    xtrain = NaN(ntot,6);
    xtrain(:,1)= 0.1 + (5-0.1)*rand(ntot,1);
    xtrain(:,2)= 0.1 + (5-0.1)*rand(ntot,1);
    xtrain(:,3)= ones(ntot,1).*0.1; % micro_sigma
    xtrain(:,4)= 10.^(-3 + (-0.155+3)*rand(1,ntot)); %fc   
    xtrain(:,5)= ones(ntot,1).*7.2;  
    xtrain(:,6)= ones(ntot,1).*0.26;  
    
elseif strcmp(tag,'Pantanal_hydrology')
    ntot = 200;
    micro_sigma_default = [];
    
    xtrain = NaN(ntot,6);
    xtrain(:,1)= 0.1 + (5-0.1)*rand(ntot,1);
    xtrain(:,2)= ones(ntot,1).*0.5; % fover
    xtrain(:,3)= ones(ntot,1).*0.1; % micro_sigma
    xtrain(:,4)= 0.001 + (0.5-0.001)*rand(ntot,1); % fc   
    xtrain(:,5)= ones(ntot,1).*7.2;  
    xtrain(:,6)= ones(ntot,1).*0.26;  
    
elseif strcmp(tag,'Pantanal_default')
    ntot = 10;
    micro_sigma_default = [];
    
    xtrain = NaN(ntot,6);
    xtrain(:,1)= ones(ntot,1).*2.5;
    xtrain(:,2)= ones(ntot,1).*0.5;
    xtrain(:,3)= ones(ntot,1).*0.1;
    xtrain(:,4)= ones(ntot,1).*0.4; 
    xtrain(:,5)= ones(ntot,1).*7.2;  
    xtrain(:,6)= ones(ntot,1).*0.26;  
elseif strcmp(tag,'Pantanal_man')
    ntot = 50;
    micro_sigma_default = [];
    
    xtrain = NaN(ntot,6);
    xtrain(:,1)= ones(ntot,1).*2.5;
    xtrain(:,2)= ones(ntot,1).*0.5;
    xtrain(:,3)= ones(ntot,1).*0.1;
    xtrain(:,4)= ones(ntot,1).*0.4; 
    xtrain(:,5)= ones(ntot,1).*7.2;  
    xtrain(:,6)= ones(ntot,1).*0.26; 

elseif strcmp(tag,'Pantanal_hydraulic')
    ntot = 200;
    micro_sigma_default = [];
    
    xtrain = NaN(ntot,6);
    xtrain(:,1)= ones(ntot,1).*2.5;
    xtrain(:,2)= ones(ntot,1).*0.5;
    xtrain(:,3)= ones(ntot,1).*0.1;
    xtrain(:,4)= ones(ntot,1).*0.4;  
    xtrain(:,5)= 2.6 + (20.2-2.6)*rand(ntot,1);   % rwid [2.6 - 20.2]
    xtrain(:,6)= 0.12 + (0.63-0.12)*rand(ntot,1); % rdep  
elseif strcmp(tag,'Pantanal_hydraulic2') || strcmp(tag,'Susquehanna_hydraulic')
    ntot = 1000;
    micro_sigma_default = [];
    
    xtrain = NaN(ntot,6);
    xtrain(:,1)= ones(ntot,1).*2.5;
    xtrain(:,2)= ones(ntot,1).*0.5;
    xtrain(:,3)= ones(ntot,1).*0.1;
    xtrain(:,4)= ones(ntot,1).*0.4;  
    xtrain(:,5)= 2.6 + (20.2-2.6)*rand(ntot,1);   % rwid [2.6 - 20.2]
    xtrain(:,6)= 0.12 + (0.63-0.12)*rand(ntot,1); % rdep  
elseif strcmp(tag,'Pantanal_calibration')
    ntot = 1;
    micro_sigma_default = ones(numc,1).*0.1;
    
    xtrain = NaN(ntot,6);
    xtrain(:,1)= ones(ntot,1).*2.5;
    xtrain(:,2)= ones(ntot,1).*0.5;
    xtrain(:,3)= ones(ntot,1).*0.1;
    xtrain(:,4)= ones(ntot,1).*0.4;  
    xtrain(:,5)= 2.6 + (20.2-2.6)*rand(ntot,1);   % rwid [2.6 - 20.2]
    xtrain(:,6)= 0.12 + (0.63-0.12)*rand(ntot,1); % rdep  
end

dx = 1/2; dy = 1/2;

xc_gridded   = ncread(elm_domain_file,'xc');
yc_gridded   = ncread(elm_domain_file,'yc');
frac_gridded = ncread(elm_domain_file,'frac');
mask_gridded = ncread(elm_domain_file,'mask');
latixy       = ncread(mosart_input_file,'latixy');
longxy       = ncread(mosart_input_file,'longxy');
areatotal2   = ncread(mosart_input_file,'areaTotal2');

xc   = xc_gridded(in);
yc   = yc_gridded(in);
frac = frac_gridded(in);
mask = mask_gridded(in);

xv = zeros(4,length(xc));
yv = zeros(4,length(xc));
xv(1,:) = xc - dx/2; xv(2,:) = xc + dx/2; 
xv(3,:) = xc + dx/2; xv(4,:) = xc - dx/2; 
yv(1,:) = yc - dx/2; yv(2,:) = yc - dx/2; 
yv(3,:) = yc + dx/2; yv(4,:) = yc + dx/2; 

out_netcdf_dir = 'inputdata_run';

if ~exist(out_netcdf_dir,'dir')
    mkdir(out_netcdf_dir);
end

mosart_usrdat_name = [tag '_template'];
fname_out1 = CreateMOSARTUgridInputForE3SM2(in, mosart_input_file, out_netcdf_dir, mosart_usrdat_name,1);
if strcmp(tag,'Pantanal_man')
    nr = ncread(fname_out1,'nr');
    nt = ncread(fname_out1,'nt');
    nh = ncread(fname_out1,'nh');
    nr_uq = repmat(nr,ntot,1);
    nt_uq = repmat(nt,ntot,1);
    nh_uq = repmat(nh,ntot,1);
    nr_prior = 0.01 : 0.01 : 0.2;
    nh_prior = 0.05 : 0.05 : 0.5;
    assert(ntot == 50);
    for i = 1 : ntot
        if i <= 20
            nr_uq((i-1)*numc+1:i*numc) = nr_prior(i);
        elseif i <= 40
            nt_uq((i-1)*numc+1:i*numc) = nr_prior(i-20);
        else
            nh_uq((i-1)*numc+1:i*numc) = nh_prior(i-40);
        end
    end
elseif strcmp(tag,'Pantanal_hydraulic')
    nr_uq = ones(numc*ntot,1).*0.05;
    nr = 0.01 + (0.2-0.01)*rand(ntot,1);
    for i = 1 : ntot
        nr_uq((i-1)*numc+1:i*numc) = nr(i);
    end
elseif strcmp(tag,'Pantanal_hydraulic2') || strcmp(tag,'Susquehanna_hydraulic')
    nr_uq = ones(numc*ntot,1).*0.05;
    nr = 0.01 + (0.2-0.01)*rand(ntot,1);
    for i = 1 : ntot
        nr_uq((i-1)*numc+1:i*numc) = nr(i);
    end
end

fname_out2 =  sprintf('%s/domain_lnd_%s%s.nc',out_netcdf_dir,tag,'_template');
if contains(tag,'Pantanal')
    area_tmp = ncread('inputdata_run/domain_lnd_Pantanal_sur_template.nc','area');
else
    [~,~,area_tmp] = xc2xv(xc,yc,0.5,0.5,true);
end

area = generate_lnd_domain(xc,yc,xv,yv,frac,mask,area_tmp,fname_out2);

% (2): generate ELM surface template input file
sand = NaN(numc,10);
clay = NaN(numc,10);
sandall = ncread(elm_surface_file,'PCT_SAND');
clayall = ncread(elm_surface_file,'PCT_CLAY');
for i = 1 : 10
    tmp = sandall(:,:,i);
    sand(:,i) = tmp(in);
    tmp = clayall(:,:,i);
    clay(:,i) = tmp(in);
end
clear sandall clayall;
% sand(~in2,:) = NaN;
% clay(~in2,:) = NaN;

fmax = ncread(elm_surface_file,'FMAX');
fmax = fmax(in);

silt = 100 - sand - clay;
bsw_mu    = 2.91 + 0.159.*clay;
sucsat_mu = 10 .* ( 10.^(1.88-0.0131.*sand) );
xksat_mu  = 0.0070556 .*( 10.^(-0.884+0.0153.*sand) );
watsat_mu = 0.01.*( 48.9 - 0.126.*sand );

fname_out3 = CreateCLMUgridSurfdatForE3SM(...
                    in,...
                    elm_surface_file,                             ...
                    out_netcdf_dir, ['_' tag '_template'],        ...
                    ones(numc,1).*2.5, ones(numc,1).*5.5e-3,      ... % fdrain,max_drain,
                    ones(numc,1).*6, ones(numc,1),                ... % ice_imped,snoalb_factor,
                    ones(numc,1).*0.5, fmax,                      ... % fover,fmax
                    bsw_mu,sucsat_mu, xksat_mu,watsat_mu,         ... % bsw,sucsat,xksat,watsat,
                    ones(numc,1).*0.4, ones(numc,1).*0.14,        ... % fc, mu
                    micro_sigma_default, []);     
%generte uq file for domain
fname_out4 = sprintf('%s/domain_lnd_%s.nc',out_netcdf_dir,[tag '_uq']);
generate_lnd_domain_uq(fname_out2,fname_out4,ntot);
%generte uq file for surface 
fdrai_uq     = ones(numc*ntot,1).*2.5;
fdqraimax_uq = ones(numc*ntot,1).*5.5e-3;
ice_imped_uq = ones(numc*ntot,1).*6;
fover_uq     = ones(numc*ntot,1).*0.5;
fmax_uq      = repmat(fmax,ntot,1);
bsw_uq       = repmat(bsw_mu,ntot,1);
sucsat_uq    = repmat(sucsat_mu,ntot,1);
xksat_uq     = repmat(xksat_mu,ntot,1);
watsat_uq    = repmat(watsat_mu,ntot,1);
fc_uq        = ones(numc*ntot,1).*0.4;
mu_uq        = ones(numc*ntot,1).*0.14;
micro_sigma_uq = NaN(numc*ntot,1);
kh2osfc_uq  = NaN(numc*ntot,1);

areaTotal2 = ncread(fname_out1,'areaTotal2');
load('w_d_global_ming.mat');

rwid = w_global(in);
rdep = d_global(in);

if contains(tag,'Pantanal')
    save('Pantanal_geometry.mat','rwid','rdep');
end
rdep_uq  = NaN(numc*ntot,1);
rwid_uq  = NaN(numc*ntot,1);

k = 1;
for i = 1 : ntot
    % fdrai
    tmp = ones(numc,1).* xtrain(i,1);
    fdrai_uq((k-1)*numc+1:k*numc,1) = tmp;
    % fover
    tmp = ones(numc,1).* xtrain(i,2);
    fover_uq((k-1)*numc+1:k*numc,1) = tmp;
    % micro_sgma
    tmp = ones(numc,1).* xtrain(i,3);
    micro_sigma_uq((k-1)*numc+1:k*numc,1) = tmp;
    % fc
    tmp = ones(numc,1).* xtrain(i,4);
    fc_uq((k-1)*numc+1:k*numc,1) = tmp;
    % rwid
    tmp = ones(numc,1).* xtrain(i,5);
    rwid_uq((k-1)*numc+1:k*numc,1) = rwid./7.2.*tmp;
    % rdep
    tmp = ones(numc,1).* xtrain(i,6);
    rdep_uq((k-1)*numc+1:k*numc,1) = rdep./0.26.*tmp;
    
    k = k + 1;
end

if strcmp(tag,'Pantanal_hydraulic2') || strcmp(tag,'Pantanal_calibration')
    load('Calibrated_hydrology.mat');
    fdrai_uq = repmat(fdrai,ntot,1);
    fover_uq = repmat(fover,ntot,1);
    micro_sigma_uq = repmat(sigma,ntot,1);
    fc_uq = repmat(fc,ntot,1);
end
if strcmp(tag,'Susquehanna_hydraulic')
    load('SUS_calibrated_hydrology.mat');
    fdrai_uq = repmat(fdrain,ntot,1);
    fover_uq = repmat(fover,ntot,1);
    micro_sigma_uq = repmat(micro_sigma,ntot,1);
    fc_uq = repmat(fc,ntot,1);
end

fname_out5 = sprintf('%s/surfdata_%s.nc',out_netcdf_dir,[tag '_uq']);
generate_lnd_surface_uq(fname_out3,fname_out5,ntot,fdrai_uq,fdqraimax_uq,ice_imped_uq,ones(numc*ntot,1),fover_uq, ...
                        fmax_uq,bsw_uq,sucsat_uq,xksat_uq,watsat_uq,fc_uq,mu_uq,micro_sigma_uq,kh2osfc_uq);
                    
%generte uq file for MOSART
fname_out6 = sprintf('%s/MOSART_%s.nc',out_netcdf_dir,[tag '_uq']);

ncid_inp = netcdf.open(fname_out1,'NC_NOWRITE');
ncid_out = netcdf.create(fname_out6,'NC_CLOBBER');
[ndims,nvars,ngatts,unlimdimid] = netcdf.inq(ncid_inp);
dimid(1) = netcdf.defDim(ncid_out,'gridcell',ntot*numc);
dimid(2) = netcdf.defDim(ncid_out,'nele',11);

% +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
%
%                           Define variables
%
% +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
for ivar = 1 : nvars
    [varname,xtype,dimids,natts] = netcdf.inqVar(ncid_inp,ivar-1);
    switch varname
        case {'ele'}
            varid(ivar) = netcdf.defVar(ncid_out,varname,xtype,dimids); 
        otherwise
            varid(ivar) = netcdf.defVar(ncid_out,varname,xtype,dimid(1)); 
    end
    varnames{ivar} = varname;
    for iatt = 1 : natts
        attname = netcdf.inqAttName(ncid_inp,ivar-1,iatt-1);
        attvalue = netcdf.getAtt(ncid_inp,ivar-1,attname);
        
        netcdf.putAtt(ncid_out,ivar-1,attname,attvalue);
    end
end
varid = netcdf.getConstant('GLOBAL');

[~,user_name]=system('echo $USER');
netcdf.putAtt(ncid_out,varid,'Created_by' ,user_name(1:end-1));
netcdf.putAtt(ncid_out,varid,'Created_on' ,datestr(now,'ddd mmm dd HH:MM:SS yyyy '));
netcdf.endDef(ncid_out);


% +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
%
%                           Copy variables
%
% +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
for ivar = 1:nvars
    [varname,vartype,vardimids,varnatts]=netcdf.inqVar(ncid_inp,ivar-1);
    data = netcdf.getVar(ncid_inp,ivar-1);
    switch varname
        case {'ID','dnID'}
            data1 = data;
            for i = 2 : ntot
                tmp = data1;
                for j = 1 : length(tmp)
                    if tmp(j) ~= -9999
                        tmp(j) = tmp(j) + (i-1)*numc;
                    end
                end
                data = [data; tmp];
            end
        case {'rdep'}
            data = rdep_uq;
        case {'rwid'}
            data = rwid_uq;
        case {'rwid0'}
            data = rwid_uq.*5;
        otherwise
            data = repmat(data,ntot,1);
    end
    netcdf.putVar(ncid_out,ivar-1,data);
end


netcdf.close(ncid_inp);
netcdf.close(ncid_out);
if strcmp(tag,'Pantanal_man')
    ncwrite(fname_out6,'nr',nr_uq);
    ncwrite(fname_out6,'nt',nt_uq);
    ncwrite(fname_out6,'nh',nt_uq);
end
if strcmp(tag,'Pantanal_hydraulic') || strcmp(tag,'Pantanal_hydraulic2') || strcmp(tag,'Susquehanna_hydraulic')
    ncwrite(fname_out6,'nr',nr_uq);
end

