% create dye tracer initialization for NEMO CREG025E setup based on 3D model grid file
% to create new dyes, add new dye case definitions further below
% =========================================================================
close all; clearvars; clc;
% =========================================================================

% set input directory and input file dummy
inDir = fullfile(filesep, 'misc', '1', 'input', 'grosse', 'NEMO_input');
gridFile = 'DOMAIN_grid.nc';

% choose domain
domain = 'CREG025E';

% define name of dye configuration to be used in file containing mask with
% all tracers (useful for checking combined masks and for post-processing)
dyeConf = 'dye06_fg001';

% define output directory
outDir = inDir;

% define tracers: one tracer per row
%  -  first column: name of tracer as to be used in netCDF file
%  - second column: tracer name used for defining initial tracer
%                   distribution. For each tracer definition, a prescribed
%                   case has to be defined further down in 'dyeCase'
dyes = {'dye01', 'CREG025E_Pacific'      ; ...  % Pacific (south of Bering Strait)
        'dye02', 'CREG025E_Atlantic'     ; ...  % North Atlantic (south of Fram Strait, Barents Sea and Labrador Sea)
        'dye03', 'CREG025E_LabSeaBaffin' ; ...  % Labrador Sea and Baffin Bay
        'dye04', 'CREG025E_CanArchi'     ; ...  % Canadian Archipelago
        'dye05', 'CREG025E_EurasianShelf'; ...  % Eurasian shelf seas (Barents, Kara, Laptev, East Siberian)
        'dye06', 'CREG025E_Arctic'};

%% DYE CASE DEFINITIONS
% =========================================================================
% - all dye cases require:
%    (1) a UNIQUE 'name', which is referred to by the entries in the 2nd
%        column in 'dyes'; must be of form [GRID NAME]_[DYE REGION NAME]
%    (2) a dye 'type' defining how the regions are described, i.e. by:
%        (a) 'ijloc': a list of (i,j) index pairs that define all input
%                     locations
%        (b) 'ijbox': a box defined by min-max ranges of i and j indices;
%                     0 values imply that min=1 and max=length(dimension),
%                     respectively
%        (c) 'llpoly': a polygon defined by lon-lat points
%        (d) 'exclDyes': a list of other dyes that should be excluded from
%                        the region of the dye to be defined
%
%        NOTE: different types can be added; you only need to code their
%              treatment in the 'switch d.type' section further below
% 
%    (3) the regional information according to the defined dye type
%
% - additional definitions of depth and/or bathymetry ranges are optional
%    (4) 'zRange': upper depth < depth <= lower depth (positive values)
%    (5) 'bRange': lower limit < bathy <= upper limit (positive values)
%
% - (de-)activation and initialization of age is done by defining 'iniAge'
%    (6) 'iniAge' not existent or < 0 => no age
%        'iniAge' >= 0                => do age and initialize with 'iniAge'
%
% - (de-)activation of test of dye mask; useful during creation of new dye
%   initializations
%    (7) 'doTest' not existent or 0 => no test required; setup approved
%        'doTest' = 1               => do test, i.e. create figure of dye
%                                      mask for visual check
% =========================================================================

dyeCase = struct([]);
% dyeCase: 'CREG025E_Pacific'
i = 1;
dyeCase(i).name = 'CREG025E_Pacific';
dyeCase(i).type = 'ijbox';
dyeCase(i).ijbox = [0, 0; 602 ,0]; % [iMin, iMax; jMin, jMax]
dyeCase(i).iniAge = 0;
dyeCase(i).doTest = 0;
% dyeCase: 'CREG025E_Atlantic'
i = i + 1;
dyeCase(i).name = 'CREG025E_Atlantic';
dyeCase(i).type = 'llpoly';
dyeCase(i).llpoly = [-55.0, 53.0; ... % lon lat
                     -43.5, 60.0; ...
                     -43.5, 67.0; ...
                     -42.7, 77.2; ...
                     -13.8, 81.5; ...
                      11.5, 79.8; ...
                      16.6, 78.3; ...
                      16.2, 77.0; ...
                      19.3, 70.0; ...
                      34.4, 58.6; ...
                      44.9, 55.7; ...
                      38.0, 25.3; ...
                     -90.0, 25.3; ...
                     -92.0, 46.8; ...
                     -70.0, 53.0; ...
                     -55.0, 53.0];
