function [G,GR] = grouptransform(T,groupVars_,method_,varargin)
%MATLAB Code Generation Library Function

%   Limitations:
%       * All inputs except T must be constant
%       * ND arrays in tables are not supported
%       * GR, the second output, is always a cell array unless groupVars
%         are constant and there are 1 or fewer groups
%   Copyright 2022-2023 The MathWorks, Inc.
%#codegen

coder.internal.prefer_const(groupVars_, method_, varargin);
ONE = coder.internal.indexInt(1);

narginchk(3,Inf);
 
% Compute matrix/table switch flag
tableFlag = isa(T,'tabular');

if tableFlag
    % Error if asking for more than 1 output for table
    nargoutchk(0,1);
end

% Parse grouping variables
[groupingDataParsed,groupVarsParsed] = parseGroupVars(groupVars_,tableFlag,'grouptransform',T);

if tableFlag
    % Create labels for groupVars in output table
    gvLabels_ = groupVarsParsed;

    % Set default values for data variables
    varNamesT = T.Properties.VariableNames;
    defaultDataVars = varNamesT;
else
    gvLabels_ = numel(groupingDataParsed);
    % in-memory version uses cellstrs for gvLabels - those can't be
    % generated at runtime in generated code
    defaultDataVars = T;
    varNamesT = [];
end

if tableFlag
    numLabels = numel(gvLabels_);
else
    numLabels = gvLabels_;
end

[groupBins, gbProvided, scalarExpandBins, scalarExpandVars, method, ...
    methodPrefix, dataVars, dvNotProvided, inclEdge, replaceDataVars] = ...
    parseOptions(varargin, tableFlag, T, varNamesT, defaultDataVars, numLabels, method_);

if gbProvided
    % Discretize grouping variables and remove repeated pairs of grouping
    % variables and group bins
    [groupingData,groupVars,gvLabels,gbForGV] = discgroupvar(groupingDataParsed,...
        groupVarsParsed,gvLabels_,groupBins,inclEdge,scalarExpandBins,...
        scalarExpandVars,'grouptransform',tableFlag);
else
    gbForGV = false(size(groupVarsParsed));
    if tableFlag
        % Remove repeated grouping variables
        [groupVars,ridx] = coder.const(@feval,'unique',groupVarsParsed,'stable');
        groupingData = {groupingDataParsed{ridx}};
        gvLabels = {gvLabels_{ridx}};
    else
        groupVars = groupVarsParsed;
        groupingData = groupingDataParsed;
        gvLabels = gvLabels_;
    end
end

% Compute grouping index and data
gvIdx = mgrp2idx(groupingData,size(T,1),true,false);

% Extract data variables
[dvData,dataLabels,numDataVars] = extractDataVars(T,groupVars,dataVars,tableFlag,dvNotProvided);

% Compute group tranformations
gStats = cell(1,numDataVars);

% gstatsLabel is only used for the table case
if tableFlag
    % Force gStatsLabel to be compile-time constant
    gStatsLabel = cell(1,numDataVars);
    coder.unroll()
    for ii = ONE:numDataVars
        gStatsLabel{ii} = coder.const([methodPrefix '_' dataLabels{ii}]);
    end
    gStatsLabel = coder.const(gStatsLabel);
else
    gStatsLabel = [];
end

[gvIdxSorted, gvIdxSortedIdx] = sort(gvIdx);
grpStart = find([1;diff(gvIdxSorted)>0]);

% Avoid vector(scalar) vs. vector(vector) indexing issues
if isscalar(grpStart)
    grpEnd = length(gvIdxSorted);
else
    grpEnd = [grpStart(2:end)-1;length(gvIdxSorted)];
end

for ii = ONE:numDataVars
    gStats{ii} = applyFunWrapper(grpStart,grpEnd,gvIdxSortedIdx, ...
        dvData{ii},method,tableFlag,dataLabels{ii},methodPrefix);
end

