classdef cvpartitionTallImpl < internal.stats.cvpartitionImpl
% Internal use only.

%   Copyright 2017-2023 The MathWorks, Inc.

    methods
       
        function cv = cvpartitionTallImpl(infoFromTallMethod,method,T,varargin)
            
            LAS = tall(true); % lazy variable to which all deferred validations are attached
            group = infoFromTallMethod{1};
            adaptor = infoFromTallMethod{2};
            partitionfun = infoFromTallMethod{3};
            clientfun = infoFromTallMethod{4};
            slicefun = infoFromTallMethod{5};
            
            hasstratify = false;
            
            if isempty(varargin)
                rs = [];
                stdargin = nargin;
            else
                if length(varargin)==1 % must be a RandStream object
                    stdargin = nargin-1;
                    rs = varargin{1};
                    if ~isa(rs,'RandStream')
                        error(message('stats:cvpartition:cvpartBadArg'));
                    end
                elseif length(varargin)==2 % must be 'Stratify' NVP
                    stra = varargin{1};
                    rs = [];
                    stdargin = nargin;
                    if ~strncmpi(stra,{'stratify'},length(stra))
                        error(message('stats:cvpartition:UnknownThirdInput'));
                    end
                    stratified = varargin{2};
                    hasstratify = true;
                    if ~islogical(stratified)
                        error(message('stats:cvpartition:InvalidStratifyType'));
                    end
                elseif length(varargin)==3 % must be 'Stratify' NVP plus RandStream object
                    stdargin = nargin-1;
                    tmp1 = varargin{1};
                    tmp2 = varargin{2};
                    tmp3 = varargin{3};
                    if strncmpi(tmp1,{'stratify'},length(tmp1)) && isa(tmp3,'RandStream')
                        stratified = varargin{2};
                        rs = tmp3;
                    elseif strncmpi(tmp2,{'stratify'},length(tmp2)) && isa(tmp1,'RandStream')
                        stratified = varargin{3};
                        rs = tmp1;
                    else
                        error(message('stats:cvpartition:UnknownThirdInput'));
                    end
                    hasstratify = true;
                    if ~islogical(stratified)
                        error(message('stats:cvpartition:InvalidStratifyType'));
                    end
                else
                    error(message('stats:cvpartition:cvpartTooManyArg'));
                end
            end
            
            if stdargin < 2
                error(message('stats:cvpartition:TooFewInputs'));
            end
            
            if ischar(method) && size(method,1) == 1
                methodNames = {'kfold','holdout','leaveout','resubstitution'};
                j = find(strncmpi(method,methodNames,length(method)));
                if length(j) > 1
                    error(message('stats:cvpartition:AmbiguousMethod', method));
                elseif isempty(j)
                    error(message('stats:cvpartition:UnknownMethod'));
                end
            else
                error(message('stats:cvpartition:InvalidType'));
            end
            
            if ~strcmp(methodNames{j},'holdout')
                error(message('stats:tall:cvpartition:HoldOutOnly'))
            end
            
            % When stratification is false, we just remove the grouping
            % information from the grouping variable and proceed with the
            % original code.
            if hasstratify && ~stratified
                group = slicefun(@removeGroupingInformation,group);
            end
            
            if ~isnan(adaptor.NDims) && all(~isnan(adaptor.Size))
                % First input is lazy, but the adaptor information can
                % tell us if the user passed N rather than group, we'll try
                % to error out immediately
                if prod(adaptor.Size)==1
                    error(message('stats:tall:cvpartition:ScalarNotAllowed'))
                end
            else
                % If not we need to anyways place a deferred check as
                % passing N is not supported.
                LAS = lazyAssert(~isscalar(group),LAS, @() error(message('stats:tall:cvpartition:ScalarNotAllowed')),clientfun);
            end
            
            cv.Type = methodNames{j};
            cv.Group = group;
            cv.N = size(cv.Group,1); % the number of observations including NaNs
            
            P  = 1/10; % the default holdout ratio
            
            if stdargin == 2 || isempty(T)
                T = P;
            elseif ~isscalar(T) || ~isnumeric(T) || T <= 0 || ~isfinite(T) || T>=1
                error(message('stats:tall:cvpartition:BadP'));
            end
            
            import internal.stats.bigdata.callOptsFromRandStream;
            [cv.indices,pss,pg]= partitionfun(callOptsFromRandStream(rs), ...
                @(info,g) stratifiedSampling(info,g,T), ...
                cv.Group);
            
            LAS = lazyAssert(floor(cv.N * T)>0, LAS, @() error(message('stats:cvpartition:PTooSmall')),clientfun);
            
            spss = sum(pss,1);
            LAS = lazyAssert(spss(1,1)==0,LAS,@() warning(message('stats:cvpartition:MissingGroupsRemoved')),clientfun);
            LAS = lazyAssert((spss(1,3)+spss(1,2))>1,LAS,@() error(message('stats:cvpartition:BadNGrp')),clientfun);

            [ngto,ngtr,ngte] = clientfun(@(pg) consolidateGroups(pg),pg);
            LAS = lazyAssert(ngte==ngto,LAS,@() warning(message('stats:cvpartition:HOTrainingZero')),clientfun);
            LAS = lazyAssert(ngtr==ngto,LAS,@() warning(message('stats:cvpartition:HOTestZero')),clientfun);
            
            cv.holdoutT = T;
            cv.NumTestSets = 1;
            cv.TestSize  = spss(1,3);
            cv.TrainSize = spss(1,2);
            
            % Propagate all lazy validations to all properties
            %cv.indices = lazyAssert(LAS,cv.indices,[],clientfun);
            cv.TestSize = lazyAssert(LAS,cv.TestSize,[],clientfun);
            cv.TrainSize = lazyAssert(LAS,cv.TrainSize,[],clientfun);
            cv.N = lazyAssert(LAS,cv.N,[],clientfun);
            
                        
        end % cvpartition constructor
        
        function disp(cv)
            isLoose = strcmp(get(0,'FormatSpacing'),'loose');
            if (isLoose)
                fprintf('\n');
            end
            disp(getString(message('stats:cvpartition:HoldoutCV')));
            fprintf('   NumObservations: %s\n','[1x1 tall]');
            fprintf('       NumTestSets: %s\n',num2str(cv.NumTestSets));
            fprintf('         TrainSize: %s\n','[1x1 tall]');
            fprintf('          TestSize: %s\n','[1x1 tall]');
            fprintf('          IsCustom: %s\n',num2str(cv.IsCustom));
        end
        
        function cv = repartition(varargin)
            error(message('stats:cvpartition:RepartitionInvalid'))
        end
              
    end

    methods (Access = protected)
        function this = gatherOne(this)
            % Gathers a single cvpartitionTallImpl. Used by
            % cvpartitionTallImpl/gather

            % First call the superclass gatherOne
            this = gatherOne@internal.stats.cvpartitionImpl(this);

            % Then convert to an internal.stats.cvpartitionInMemoryImpl
            warning("off","MATLAB:structOnObject")
            cleanup = onCleanup(@() warning("on","MATLAB:structOnObject"));
            s = struct(this);
            s = rmfield(s,"NumObservations");
            s.BackdoorConstructor = true;
            this = internal.stats.cvpartitionInMemoryImpl(s);
        end
    end