dyeCase(i).iniAge = 0;
dyeCase(i).doTest = 0;
% dyeCase: 'CREG025E_LabSeaBaffin' => Labrador Sea and Baffin Bay
i = i + 1;
dyeCase(i).name = 'CREG025E_LabSeaBaffin';
dyeCase(i).type = 'llpoly';
dyeCase(i).llpoly = [-55.0, 53.0; ... % lon lat
                     -43.5, 60.0; ...
                     -43.5, 67.0; ...
                     -57.0, 82.0; ...
                     -63.0, 82.4; ...
                     -70.7, 81.1; ...
                     -83.1, 77.2; ...
                     -78.2, 76.7; ...
                     -79.5, 76.0; ...
                     -80.5, 75.4; ...
                     -80.5, 75.2; ...
                     -79.5, 74.9; ...
                     -77.6, 73.4; ...
                     -76.3, 73.1; ...
                     -76.2, 72.9; ...
                     -76.3, 71.3; ...
                     -75.0, 70.0; ...
                     -68.6, 66.9; ...
                     -69.5, 63.3; ...
                     -66.5, 62.1; ...
                     -64.5, 60.1; ...
                     -62.2, 53.1; ...
                     -55.0, 53.0];
dyeCase(i).iniAge = 0;
dyeCase(i).doTest = 0;
% % dyeCase: 'CREG025E_CanArchi' => Canadian Archipelago
i = i + 1;
dyeCase(i).name = 'CREG025E_CanArchi';
dyeCase(i).type = 'llpoly';
dyeCase(i).llpoly = [-117.0, 58.0; ... % lon lat
                     -125.0, 66.0; ...
                     -128.0, 70.6; ...
                     -125.7, 72.1; ...
                     -121.0, 72.8; ...
                     -124.2, 74.4; ...
                     -122.3, 76.1; ...
                     -117.0, 77.1; ...
                     -116.4, 77.7; ...
                     -113.4, 78.4; ...
                     -110.4, 78.6; ...
                     -104.8, 79.3; ...
                      -99.3, 80.1; ...
                      -95.5, 80.5; ...
                      -93.3, 81.3; ...
                      -90.7, 81.7; ...
                      -86.0, 81.0; ...
                      -63.0, 82.4; ...
                      -70.7, 81.1; ...
                      -83.1, 77.2; ...
                      -78.2, 76.7; ...
                      -79.5, 76.0; ...
                      -80.5, 75.4; ...
                      -80.5, 75.2; ...
                      -79.5, 74.9; ...
                      -77.6, 73.4; ...
                      -76.3, 73.1; ...
                      -76.2, 72.9; ...
                      -76.3, 71.3; ...
                      -75.0, 70.0; ...
                      -68.6, 66.9; ...
                      -69.5, 63.3; ...
                      -66.5, 62.1; ...
                      -64.5, 60.1; ...
                      -62.2, 53.1; ...
                      -55.0, 53.0; ...
                      -94.2, 47.0; ...
                     -103.0, 65.0; ...
                     -108.0, 66.0; ...
                     -117.0, 58.0];
