% script to create a modified river forcing file based on an existing one
% by Fabian Grosse (fabian.grosse[at]dal.ca)
% last modified on: 28-Feb-2018

% define input/output file location, and input and output filenames
fileDir = fullfile(filesep, 'misc', '1', 'input', 'grosse', 'ROMS_forcing', 'roms854');
inFile = 'mch_river_1982_2016_Rdet_prov.nc';
outFile = 'mch_river_1982_2016_Rdet_prov_DISCHARGE_Miss50pc_Atch50pc_PTRACER_AGE.nc';

% Define river specifications to be changed by using netCDF variable name, 
% netCDF river mask variable name, and the mathematical operation and the
% value to be applied.
% Allowed operations are:
%  (1) 'MULT'  : multiplication with the provided factor (i.e. the value)
%  (2) 'WEIGH' : weighting with the provided value relative to the overall
%                sum of the selected variable (NOTE: only applicable to discharge)
%  (3) 'NEWVAR': add new variable to forcing file with one constant value
%                for all locations of the selected river (i.e. the defined
%                mask)
%
% NOTE: multiple DIFFERENT operations (e.g., weighting and multiplication) may be
%       applied to one variable

rivChanges = {'river_transport', 'river_Miss', 'WEIGH' , 0.5; ...
              'river_transport', 'river_Atch', 'WEIGH' , 0.5; ...
              'river_dye_01'   , 'river_Miss', 'NEWVAR', 1.0; ... % passive tracer
              'river_dye_02'   , 'river_Miss', 'NEWVAR', 0.0; ... % age tracer
              'river_dye_03'   , 'river_Atch', 'NEWVAR', 1.0; ... % passive tracer
              'river_dye_04'   , 'river_Atch', 'NEWVAR', 0.0};    % age tracer

% provide information to be included in output netCDF file's attributes
author = 'Fabian Grosse';
modificationInfo = 'river discharge is distributed 50:50 between Mississippi and Atchafalaya; passive tracers and age tracers added for both rivers';
          

%% START PROGRAM
% check changes to be applied
allowedChanges = {'MULT', 'WEIGH', 'NEWVAR'};
allowedStr = sprintf('''%s''', allowedChanges{1});
nAllowed = length(allowedChanges);
if  nAllowed==1
    allowedStr = sprintf('%s is', allowedStr);
elseif nAllowed==2
    allowedStr = sprintf('%s and ''%s'' are', allowedStr, allowedChanges{2});
else
    for ia = 2:nAllowed-1
        allowedStr = sprintf('%s, ''%s''', allowedStr, allowedChanges{ia});
    end
    allowedStr = sprintf('%s and ''%s'' are', allowedStr, allowedChanges{nAllowed});
end
nChanges = size(rivChanges,1);
% 1. Check validity of operations. Is same operation applied more than once?
for ic = 1:nChanges
    if  sum(contains(allowedChanges,rivChanges{ic,3}))==0
        error('Invalid operation: %s. Only %s defined.', rivChanges{ic,3}, allowedStr);
    end
    if  ~strcmp(rivChanges{ic,1}, 'river_transport') && strcmp(rivChanges{ic,3}, 'WEIGH')
        error('Weighting is only allowed for ''river_transport''.');
    end
    currChange = strrep(strcat(rivChanges{ic,1:3}), ' ', '');
    for icc = ic+1:nChanges
        compChange = strrep(strcat(rivChanges{icc,1:3}), ' ', '');
        if  strcmp(currChange,compChange)
            error('Same operation applied multiple times to same variable: %s, %s, %s.', rivChanges{ic,1:3});
        end
    end
end
% 2. Do weights for each variable add up to 1?
rivVars = unique(rivChanges(:,1));
nVars = length(rivVars);
for iVar = 1:nVars
    iRivWeighs = strcmp(rivChanges(:,3),'WEIGH') & strcmp(rivChanges(:,1),rivVars{iVar});
    if  (sum(iRivWeighs)>0) && (sum([rivChanges{iRivWeighs,4}])~=1)
        error('Weights for %s do not add up to 1.', rivVars{iVar});
    end
end