% Prepare discretized groupingData
if (gbProvided && tableFlag) || nargout > 1
    if tableFlag
        tmp = coder.const(gbForGV);
        if coder.const(feval('any',gbForGV))
            gvLabelsFinal = {gvLabels{tmp}};
            groupingDataFinal = {groupingData{tmp}};
        else
            % Avoid trying to index with an empty variable
            gvLabelsFinal = {};
            groupingDataFinal = {};
        end
    else
        gvLabelsFinal = gvLabels;
        % Set groups
        numGroupVars = gvLabelsFinal;
        % GR will be a cell array unless the number of groups is
        % constantly 1
        if coder.internal.isConstTrue(numGroupVars == 1)
            groupingDataFinal = groupingData{1};
        else
            groupingDataFinal = groupingData;
        end
        if numGroupVars == 0
            % groupVars must be 0x0 or mx0 here because for 0xn numGroupVars is n
            GR_ = constructEmpty(coder.internal.scalarEg(groupVars),size(T,1),0);
            if issparse(groupVars)
                GR = {sparse(GR_)};
            else
                GR = {GR_};
            end
        else
            GR = groupingDataFinal;
        end
    end
end

% Package outputs
% namelengthmax is not supported for codegen
NAMELENGTHMAX = coder.const(feval('namelengthmax'));
if replaceDataVars % 'ReplaceValues' true
    if tableFlag
        DIMNAMES = coder.const(T.Properties.DimensionNames);
        if gbProvided
            data = cell(1, width(T) + numel(groupingDataFinal));
            VARNAMES = coder.const(feval('cat',2,varNamesT,gvLabelsFinal));
            uniqueLabels = coder.const(feval('matlab.lang.makeUniqueStrings',VARNAMES,DIMNAMES,NAMELENGTHMAX));
        else
            data = cell(1, width(T));
            uniqueLabels = varNamesT;
        end

        [~, dataVarsIdx] = coder.const(@feval,'ismember',dataLabels,varNamesT);
        dataVarsCount = ONE;
        coder.unroll()
        for ii = ONE:width(T)
            if dataVarsCount <= numel(dataVarsIdx) && ii == dataVarsIdx(dataVarsCount)
                % Replace variable with computed data
                data{ii} = gStats{:,dataVarsCount};
                dataVarsCount = dataVarsCount + ONE;
            else
                % Variable is not replaced
                data{ii} = T{:,ii};
            end
        end
        if gbProvided
            for ii =  ONE:numel(groupingDataFinal)
                data{ii + width(T)} = groupingDataFinal{ii};
            end
        end
        G = initTableOrTimetable(T,data,height(T),numel(data),uniqueLabels,DIMNAMES);
    else
        % Try to extract G from Data
        G = [gStats{:}];
    end
else % 'ReplaceValues' false
    if tableFlag
        DIMNAMES = coder.const(T.Properties.DimensionNames);
        if gbProvided
            % Make sure all labels are unique
            VARNAMES = coder.const(feval('cat',2,varNamesT,gvLabelsFinal,gStatsLabel));
            uniqueLabels = coder.const(feval('matlab.lang.makeUniqueStrings',VARNAMES,DIMNAMES,NAMELENGTHMAX));
            % groupingDataFinal is always a cell array when the input is a
            % table
            numelGroupingDataFinal = numel(groupingDataFinal);
            numelGStats = numel(gStats);
            data = cell(1,numelGroupingDataFinal+numelGStats);
            [data{ONE:numelGroupingDataFinal}] = deal(groupingDataFinal{:});
            [data{ONE+numelGroupingDataFinal:end}] = deal(gStats{:});
            GTmp = initTableOrTimetable(T,data,height(T),numel(data),{uniqueLabels{numel(varNamesT)+1:end}},DIMNAMES);
        else
            % Make sure all labels are unique
            VARNAMES = coder.const(feval('cat',2,varNamesT,gStatsLabel));
            uniqueLabels = coder.const(feval('matlab.lang.makeUniqueStrings',VARNAMES,DIMNAMES,NAMELENGTHMAX));
            GTmp = initTableOrTimetable(T,gStats,height(T),numel(gStats),{uniqueLabels{numel(varNamesT)+1:end}},DIMNAMES);
        end
        if isa(T,'timetable') && isempty(T)
            G = T;
        else
            G = [T,GTmp];
        end
    else
        % Try to extract G from Data
        G = [T,gStats{:}];
    end