dyeCase(i).iniAge = 0;
dyeCase(i).doTest = 0;
% dyeCase: 'CREG025E_EurasianShelf'
i = i + 1;
dyeCase(i).name = 'CREG025E_EurasianShelf';
dyeCase(i).type = 'llpoly';
dyeCase(i).llpoly = [ 11.5, 79.8; ... % lon lat
                      16.6, 78.3; ...
                      16.2, 77.0; ...
                      19.3, 70.0; ...
                      34.4, 58.6; ...
                      65.5, 63.0; ...
                     104.0, 54.5; ...
                     141.0, 65.0; ...
                     180.0, 67.0; ...
                     180.0, 71.1; ...
                     180.0, 75.0; ...
                     175.2, 76.1; ...
                     173.8, 76.4; ...
                     170.1, 77.1; ...
                     169.0, 77.8; ...
                     169.0, 78.2; ...
                     165.1, 78.3; ...
                     163.9, 78.4; ...
                     161.4, 78.9; ...
                     159.6, 79.0; ...
                     157.5, 79.3; ...
                     154.1, 79.5; ...
                     149.6, 80.0; ...
                     143.2, 79.6; ...
                     140.8, 79.1; ...
                     139.5, 79.0; ...
                     138.7, 78.8; ...
                     134.1, 78.3; ...
                     132.6, 78.1; ...
                     128.4, 77.6; ...
                     127.1, 77.4; ...
                     124.7, 77.1; ...
                     122.0, 77.3; ...
                     117.7, 77.3; ...
                     116.5, 77.5; ...
                     114.3, 78.1; ...
                     112.0, 78.5; ...
                     107.9, 79.1; ...
                     105.4, 79.7; ...
                     103.1, 80.3; ...
                     100.7, 81.0; ...
                      97.6, 81.3; ...
                      94.6, 81.8; ...
                      92.4, 81.9; ...
                      89.1, 81.9; ...
                      85.8, 82.0; ...
                      82.2, 82.1; ...
                      79.2, 82.1; ...
                      73.8, 81.7; ...
                      66.7, 82.0; ...
                      63.1, 82.1; ...
                      60.1, 82.5; ...
                      58.0, 82.6; ...
                      56.1, 82.7; ...
                      52.0, 82.5; ...
                      50.0, 82.4; ...
                      47.3, 82.3; ...
                      44.0, 82.0; ...
                      37.2, 81.3; ...
                      34.4, 81.3; ...
                      33.5, 81.3; ...
                      30.1, 81.4; ...
                      28.4, 81.2; ...
                      24.6, 81.2; ...
                      20.6, 80.5; ...
                      15.7, 79.0; ...
                      11.5, 79.8];
dyeCase(i).iniAge = 0;
dyeCase(i).doTest = 0;
% dyeCase: 'CREG025E_Arctic'
i = i + 1;
dyeCase(i).name = 'CREG025E_Arctic';
dyeCase(i).type = 'exclDyes';
dyeCase(i).exclDyes = {'CREG025E_Pacific', 'CREG025E_Atlantic', ...
                       'CREG025E_LabSeaBaffin', 'CREG025E_CanArchi', ...
                       'CREG025E_EurasianShelf'};
dyeCase(i).iniAge = 0;
dyeCase(i).doTest = 0;

% add other dye cases here as required

% =========================================================================
%% START PROGRAM
nDyes = size(dyes,1);