% make copy of netCDF input file with output filename
outFile = fullfile(fileDir,outFile);
if exist(outFile, 'file'), error('File already exists. Change output filename or remove existing file.'); end
[comStat, comMess] = unix(sprintf('cp %s %s', fullfile(fileDir,inFile), outFile));
if comStat~=0, error('%s', comMess); end

% get netCDF file information
ncInfo = ncinfo(outFile);
% check if all variables to be modified (and not newly added) are available in netCDF file
for iVar = 1:nVars
    % skip variables to be newly added
    iMod = ~strcmp(rivChanges(:,3),'NEWVAR') & strcmp(rivChanges(:,1),rivVars{iVar});
    if sum(iMod)==0, continue, end
    iNC = strcmp({ncInfo.Variables.Name}, rivVars{iVar});
    if sum(iNC)==0, error('Variable not found: %s.', rivVars{iVar}); end
end

% apply defined operations
% 1. multiplications
for iVar = 1:nVars
    iMultVars = find(strcmp(rivChanges(:,1), rivVars{iVar}) & strcmp(rivChanges(:,3), 'MULT'));
    if isempty(iMultVars), continue, end
    varDataIn = ncread(outFile, rivChanges{iMultVars(1),1});
    varDataOut = varDataIn;
    for iMask = 1:length(iMultVars)
        varMask = logical(ncread(outFile, rivChanges{iMultVars(iMask),2}));
        % apply factor to variable
        if  sum(numel(varDataIn)-numel(varMask))==0
            % same dimensions for variable and mask
            varDataOut(varMask) = varDataOut(varMask) * rivChanges{iMultVars(iMask),4};
        else
            varMask = squeeze(sum(varMask,2))>0;
            varDataOut(varMask) = varDataOut(varMask) * rivChanges{iMultVars(iMask),4};
        end
    end
    ncwrite(outFile, rivChanges{iMultVars(1),1}, varDataOut);
end

% 2. weighting
% step 1: get weights for different input locations of same river
% step 2: distribute overall variable to rivers according to defined weights
% step 3: distribute variable for each river over input locations according
% to 'internal' weights calculated during step 1
for iVar = 1:nVars
    iWeightVars = find(strcmp(rivChanges(:,1), rivVars{iVar}) & strcmp(rivChanges(:,3), 'WEIGH'));
    if isempty(iWeightVars), continue, end
    % read variable to be modified and get absolute values and signs
    varDataIn = ncread(outFile, rivChanges{iWeightVars(1),1});
    varDataSign = sign(varDataIn);
    varDataAbs = abs(varDataIn);
    nDims = ndims(varDataIn);
    if nDims~=2, error('Weighting only defined for 2D variables (dims: river locations & time).'); end
    % get rivers' internal weights
    nRivLocs = size(varDataIn,1);
    nRivs = length(iWeightVars);
    varMask = cell(1,nRivs);
    rivWeights = NaN(size(varDataIn));
    for iRiv = 1:nRivs
        varMask{iRiv} = logical(ncread(outFile, rivChanges{iWeightVars(iRiv),2}));
        if  sum(numel(varDataIn)-numel(varMask{iRiv}))~=0
            % different dimensions for variable and mask
            varMask{iRiv} = squeeze(sum(varMask{iRiv},2))>0;
        end
        for it = 1:size(varDataIn,nDims)
            maskVec = squeeze(varMask{iRiv}(:,it));
            rivWeights(maskVec,it) = varDataAbs(maskVec,it)/sum(varDataAbs(maskVec,it));
            % sum of weights might differ slightly from 1 (due to single precision)
            % if so, correct to get sum = 1
            if  sum(rivWeights(maskVec,it))~=1
                rivWeights(maskVec,it) = rivWeights(maskVec,it)/sum(rivWeights(maskVec,it));
            end
        end
    end
    
    % get overall variable
    totalVar = sum(varDataAbs,1);
    
    % calculate output data
    % 1. apply internal weights to total mass
    varDataOut = rivWeights .* repmat(totalVar,nRivLocs,1); % all positive
    % 2. apply user-defined weights to each river group
    for iRiv = 1:nRivs
        varDataOut(varMask{iRiv}) = varDataOut(varMask{iRiv}) * rivChanges{iWeightVars(iRiv),4};
    end
    % 3. check if overall sum (over all rivers and all time steps) of
    %    output and input values are identical (within the order of data precision)
    %    if difference > 0 but within data precision, apply correction to
    %    minimize difference
    varSumOut = sum(varDataOut(:));
    varSumIn = sum(varDataAbs(:));
    if  varSumOut~=varSumIn
        if  abs(varSumOut-varSumIn)/varSumIn>1.e-5
            error('Difference between total input and output variables larger than data precision.');
        else
            warning('Small difference between total input and output variables (smaller than data precision). Apply correction.');
            varDataOut = varDataOut * varSumIn/varSumOut;
        end
    end
    % 4. apply sign to output variable and write data
    varDataOut = varDataOut .* varDataSign;
    ncwrite(outFile, rivChanges{iWeightVars(1),1}, varDataOut);
