function [targetcols,mergedColTerms,mergedLevels] = enlargeVIArrays(val1,mergedVariableInfo,terms)
%enlargeArray Enlarge arrays to match the union of all categorical levels.

% Inputs:
%   val1 = original VariableInfo table
%   mergedVariableInfo = merged VariableInfo table
%   terms = terms matrix
%
% Given an already-merged VariableInfo table, find the following for the
% merged data:
%   targetcols - columns in the merged array to contain X for this chunk
%   mergedColTerms - merged columns-to-terms vector
%   mergedLevels - all combinations of categorical levels

%   Copyright 2016-2017 The MathWorks, Inc.

oldVariableInfo = val1.VariableInfo;
oldColTerms = val1.Cols2Terms;
iscat = oldVariableInfo.IsCategorical;

% Update information about the variables
[numMergedCats,numOldCats,newindices] = ...
    updateVariables(oldVariableInfo,mergedVariableInfo);

% Loop over terms once to find the total number of design matrix columns
ncols = 0;
nterms = size(terms,1);
for j=1:nterms
    idx = (terms(j,:)~=0)';
    ncols = ncols + prod(numMergedCats(idx));
end

% Loop over terms to find out what goes in each column
mergedColTerms = zeros(1,ncols);
mergedLevels = cell(1,nterms);
targetcols = zeros(1,size(oldColTerms,2));
base = 0;
for j=1:nterms
    oldColumns = oldColTerms==j;
    idx = (terms(j,:)~=0)' & iscat;
    isCatTerm = any(idx);
    if isCatTerm
        [termcols,numtermcols,mergedLevels{j}] = ...
            findTargetCols(numMergedCats(idx),numOldCats(idx),newindices(idx));
    else
        termcols = 1;
        numtermcols = 1;
    end
    targetcols(oldColumns) = base + termcols;
    mergedColTerms(1,base+(1:numtermcols)) = j;
    base = base + numtermcols;
end
end

%% Update single variable info to match merged info
function [numMergedCats,numOldCats,newindices] = updateVariables(oldVariableInfo,mergedVariableInfo)
% Get information needed to update each variable during merge
nvars = size(mergedVariableInfo,1);
iscat = oldVariableInfo.IsCategorical;
oldRange = oldVariableInfo.Range;
numMergedCats = ones(nvars,1);
numOldCats = ones(nvars,1);
newindices = cell(nvars,1);
for j=1:nvars
    if iscat(j)
        % Count the number of categories of each variable. Convert old
        % levels to indices into new levels.
        mergedVarLevels = mergedVariableInfo.Range{j};
        if iscategorical(mergedVarLevels) && isordinal(mergedVarLevels)
            mergedVarLevels = categories(mergedVarLevels);
        end
        oldLevels = oldRange{j};
        
        if ischar(mergedVarLevels)
            nlevels = size(mergedVarLevels,1);
            [~,indices] = ismember(oldLevels,mergedVarLevels,'rows');
        else
            nlevels = length(mergedVarLevels);
            [~,indices] = ismember(oldLevels,mergedVarLevels);
        end
        numMergedCats(j) = nlevels;
        numOldCats(j) = length(indices);
        newindices{j} = indices(:);
    end
end
end

%% Find target columns for merging
function [colidx,ncols,mergedDesign] = findTargetCols(newcats,oldcats,catlevels)
% Find location for old categorical columns in merged design.
mergedDesign = fullfact(newcats);
if oldcats>0
    oldDesign = fullfact(oldcats);
else
    oldDesign = mergedDesign(1:0,:); % take zero rows
end
for j=1:length(catlevels)
    levj = catlevels{j};
    oldDesign(:,j) = levj(oldDesign(:,j));
end

[~,colidx] = ismember(oldDesign,mergedDesign,'rows');
ncols = prod(newcats);
end