b_read = false;
z_read = false;
dyeMask = cell(1,nDyes);
for iDye = 1:nDyes
    indDye = strcmp({dyeCase.name},dyes{iDye,2});
    if sum(indDye)==0, error('Dye case not defined: %s.', dyes{iDye,2}); end
    nDye = find(indDye==1, 1, 'first');
    d = dyeCase(indDye);
    % check domains of dye and grid
    i_ = min(strfind(d.name, '_'));
    dyeName = strtrim(d.name(i_+1:end));
    if  ~strcmpi(domain,d.name(1:i_-1))
        error('Selected dye (''%s'') not defined for selected domain (''%s'').', ...
              strtrim(d.name), domain);
    end
    if  ~isfield(d,'type')
        error('%s - %s: dyeCase.type field does not exist.', dyes{iDye,:});
    else
        eval(sprintf('lType = isfield(d,''%s'');', d.type));
        if ~lType, error('%s - %s: dyeCase.%s field does not exist.', dyes{iDye,:}, d.type); end
    end
    % read basic grid information
    maskFile = strrep(fullfile(inDir,gridFile), 'DOMAIN', domain);
    if ~exist(maskFile, 'file'), error('File does not exist: %s.', maskFile); end
    if  iDye==1
        mask = ncread(maskFile, 'mask_t');
        lon = ncread(maskFile, 'nav_lon');
        lat = ncread(maskFile, 'nav_lat');
        iMax = size(mask,1);
        jMax = size(mask,2);
        dyeMaskAll = mask;
        dyeMaskAll(mask==0) = NaN;
        dyeMaskAll(mask==1) = 0;
    end
    % initialize dye mask
    dyeMask{nDye} = mask;
    % apply bathymetry range if defined
    if  isfield(d,'bRange')
        if  sum(d.bRange~=0)>0
            if  ~b_read
                b = ncread(maskFile,'bathy');
                b_read = true;
            end
            if  d.bRange(1)<=0
                b1 = -1;
            else
                b1 = d.bRange(1);
            end
            if  (d.bRange(2)<=0) || (d.bRange(2)>=nanmax(b(:)))
                b2 = nanmax(b(:));
            else
                b2 = d.bRange(2);
            end
            bMask = b > b1 & b <= b2;
            for i = 1:size(dyeMask{nDye},3)
                tmp = squeeze(dyeMask{nDye}(:,:,i));
                tmp(~bMask & tmp==1) = 0;
                dyeMask{nDye}(:,:,i) = tmp;
                clear tmp
            end
            clear b1 b2 bMask
            if iDye==nDyes, clear b; end
        end
    end
    % apply depth range if defined
    if  isfield(d,'zRange')
        if  sum(d.zRange~=0)>0
            if  ~z_read
                z_read = true;
                z = ncread(maskFile,'depth_t');
            end
            if  d.zRange(1)<=0
                z1 = 0;
            else
                z1 = d.zRange(1);
            end
            if  (d.zRange(2)<=0) || (d.zRange(2)>=nanmax(z(:)))
                z2 = nanmax(z(:));
            else
                z2 = d.zRange(2);
            end
            zMask = z > z1 & z <= z2;
            dyeMask{nDye}(~zMask & dyeMask{nDye}==1) = 0;
            clear z1 z2 zMask
            if iDye==nDyes, clear z; end
        end
    end
    switch d.type
        case 'ijloc'
            for i = 1:size(d.ijloc,1)
                dyeMask{nDye}(d.ijloc(i,1),d.ijloc(i,2),:) = 1;
            end
            dyeMask{nDye}(dyeMask{nDye}~=1) = 0;
        case 'ijbox'
            ii = repmat(1:iMax, jMax, 1);
            jj = repmat((1:jMax)', 1, iMax);
            if  d.ijbox(1,1)<=0
                i1 = 1;
            else
                i1 = min([iMax, d.ijbox(1,1)]);
            end
            if  (d.ijbox(1,2)<=0) || (d.ijbox(1,2)>=iMax)
                i2 = iMax;
            else
                i2 = min([iMax, d.ijbox(1,2)]);
            end
            if  d.ijbox(2,1)<=0
                j1 = 1;
            else
                j1 = min([jMax, d.ijbox(2,1)]);
            end
            if  (d.ijbox(2,2)<=0) || (d.ijbox(2,2)>=jMax)
                j2 = jMax;
            else
                j2 = min([jMax, d.ijbox(2,2)]);
            end
            ijMask = ii>=i1 & ii<=i2 & jj>=j1 & jj<=j2;
            for i = 1:size(dyeMask{nDye},3)
                tmp = squeeze(dyeMask{nDye}(:,:,i));
                tmp(~ijMask') = 0;
                dyeMask{nDye}(:,:,i) = tmp;
                clear tmp
            end
            clear ii jj i1 i2 j1 j2 ijMask
        case 'llpoly'
            polyMask = zeros(size(squeeze(dyeMask{nDye}(:,:,1))));
            for i = 1:size(lon,1)
                for j = 1:size(lon,2)
                    if  inpolygon(lon(i,j),lat(i,j), d.llpoly(:,1), d.llpoly(:,2))
                        polyMask(i,j) = 1;
                    end
                end
            end
            for i = 1:size(dyeMask{nDye},3)
                tmp = squeeze(dyeMask{nDye}(:,:,i));
                tmp(polyMask==0) = 0;
                dyeMask{nDye}(:,:,i) = tmp;
                clear tmp
            end
            clear polyMask usePoly
        case 'exclDyes'
            for id = 1:length(d.exclDyes)
                iExcl = strcmpi(d.exclDyes{id}, {dyes{:,2}});
                if sum(iExcl)==0, error('Dye to be excluded not found: %s', d.exclDyes{id}); end
                dyeMask{nDye}(dyeMask{iExcl}==1) = 0;
            end
        otherwise
            error('%s: dye.type = %s not defined.', d.name, d.type);
    end
    
    % store all dyes in one array
    dyeMaskAll(dyeMask{nDye}==1) = dyeMask{nDye}(dyeMask{nDye}==1)*nDye;
    
    if  isfield(d, 'doTest')
        if  d.doTest==1
            hFig = figure(nDye);
            set(hFig, 'Position', [50, 50, 1000, 1200]);
            plotMask = squeeze(dyeMask{nDye}(:,:,1));
            plotMask(mask(:,:,1)==0) = NaN;
            pcolor(plotMask');
            shading flat
            colormap(parula(2));
            isOK = input('Check figure. Is dye mask satisfactory? Yes (1) or no (0)? ');
            if isOK~=1, error('Correct section from-to locations for %s.', secDef(id).name); end
            close(hFig);
            clear isOK
        end
    end
    
    % do age
    doAge = false;
    if  isfield(d,'iniAge')
        if  d.iniAge>=0
            ageMask = dyeMask{nDye}*d.iniAge;
            doAge = true;
        end
    end
    
    % create netCDF file
    fileName = fullfile(outDir, sprintf('%s_%s_%s.nc', domain, dyes{iDye,1}, dyeName));
    cMode = bitor(netcdf.getConstant('64BIT_OFFSET'), netcdf.getConstant('CLOBBER'));
    ncid = netcdf.create(fileName, cMode);
    xid = netcdf.defDim(ncid, 'x', size(dyeMask{nDye},1));
    yid = netcdf.defDim(ncid, 'y', size(dyeMask{nDye},2));
    zid = netcdf.defDim(ncid, 'z', size(dyeMask{nDye},3));
    dye_id = netcdf.defVar(ncid, dyes{iDye,1}, 'NC_FLOAT', [xid, yid, zid]);
    netcdf.putAtt(ncid, dye_id, 'long_name' , sprintf('concentration of %s', dyes{iDye,1}));
    netcdf.putAtt(ncid, dye_id, 'dye_name' , sprintf('%s_%s', dyes{iDye,:}));
    netcdf.putAtt(ncid, dye_id, 'unit' , 'kg/m3');
    if  doAge
        age_id = netcdf.defVar(ncid, sprintf('%s_age', dyes{iDye,1}), 'NC_FLOAT', [xid, yid, zid]);
        netcdf.putAtt(ncid, age_id, 'long_name' , sprintf('age of %s', dyes{iDye,1}));
        netcdf.putAtt(ncid, age_id, 'unit' , '(kg s)/m3');
    end
    netcdf.endDef(ncid);
    % write data
    netcdf.putVar(ncid, dye_id, dyeMask{nDye});
    if doAge, netcdf.putVar(ncid, age_id, ageMask); end
    netcdf.close(ncid);
end

% store mask with all dyes in new netcdf file
fileName = fullfile(outDir, sprintf('%s_allDyes_%s.nc', domain, dyeConf));
cMode = bitor(netcdf.getConstant('64BIT_OFFSET'), netcdf.getConstant('CLOBBER'));
ncid = netcdf.create(fileName, cMode);
xid = netcdf.defDim(ncid, 'x', size(dyeMaskAll,1));
yid = netcdf.defDim(ncid, 'y', size(dyeMaskAll,2));
zid = netcdf.defDim(ncid, 'z', size(dyeMaskAll,3));
dye_id = netcdf.defVar(ncid, 'dyeMask', 'NC_FLOAT', [xid, yid, zid]);
netcdf.putAtt(ncid, dye_id, 'long_name' , 'mask with all dye regions');
netcdf.putAtt(ncid, dye_id, 'unit' , 'dimensionless');
netcdf.endDef(ncid);
% write data
netcdf.putVar(ncid, dye_id, dyeMaskAll);
netcdf.close(ncid);