end

% 3. add new variables to file
for iVar = 1:nVars
    iAddVars = find(strcmp(rivChanges(:,1), rivVars{iVar}) & strcmp(rivChanges(:,3), 'NEWVAR'));
    if isempty(iAddVars), continue, end
    varDataIn = ncread(outFile, rivChanges{iAddVars(1),2});
    iNC = strcmp({ncInfo.Variables.Name}, rivChanges{iAddVars(1),2});
    dimNames = {ncInfo.Variables(iNC).Dimensions.Name};
    dimLength = {ncInfo.Variables(iNC).Dimensions.Length};
    outDims = cell(1,2*length(dimNames));
    for id = 1:length(dimNames)
        outDims{2*id-1} = dimNames{id};
        outDims{2*id} = dimLength{id};
    end
    varDataOut = varDataIn;
    for iMask = 1:length(iAddVars)
        varMask = logical(ncread(outFile, rivChanges{iAddVars(iMask),2}));
        % apply factor to variable
        if  sum(numel(varDataIn)-numel(varMask))==0
            % same dimensions for variable and mask
            varDataOut(varMask) = rivChanges{iAddVars(iMask),4};
        else
            varMask = squeeze(sum(varMask,2))>0;
            varDataOut(varMask) = rivChanges{iAddVars(iMask),4};
        end
    end
    nccreate(outFile, rivChanges{iAddVars(1),1}, 'DataType', 'single', 'Dimensions', outDims)
    ncwrite(outFile, rivChanges{iAddVars(1),1}, varDataOut);
end

% update global attributes
iAuth = strcmp({ncInfo.Attributes.Name}, 'author');
oldAuthor = ncInfo.Attributes(iAuth).Value;
if  ~contains(oldAuthor,author)
    authors = sprintf('%s and %s', oldAuthor, author);
    ncwriteatt(outFile, '/', 'author', authors);
end
% add modification information
[~, host] = unix('echo $HOSTNAME');
host = strrep(host, newline, '');
modTime = datestr(datenum(clock), 'dd-mmm-yyyy HH:MM:SS');
iMod = strcmp({ncInfo.Attributes.Name}, 'modified');
if  sum(iMod)==0
    % 'modified' attribute does not exist => create new one
    ncwriteatt(outFile, '/', 'modified', sprintf('by %s on %s on %s', author, host, modTime));
else
    % 'modified' attribute does exist => update existing one
    oldMod = ncInfo.Attributes(iMod).Value;
    newMod = sprintf('%s\nby %s on %s on %s', strrep(oldMod, newline, ''), author, host, modTime);
    ncwriteatt(outFile, '/', 'modified', newMod);
end
iMod = strcmp({ncInfo.Attributes.Name}, 'modification');
if  sum(iMod)==0
    % 'modification' attribute does not exist => create new one
    ncwriteatt(outFile, '/', 'modification', modificationInfo);
else
    % 'modification' attribute does exist => update existing one
    oldMod = ncInfo.Attributes(iMod).Value;
    newMod = sprintf('%s\n%s', strrep(oldMod, newline, ''), modificationInfo);
    ncwriteatt(outFile, '/', 'modification', newMod);
end