end

end

%--------------------------------------------------------------------------

function [groupBins, gbProvided, scalarExpandBins, scalarExpandVars, method, ...
    methodPrefix, dataVars, dvNotProvided, inclEdge, ReplaceValues] = ...
    parseOptions(inputs, tableFlag, T, varNamesT, defaultDataVars, numLabels, method_)

coder.internal.prefer_const(inputs, tableFlag, varNamesT, defaultDataVars, numLabels, method_);
ONE = coder.internal.indexInt(1);

if isgroupbins(method_,'grouptransform')
    [groupBins,scalarExpandBins,scalarExpandVars] = parsegroupbins(method_,numLabels,'grouptransform');
    gbProvided = true;
    % Next input argument (first element of inputs) must be method
    coder.internal.errorIf(isempty(inputs), 'MATLAB:grouptransform:MethodNotProvided')
    [method,methodPrefix] = parseMethod(inputs{ONE},tableFlag);
    indStart = coder.internal.indexInt(2);
else
    [method,methodPrefix] = parseMethod(method_,tableFlag);
    groupBins = {};
    scalarExpandBins = false;
    scalarExpandVars = false;
    gbProvided = false;
    indStart = ONE;
end

% Parse datavars
if tableFlag && indStart<= numel(inputs) && isDataVars(inputs{indStart},numel(inputs),indStart)
    dataVars = parseDataVars(T,inputs{indStart},varNamesT);
    dvNotProvided = false;
    indStart_ = indStart + ONE;
else
    dataVars = defaultDataVars;
    dvNotProvided = true;
    indStart_ = indStart;
end

if indStart_ <= numel(inputs)
    % Parse name-value pairs
    poptions = struct( ...
             'CaseSensitivity',false, ...
             'PartialMatching','unique', ...
             'StructExpand',false, ...
             'IgnoreNulls',false, ...
             'SupportOverrides',true);
    
    pstruct = coder.internal.parseInputs({},{'IncludedEdge','ReplaceValues'},poptions,inputs{indStart_:end});
    
    coder.internal.assert(pstruct.IncludedEdge == 0 || gbProvided, 'MATLAB:grouptransform:IncludedEdgeNoGroupBins');
    inclEdge = coder.internal.getParameterValue(pstruct.IncludedEdge,'left',inputs{indStart_:end});
    % inclEdge value will be parsed by discretize
    
    ReplaceValues_ = coder.internal.getParameterValue(pstruct.ReplaceValues,true,inputs{indStart_:end});
    coder.internal.assert(isscalar(ReplaceValues_) && (islogical(ReplaceValues_) || isnumeric(ReplaceValues_)), ...
                'MATLAB:table:InvalidLogicalVal', 'ReplaceValues');
    ReplaceValues = logical(ReplaceValues_);
else
    inclEdge = 'left';
    ReplaceValues = true;
end
end

%--------------------------------------------------------------------------

function [method,methodPrefix] = parseMethod(method_,tableFlag)
coder.internal.prefer_const(method_,tableFlag);
coder.internal.assert(coder.internal.isConst(method_),'MATLAB:grouptransform:CodegenConstMethod');
if ischar(method_)
    coder.internal.assert(isrow(method_), 'MATLAB:grouptransform:InvalidMethodOption');
    methodCharOrFunctionHandle = method_;
else
    coder.internal.assert(isscalar(method_), 'MATLAB:grouptransform:InvalidMethodOption');
    if isstring(method_)
        methodCharOrFunctionHandle = char(method_);
    else
        methodCharOrFunctionHandle = method_;
    end
end

