function md=FlattenModelStructure(md,operation)

% FlattenModelStructure

% Mike Wolovick, 8/1/2022

% This function takes in an ISSM model structure and "flattens" it.  IE,
% shrinks all unneccessarily 3D variables into 2D variables.  Alternately,
% it can take a flattened model structure and unflatten it.

% This is different from the ISSM built-in "collapse" command because this
% function preserves both the 3D mesh and the variables that actually
% contain 3D information.

% Inputs:

% md               ISSM model structure
% operation     string equal to 'flatten' or 'expand'

% Output:
% md              ISSM model structure

%% Variables List:

% This section is a hardcoded list of variables that can be flattened.

% The format is a cell array of strings, including the substructures, but
% not including the inital "md.".  For example:
% FlattenVariables={friction.coefficient';'results.TransientSolution.Surface'};

% If you want to include variables that are not always part of a model
% structure, or that are only sometimes fully 3D, that is fine.  The
% function always checks to make sure that the relevant substructure and
% variable exists before operating on it.  In addition, the function checks
% if the variable goes on vertices or elements, and it skips variables that
% are a length it is not expecting.  The function also
% checks if a variable is a cell array, and, if so, flattens every element
% of the array (helpful for thermal forcing).  Finally, if there are
% multiple elements of the last substructure, the model loops through all
% of them (so it loops through all of md.results.TransientSolution).

% The function assumes that all of these variables are vertically uniform,
% and takes the value from the basal layer.  The only exception is
% md.thermal.spctemperature, for which the function takes the surface layer
% and fills in NaN for everything else.

FlattenVariables={'mesh.x';...
    'mesh.y';...
    'mesh.vertexonboundary';...
    'mask.ocean_levelset';...
    'mask.ice_levelset';...
    'geometry.surface';...
    'geometry.thickness';...
    'geometry.base';...
    'geometry.bed';...
    'smb.mass_balance';...
    'basalforcings.groundedice_melting_rate';...
    'basalforcings.floatingice_melting_rate';...
    'basalforcings.geothermalflux';...
    'basalforcings.basin_id';...
    'basalforcings.tf';...
    'basalforcings.deepwater_melting_rate';...
    'basalforcings.deepwater_elevation';...
    'basalforcings.upperwater_melting_rate';...
    'basalforcings.upperwater_elevation';...
    'basalforcings.perturbation_melting_rate';...
    'friction.coefficient';...
    'friction.p';...
    'friction.q';...
    'friction.effective_pressure';...
    'friction.C';...
    'friction.m';...
    'flowequation.vertex_equation';...
    'flowequation.element_equation';...
    'flowequation.borderSSA';...
    'flowequation.borderHO';...
    'flowequation.borderFS';...
    'initialization.watercolumn';...
    'hydrology.spcwatercolumn';...
    'masstransport.spcthickness';...
    'thermal.spctemperature';...
    'calving.calvingrate';...
    'calving.stress_threshold_floatingice';...
    'calving.stress_threshold_groundedice';...
    'levelset.spclevelset';...
    'frontalforcings.meltingrate';...
    'inversion.cost_functions_coefficients';...
    'inversion.min_parameters';...
    'inversion.max_parameters';...
    'inversion.vx_obs';...
    'inversion.vy_obs';...
    'inversion.vel_obs';...
    'inversion.thickness_obs';...
    'inversion.surface_obs';...
    'results.StressbalanceSolution.FrictionCoefficient';...
    'results.StressbalanceSolution.MaterialsRheologyBbar';...
    'results.TransientSolution.Watercolumn';...
    'results.TransientSolution.BasalforcingsGroundediceMeltingRate';...
    'results.TransientSolution.BasalforcingsFloatingiceMeltingRate';...
    'results.TransientSolution.BasalforcingsIsmip6TfShelf';...
    'results.TransientSolution.Thickness';...
    'results.TransientSolution.Surface';...
    'results.TransientSolution.Base';...
    'results.TransientSolution.MaskOceanLevelset';...
    'results.TransientSolution.MaskIceLevelset';...
    'results.TransientSolution.SmbMassBalance';...
    'results.TransientSolution.IceMaskNodeActivation';...
    'results.TransientSolution.CalvingCalvingrate';...
    'results.TransientSolution.CalvingMeltingrate'};  % Cell array of strings
    
% Check requested operation:
if ~strcmp(operation,'flatten') && ~strcmp(operation,'expand')
    error('Parameter "operation" must be a string equal to "flatten" or "expand".')
end

% Check that the mesh is 3D:
if dimension(md.mesh)~=3
    warning('Flattening or expanding a model structure only works for a 3D mesh, exiting without doing anything');
    return
end

%% Work:

% Get mesh sizes:
numvertices2d=md.mesh.numberofvertices2d;
numelements2d=md.mesh.numberofelements2d;
numvertices=md.mesh.numberofvertices;
numelements=md.mesh.numberofelements;
numlayers=md.mesh.numberoflayers;

% Loop through variables:
numvariables=size(FlattenVariables,1);
for ii=1:numvariables
    % Locate periods in variable name:
    periodnums=strfind(FlattenVariables{ii},'.');
    periodnums=[0,periodnums];
    % Check that model substructure exists:
    thissubstructureexists=1;
    for jj=2:length(periodnums)
        try 
            eval(['test=length(md.',FlattenVariables{ii}(1:periodnums(jj)-1),');'])
        catch
            thissubstructureexists=0;
            break
        end
    end
    if ~thissubstructureexists
        continue
    end
    % Check that variable exists and get length:
    try
        eval(['test=length(md.',FlattenVariables{ii}(1:periodnums(end)-1),'(1).',FlattenVariables{ii}(periodnums(end)+1:end),');'])
    catch
        continue
    end
    % Get length of this substructure:
    eval(['thislength1=length(md.',FlattenVariables{ii}(1:periodnums(end)-1),');'])
    % Loop through substructures:
    for jj=1:thislength1
        % Check if this variable is a cell:
        eval(['thisiscell=iscell(md.',FlattenVariables{ii}(1:periodnums(end)-1),'(1).',FlattenVariables{ii}(periodnums(end)+1:end),');'])
        if thisiscell
            eval(['thislength2=length(md.',FlattenVariables{ii}(1:periodnums(end)-1),'(jj).',FlattenVariables{ii}(periodnums(end)+1:end),');'])
        else
            thislength2=1;
        end
        % Loop through cell elements:
        for kk=1:thislength2
            % Generate cell string:
            if thisiscell
                thiscellstring=['{',num2str(kk),'}'];
            else
                thiscellstring='';
            end
            % Get size of this variable:
            eval(['[thisysize,thisxsize]=size(md.',FlattenVariables{ii}(1:periodnums(end)-1),'(jj).',FlattenVariables{ii}(periodnums(end)+1:end),thiscellstring,');'])
            % Check if we are flattening or expanding:
            if strcmp(operation,'flatten')
                % Check size to determine if we have a time vector:
                if thisysize==numvertices+1 || thisysize==numelements+1
                    % Get the time vector:
                    eval(['ThisTimeVector=md.',FlattenVariables{ii}(1:periodnums(end)-1),'(jj).',FlattenVariables{ii}(periodnums(end)+1:end),thiscellstring,'(end,:);'])
                elseif thisysize==numvertices || thisysize==numelements
                    % Empty time vector:
                    ThisTimeVector=[];
                else
                    % Size is wrong to flatten:
                    %disp(['Variable md.',FlattenVariables{ii}(1:periodnums(end)-1),'(',num2str(jj),').',FlattenVariables{ii}(periodnums(end)+1:end),thiscellstring,' is the wrong size to flatten.'])
                    continue
                end
                % Check if this is surface temperature or not:
                if ~strcmp(FlattenVariables{ii},'thermal.spctemperature')
                    % Check if this is on vertices or elements:
                    if thisysize==numvertices || thisysize==numvertices+1
                        % Flatten the variable, take values from the bed:
                        eval(['md.',FlattenVariables{ii}(1:periodnums(end)-1),'(jj).',FlattenVariables{ii}(periodnums(end)+1:end),thiscellstring,...
                            '=[md.',FlattenVariables{ii}(1:periodnums(end)-1),'(jj).',FlattenVariables{ii}(periodnums(end)+1:end),thiscellstring,'(1:numvertices2d,:);ThisTimeVector];'])
                    else
                        % Flatten the variable, take values from the bed:
                        eval(['md.',FlattenVariables{ii}(1:periodnums(end)-1),'(jj).',FlattenVariables{ii}(periodnums(end)+1:end),thiscellstring,...
                            '=[md.',FlattenVariables{ii}(1:periodnums(end)-1),'(jj).',FlattenVariables{ii}(periodnums(end)+1:end),thiscellstring,'(1:numelements2d,:);ThisTimeVector];'])
                    end
                else
                    % Flatten the variable, take values from the surface:
                    eval(['md.',FlattenVariables{ii}(1:periodnums(end)-1),'(jj).',FlattenVariables{ii}(periodnums(end)+1:end),thiscellstring,...
                        '=[md.',FlattenVariables{ii}(1:periodnums(end)-1),'(jj).',FlattenVariables{ii}(periodnums(end)+1:end),thiscellstring,'((numlayers-1)*numvertices2d+1:numlayers*numvertices2d,:);ThisTimeVector];'])
                end
            else
                % Check size to determine if we have a time vector:
                if thisysize==numvertices2d+1 || thisysize==numelements2d+1
                    % Get the time vector:
                    eval(['ThisTimeVector=md.',FlattenVariables{ii}(1:periodnums(end)-1),'(jj).',FlattenVariables{ii}(periodnums(end)+1:end),thiscellstring,'(end,:);'])
                elseif thisysize==numvertices2d || thisysize==numelements2d
                    % Empty time vector:
                    ThisTimeVector=[];
                else
                    % Size is wrong to flatten:
                    %disp(['Variable md.',FlattenVariables{ii}(1:periodnums(end)-1),'(',num2str(jj),').',FlattenVariables{ii}(periodnums(end)+1:end),thiscellstring,' is the wrong size to expand.'])
                    continue
                end
                % Check if this is surface temperature or not:
                if ~strcmp(FlattenVariables{ii},'thermal.spctemperature')
                    % Check if this is on vertices or elements:
                    if thisysize==numvertices2d || thisysize==numvertices2d+1
                        % Expand the variable, vertically uniform:
                        eval(['md.',FlattenVariables{ii}(1:periodnums(end)-1),'(jj).',FlattenVariables{ii}(periodnums(end)+1:end),thiscellstring,...
                            '=[repmat(md.',FlattenVariables{ii}(1:periodnums(end)-1),'(jj).',FlattenVariables{ii}(periodnums(end)+1:end),thiscellstring,'(1:numvertices2d,:),[numlayers,1]);ThisTimeVector];'])
                    else
                        % Expand the variable, vertically uniform:
                        eval(['md.',FlattenVariables{ii}(1:periodnums(end)-1),'(jj).',FlattenVariables{ii}(periodnums(end)+1:end),thiscellstring,...
                            '=[repmat(md.',FlattenVariables{ii}(1:periodnums(end)-1),'(jj).',FlattenVariables{ii}(periodnums(end)+1:end),thiscellstring,'(1:numelements2d,:),[numlayers-1,1]);ThisTimeVector];'])
                    end
                else
                    % Expand the variable, NaNs below the surface:
                    eval(['md.',FlattenVariables{ii}(1:periodnums(end)-1),'(jj).',FlattenVariables{ii}(periodnums(end)+1:end),thiscellstring,...
                        '=[NaN*ones((numlayers-1)*numvertices2d,thisxsize);md.',FlattenVariables{ii}(1:periodnums(end)-1),'(jj).',FlattenVariables{ii}(periodnums(end)+1:end),thiscellstring,'(1:numvertices2d,:);ThisTimeVector];'])
                end
            end
        end
    end
end

    
    
    
    
    
    
    