end


function [ngto,ngtr,ngte] = consolidateGroups(pg)
    gte = unique(vertcat(pg{:,2}));
    gtr = unique(vertcat(pg{:,3}));
    gto = unique(vertcat(gte,gtr));
    ngtr = numel(gtr);
    ngte = numel(gte);
    ngto = numel(gto);
end

function [hasFinished,idx,s,g] = stratifiedSampling(info,group,p)

    hasFinished = info.IsLastChunk;
    n = size(group,1);
    if isempty(group)
        gidx=[];  % This happens if we get an empty chunk
    else
        gidx = grp2idx(group(:,1)); % canonicalization of the grouping variable
    end
    gidxnn = gidx(~isnan(gidx));
    if isempty(gidxnn)
        sizeGroups = []; % This happens if we get a chunk with all missing
    else
        sizeGroups = accumarray(gidxnn,1); % tally groups, ignoring NaNs
    end
    
    % Given the size of each group, figure out the size of the training set
    % for each group:
    sizeTest = floor(sizeGroups.*p) + double(rem(sizeGroups.*p,1)>rand(numel(sizeGroups),1)); % number of samples for the test set for each group
    
    % Output idx is a [ChunkSize x 1] vector with 1's for training, 2's for
    % testing, and NaN's for missing
    idx = nan(n,1); % prefill with NaNs
    for i=1:numel(sizeGroups)
        % create a random permutation of all the indices that belong to
        % group gidx==i, permIdxs is a 
        % [Number_of_samples_that_belong_to_group_i  x 2] matrix, the first
        % column has random numbers that are used for sorting and the
        % second column has the linear indices to the original input vector
        permIdxs = sortrows([rand(sizeGroups(i),1),[zeros(0,1);find(gidx==i)]]);
        % Fill idx with 1's and 2's
        idx(permIdxs(1:sizeTest(i),2))=2; % to the test set
        idx(permIdxs(sizeTest(i)+1:end,2))=1; % to the training set
    end
    
    % Output s is a [1x3] vector with a tally of num of nans, training and
    % test set sizes:
    s = sum([isnan(idx) idx==[1 2]],1);
    
    % Output g is a [1x3] cell with unique lists of groups found in this
    % chunk for each of the 3 cases in the Output s
    if ischar(group)
        g = {unique(group(isnan(idx),:),'row'),unique(group(idx==1,:),'row'),unique(group(idx==2,:),'row')};
    else
        g = {unique(group(isnan(idx),:)),unique(group(idx==1,:)),unique(group(idx==2,:))};
    end
end

function y = removeGroupingInformation(x)
% x is any grouping variable, y is all ones of the same height, we only
% preserve "missing" and mark them as NaNs in y.
if isempty(x)
    y = ones(0,1);
else
    if ischar(x)
        x = cellstr(x);
    end
    y = ones(numel(x),1);
    y(ismissing(x)) = nan;
end
end

function tValue = lazyAssert(tSwitch,tValue,errorHandle,clientfun)
%LAZYASSERT Performs a validation in a lazy manner
%
% Syntax:
%  tValue = lazyAssert(tSwitch,tValue,errorHandle)
%
% Inputs:
%  - tSwitch: A tall logical scalar.
%  - tValue: The tall value to which we attach the lazy dependence
%  - errorHandle: Function handle that throws the right error
%
% Outputs:
%  - tValue: Same tValue but with the lazy dependency attached

tValue = clientfun(@(tof,A) ilazyAssert(tof,A,errorHandle),tSwitch,tValue);
    function A = ilazyAssert(tof,A,eh)
        if ~tof
            feval(eh);
        end
    end
end

% LocalWords:  cvpartBadArg NVP cvpartTooManyArg BadNGrp clientfun prefill permIdxs LAZYASSERT