% Change each option to a function handle and set names appropriately
if ischar(methodCharOrFunctionHandle)
    if strncmpi(methodCharOrFunctionHandle,'zscore',1)
        method = @(x) normalizeArray(x,'zscore','std',0,1,[],[],[],tableFlag);
        methodPrefix = 'zscore';
    elseif strncmpi(methodCharOrFunctionHandle,'norm',3)
        method = @(x) normalizeArray(x,'norm','value',2,1,[],[],[],tableFlag);
        methodPrefix = 'norm';
    elseif strncmpi(methodCharOrFunctionHandle,'meancenter',5)
        method = @(x) normalizeArray(x,'center','mean',0,1,[],[],[],tableFlag);
        methodPrefix = 'meancenter';
    elseif strncmpi(methodCharOrFunctionHandle,'rescale',1)
        method = @(x) normalizeArray(x,'range','value',[0 1],1,[],[],[],tableFlag);
        methodPrefix = 'rescale';
    elseif strncmpi(methodCharOrFunctionHandle,'meanfill',5)
        method = @(x) fillmissing(x,'constant',mean(x,1,'omitnan'),1);
        methodPrefix = 'meanfill';
    elseif strncmpi(methodCharOrFunctionHandle,'linearfill',1)
        method = @(x) fillmissing(x,'linear',1);
        methodPrefix = 'linearfill';
    else
        coder.internal.assert(false,'MATLAB:grouptransform:InvalidMethodOption');
    end
else
    if ~isa(methodCharOrFunctionHandle,'function_handle')
        error(message('MATLAB:grouptransform:InvalidMethodOption'));
    end
    method = methodCharOrFunctionHandle;
    methodPrefix = 'fun';
end
end

%--------------------------------------------------------------------------

function tf = isDataVars(in,numInputs,ind)
coder.internal.prefer_const(in,numInputs,ind);
% Returns true if in is a char, cellstr, string scalar, numeric vector,
% or logical vector. Errors for funtion handle.
% Check chars and string scalars for valid names
coder.internal.errorIf(isa(in,'function_handle'),'MATLAB:grouptransform:CodegenFuncHandleDataVars');
if ischar(in) || isstring(in)
    % Check that in isn't the start of a N-V pair
    % in should have already been checked for const-ness
    if rem(numInputs - ind,2) == 1
        % Only check for possible parameter names if there are an even
        % number of inputs left
        coder.internal.assert(coder.internal.isConst(in),'MATLAB:grouptransform:CodegenConstDataVars');
        matchesName = coder.const(feval('strncmpi',in,{'IncludedEdge','ReplaceValues'},12));
        tf = ~coder.const(feval('any',matchesName));
    else
        tf = true;
    end
elseif iscellstr(in) || islogical(in) || isnumeric(in)
    tf = true;
    return;
else
    tf = false;
end
end

%--------------------------------------------------------------------------

function dataVars = parseDataVars(T,in,varNamesT)
    % The only allowable formats for dataVars are char arrays, string
    % scalars, cellstrs, logical vectors, and numeric vectors
    coder.internal.prefer_const(in,varNamesT);
    coder.internal.assert(coder.internal.isConst(in),'MATLAB:grouptransform:CodegenConstDataVars');
    dataVarsIdx = checkDataVariables(T, in, 'grouptransform', 'Data');
    dataVars = coder.const(feval('unique',{varNamesT{dataVarsIdx}},'stable'));
end

%--------------------------------------------------------------------------

function a = applyFunByGroup(grpStart,grpEnd,sortOrd,x,fun,tableFlag,dataLabels)
% APPLYFUNBYGROUP Apply function to data using sorted groups

coder.internal.prefer_const(fun,tableFlag,dataLabels);

% Apply the function to the first group to get the type/size of the output
% to preallocate

grpSz = grpEnd(1) - grpStart(1) + 1;
gvals = x(sortOrd(grpStart(1):grpEnd(1)),:);
d1 = fun(gvals);
szD1 = size(d1,1);
szD2 = size(d1,2);

coder.internal.errorIf(iscell(d1),'MATLAB:grouptransform:CodegenMethodCantReturnCell');
coder.internal.errorIf(grpSz ~= szD1 && ~(isscalar(d1) || (isempty(d1) && szD1 == 1) || (tableFlag && szD1 == 1)), ...
    'MATLAB:grouptransform:MethodOutputInvalidSize');
coder.internal.errorIf(~tableFlag && (length(size(d1)) ~= 2 || ~(size(d1,2) == 0 || size(d1,2) == 1)), ...
    'MATLAB:grouptransform:MethodOutputInvalidSize');

% Allocate output - use d to figure out type, but since d could have height
% of the number of group elements, just use a sample to get the right size
if ~isempty(d1)
    a = repmat(d1(1),[size(x,1),szD2]);
else
    a = repmat(d1(1,:),size(x,1),1);
end

% Fill output with result of function applied to first group
if (tableFlag && szD1 == 1) || (~tableFlag && issparse(d1) && isempty(d1) && szD1 == 1)
    d1_ = repmat(d1,grpSz,1);
else
    d1_ = d1;
end

a(sortOrd(grpStart(1):grpEnd(1)),:) = d1_;

% Apply function to remaining groups
for g = 2:length(grpStart)
    grpSz = grpEnd(g) - grpStart(g) + 1;
    gvals = x(sortOrd(grpStart(g):grpEnd(g)),:);
    dn = fun(gvals);
    szD1 = size(dn,1);

    coder.internal.errorIf(grpSz ~= szD1 && ~(isscalar(dn) || ...
        (isempty(dn) && szD1 == 1) || (tableFlag && szD1 == 1)), ...
        'MATLAB:grouptransform:MethodOutputInvalidSize');

    if tableFlag
        coder.internal.errorIf(~isequal(size(dn,2:ndims(dn)),szD2), 'MATLAB:grouptransform:ApplyDataVarsError',dataLabels);
    else
        coder.internal.errorIf(~isequal(size(dn,2:ndims(dn)),szD2), 'MATLAB:grouptransform:ApplyDataVecsError',dataLabels);
    end

    if szD1 == 1 && (tableFlag || (~tableFlag && issparse(dn) && isempty(dn)))
        dn_ = repmat(dn,grpSz,1);
    else
        dn_ = dn;
    end

    rowInd = sortOrd(grpStart(g):grpEnd(g));
    % This works around a scalar expansion limition for Coder
    if isscalar(dn_)
        % When a and dn_ are both varsize, scalar expansion only works when
        % the rhs is guaranteed to be scalar
        a(rowInd,:) = dn_(1);
    else
        a(rowInd,:) = dn_;
    end
end
end

%--------------------------------------------------------------------------

function d = applyFunWrapper(grpStart,grpEnd,gvIdxSortedIdx,x,method,tableFlag,dataLabel,methodPrefix)
% Error messages here will differ from the in-memory version
if isempty(gvIdxSortedIdx)
    % Need to special case mean fill of an empty since mean of an
    % empty will return NaN
    d_ = method(x);
    coder.internal.errorIf(iscell(d_),'MATLAB:grouptransform:CodegenMethodCantReturnCell');
    if strcmp(methodPrefix,'meanfill') && isempty(x)
        % This is equivalent to returning x in this case, but it avoids a
        % possible size mis-match
        d_ = repmat(d_,size(x,1),size(x,2));
    end
    if ~tableFlag
        coder.internal.errorIf(length(size(d_)) ~= 2 || ~(size(d_,2) == 0 || size(d_,2) == 1), ...
            'MATLAB:grouptransform:MethodOutputInvalidSize');
    end
    coder.internal.errorIf((size(d_,1) ~= size(x,1)) && ~(isscalar(d_) && size(x,1) == 0) ...
        && size(d_,1) ~= 0, 'MATLAB:grouptransform:MethodOutputInvalidSize');
    if size(d_,1) ~= size(x,1)
        % Make sure that the output is always the same height as
        % the input when, for example, d is a scalar we need to
        % shrink it to an empty
        d = repmat(d_,size(x,1),1);
    else
        d = d_;
    end
else
    % Otherwise apply function by group
    d = applyFunByGroup(grpStart,grpEnd,gvIdxSortedIdx,x,method,tableFlag,dataLabel);
end
end

%--------------------------------------------------------------------------

function G = initTableOrTimetable(T, vars, numRows, numVars, varNames, dimNames)
coder.inline('always');
coder.internal.prefer_const(numRows, numVars, varNames, dimNames);
if istimetable(T)
    G = timetable.init(vars, numRows, T.Properties.RowTimes, numVars, varNames, dimNames);
else
    G = table.init(vars, numRows, T.Properties.RowNames, numVars, varNames, dimNames);
end
end