function [s,idx] = datasample(data,k,varargin)
%DATASAMPLE Randomly sample from data.
%
%   Supported syntaxes:
%
%   [Y,I] = datasample(DATA,K,'Replace',false)
%   [...] = datasample(DATA,K,1,'Replace',false)
%   [...] = datasample(S,...,'Replace',false)
%   
%   Supported optional parameter:
%       'Weights'
%
%   I is a tall logical index.
%   Output is not shuffled.
%   datasample requires a random number generator with multiple stream
%   support, if none is provided datasample defaults to 'mrg32k3a' which in
%   turn is seeded by MATLAB's global random number generator. For
%   reproducibility pass a properly seeded RandStream in the first input.

%   Copyright 2016-2017 The MathWorks, Inc.

if nargin > 2
    [varargin{:}] = convertStringsToChars(varargin{:});
end

nargs = nargin;

% Process the stream argument, if present.
if isnumeric(data) || ~isa(data,'RandStream')
    rs = []; % the use default stream
else
    % shift right to drop RandStream from the argument list
    nargs = nargs - 1;
    rs = data;
    data = k;
    if nargs > 1
        k = varargin{1};
        varargin(1) = [];
    end
end

if nargs < 2
    error(message('stats:datasample:TooFewInputs'));
end

if ~(isscalar(k) && isnumeric(k) && (k==round(k)) && (k >= 0))
    error(message('stats:datasample:InvalidK'));
end

% Pull out the DIM arg if present, or determine which dimension to use.
if nargs > 2 && ~(ischar(varargin{1}) || isstring(varargin{1}))
    dim = varargin{1};
    varargin(1) = [];
else
    dim = 1;
end

% Dimension argument must not be tall and must be a positive integer
if (istall(dim) || ~isnumeric(dim) || ~isscalar(dim) || dim<1 || floor(dim)~=dim)
    error(message('stats:tall:datasample:DimensionMustBePositiveInteger'));
end

if dim~=1
    error(message('stats:tall:datasample:TallDimensionOnly'))
end

% Parse varargin
pnames = {'replace' 'weights' '_forceImpl1' };
dflts =  {true      []        false};
[replace,weights,forceImpl1] = statslib.internal.parseArgs(pnames, dflts, varargin{:});

if replace
    error(message('stats:tall:datasample:NoSampleWithReplacement'))
end

if istall(weights)
    [data,weights] = validateSameTallSize(data,weights);
    [s, idx] = weightedDataSample(data,k,weights,rs);
    return;
else
    if ~isempty(weights)
        validateattributes(weights,{'tall'},{},'tall/datasample','Weights');
    end
end

% Pick which of the 3 implementations to use
implementationLimit_1 = 10000;
implementationLimit_2 = 1000;
if (k>implementationLimit_1) || (nargout>1)
    [s,idx] = datasample_impl_1(data,k,rs);
    idx.Adaptor = matlab.bigdata.internal.adaptors.getAdaptorForType('logical');
    if ~isnan(data.Adaptor.NDims)
        sz = data.Adaptor.Size;
        sz(2) = 1;
        idx.Adaptor = setKnownSize(idx.Adaptor, sz);
    end
elseif (k<=implementationLimit_2)
    s = datasample_impl_2(data,k,rs);
else
    s = datasample_impl_3(data,k,rs,forceImpl1);
end

s.Adaptor = getNewAdaptorAndSetSize(data,k);

end

function newAdaptor = getNewAdaptorAndSetSize(t,k)
% Retrofit adaptor information

% Here we need a copy of the adaptor with a different tall size, so first we
% must reset the tall size to unlink it, and then set it to the known value (if
% non-NaN).
newAdaptor = resetTallSize(t.Adaptor);
if ~isnan(k)
    newAdaptor = setTallSize(newAdaptor, k);
end
end

function [s,idx] = datasample_impl_1(data,k,rs)
% Implementation 1: sample is most likely tall. Worst case scenario is two
% passes. In the first pass the size of each chunk and the number of
% chunks are collected, then the client figures out how many samples are
% required in each chunk replicating a uniform sampling over the whole
% array. The second pass uniformly samples the data in each chunk. For
% this approach it is required is that chunks are consistent between
% different data passes.
import internal.stats.bigdata.callOptsFromRandStream;
MaxNumSlices = 20000; % on-a-par with ReadSize for now

[chunkIDs,chunkSizes] = fixedchunkfun( @(info,x)  chunkInfoFun(info,x), MaxNumSlices, {zeros(0,2),zeros(0,1)}, data);
nPerChunk  = clientfun(callOptsFromRandStream(rs), ...
    @(siz) nPerChunkFun(siz,k), chunkSizes);

nPerChunk_ = matlab.bigdata.internal.broadcast(nPerChunk);
chunkIDs_  = matlab.bigdata.internal.broadcast(chunkIDs);

idx = fixedchunkfun(callOptsFromRandStream(rs), ...
    @exactRandSampleFun, ...
    MaxNumSlices, zeros(0,1,'logical'), data, chunkIDs_, nPerChunk_);

% n-d subsref
s = chunkfun(@(x,idx) ndSubsrefInDim1(x,idx),data,idx);
s.Adaptor = getNewAdaptorAndSetSize(data,k);
end

function s = datasample_impl_2(data,k,rs)
% Implementation 2: very small samples, k is very small (k*numChunks fits
% in the client memory) and we only need the samples -s- (but not the
% indices -idx-). Worst case scenario can be reduced to one pass on the
% data. We sample each chunk using k. Samples and actual chunk sizes are
% returned to the client and the client guarantees k samples uniformly
% sampled from the original data.
import internal.stats.bigdata.callOptsFromRandStream;
[S,nS,chunkSizes] = partitionfun( callOptsFromRandStream(rs), ...
    @(info,x) constantRandSampleFun(info,x,k), data);
S.Adaptor = getNewAdaptorAndSetSize(data,NaN);

s = clientfun( callOptsFromRandStream(rs), ...
    @(S,nS,siz) samplekFromSampleSets(S,nS,siz,k), ...
    S, nS, chunkSizes);
s.Adaptor = getNewAdaptorAndSetSize(data,k);
end

function s = datasample_impl_3(data,k,rs,forceImpl1)
% Implementation 3: small samples, k is small (sample fits in the client
% memory) and we only need the samples -s- (but not the indices -idx-).
% If size(data,1) is available worst case scenario can be one pass,
% otherwise and extra pass to get N is needed.
import internal.stats.bigdata.callOptsFromRandStream;

N = size(data,1);
[~, ub] = sampleSizeBounds(k./N, N);
ub_  = matlab.bigdata.internal.broadcast(ub);

[S,nS] = partitionfun( callOptsFromRandStream(rs), ...
    @rateRandSampleFun, data, ub_ );
S.Adaptor = getNewAdaptorAndSetSize(data,NaN);

[s,tnS] = clientfun( callOptsFromRandStream(rs), ...
    @(S,nS,siz) samplekFromSampleSets(S,nS,siz,k), ...
    S, nS, N);
s.Adaptor = getNewAdaptorAndSetSize(data,k);
s_is_good = tnS==k;
if forceImpl1
   s_is_good = tall(false);
end
s_safe = datasample_impl_1(data,k,rs);
s = ternary(s_is_good ,s,{s_safe}); % { } removes dependency between s and  
                                    % s_safe in the evaluation tree so
                                    % passes over the data to calculate s
                                    % do not trigger the passes to
                                    % calculate s_safe, there is a specific
                                    % gather inside ternary() that will
                                    % trigger that branch in case is needed
                                    
end

function s = ndSubsrefInDim1(x,idx)
   ss = substruct('()',[{idx} repmat({':'},1,ndims(x)-1)]);
   s = subsref(x,ss);
end

function tValue = ternary(tSwitch,tValueIfTrue,tValueIfFalse)
%TERNARY Select between two tall values depending on a tall logical scalar value.
%
% Syntax:
%  tValue = ternary(tSwitch,tValueIfTrue,tValueIfFalse)
%
% Inputs:
%  - tSwitch: A tall logical scalar.
%  - tValueIfTrue: The tall value to return if tSwitch is true.
%  - tValueIfFalse: The tall value to return if tSwitch is false.
%
% Outputs:
%  - tValue: Either tValueIfTrue or tValueIfFalse depending on tSwitch.

tValue = clientfun(@(tof,oA,oB) iTernary(tof,oA,oB),tSwitch,tValueIfTrue,tValueIfFalse);
    function out = iTernary(tof,oA,oB)
        if tof
            out = oA;
        else
            out = oB;
        end
        % Inputs wrapped in a cell need to be triggered here
        if iscell(out) && numel(out)>0 && istall(out{1})
            out = gather(out{1});
        end
    end
end

function  varargout = fixedchunkfun(opts,fcn,FixedNumSlices,OutputEmptyTemplate,varargin)
%TODO (g1555526): Use tall/fixedchunkfun now that it is available.

% Utility function that uses internal.stats.bigdata.FixedChunkingFunction
% to create a function handle which in turn will be used by partitionfun
% but where we enforce fixed number of rows in every chunk within a
% partition except for the last chunk.

% Strip out opts and fcn
[opts,fcn,FixedNumSlices,OutputEmptyTemplate,varargin] = ...
    matlab.bigdata.internal.util.stripOptions(opts,fcn,FixedNumSlices,OutputEmptyTemplate,varargin{:});

    bi = cellfun(@(x) isa(x,'matlab.bigdata.internal.BroadcastArray'),varargin);
    fixedChunk = internal.stats.bigdata.FixedChunkingFunction(fcn,FixedNumSlices,bi,OutputEmptyTemplate);
    fcnHandle  = matlab.bigdata.internal.FunctionHandle(fixedChunk,'MaxNumSlices',fixedChunk.FixedChunkSize);
    [varargout{1:nargout}] = partitionfun(opts,fcnHandle,varargin{:});
end

function [s,tnS] = samplekFromSampleSets(S,nS,siz,k)
if numel(siz)==1
    nS = sum(nS);
end

NPC =  min(nS,nPerChunkFun(siz,k)); % use min to prevent the very unlikely case than we want to get
                                    % more samples than what we have available in a subset
csnS = [0;cumsum(nS)];
idx = false(csnS(end),1);
for i = 1:numel(NPC)
    idx(randsample(nS(i),NPC(i)) + csnS(i)) = true;
end
s = ndSubsrefInDim1(S,idx);
tnS = sum(idx);
end

function NPC = nPerChunkFun(siz,n)
% Find number of samples required in every chunk (NPC). NPC will be the
% same in every chunk because all the workers setup the global stream as brs
nChunks = size(siz,1);
N = sum(siz);
if n>N
    error(message('stats:datasample:SampleTooLarge'))
end
NPC = zeros(nChunks,1);
SNPC = 0;

while SNPC<n
    % By using Meng et.al. upper bound instead of just n/N,
    % the body loop is executed once 99.99% of the times.
    [~,ub] = sampleSizeBounds((n-SNPC)./N, N);
    NPC = NPC + binornd(siz-NPC,min(1,ub));
    SNPC = sum(NPC);
end
while SNPC>n
    NPC = NPC - min(NPC,mnrnd(SNPC-n,NPC./SNPC)');
    SNPC = sum(NPC);
end
end

function [hasFinished,id,siz] = chunkInfoFun(info,x)
hasFinished = info.IsLastChunk;
id = [info.PartitionId info.FixedSizeChunkID];
siz = size(x,1);
end

function [hasFinished,S,nS,siz] = rateRandSampleFun(info,x,p)
hasFinished = info.IsLastChunk;
siz = size(x,1);
% Create a logical index to subsample the data in each chunk.
idx = rand(siz,1)<=p;
S = ndSubsrefInDim1(x,idx);
nS = sum(idx);
end

function [hasFinished,S,nS,siz] = constantRandSampleFun(info,x,k)
hasFinished = info.IsLastChunk;
siz = size(x,1);
if siz>k
    % Create a logical index to subsample the data in each chunk.
    idx = false(siz,1);
    idx(randsample(siz,k)) = true;
    S = ndSubsrefInDim1(x,idx);
    nS = k;
else
    S = x;
    nS = siz;
end

end

function [hasFinished,idx] = exactRandSampleFun(info,x,cids,NPC)
hasFinished = info.IsLastChunk;
% Identify the chunk in which we currently are...
chunkIdx = knnsearch(cids,[info.PartitionId info.FixedSizeChunkID]);
% Number to samples to draw in this chunk
np = NPC(chunkIdx);
idx = false(size(x,1),1);
idx(randsample(size(x,1),np)) = true;
end

function [lower, upper] = sampleSizeBounds(p, n)
%sampleSizeBounds computes the acceptance and reject sampling thresholds
%  [LOWER, UPPER] = sampleSizeBounds(P, N) computes the accepting and
%  rejecting threshold for uniform random sampling without replacement of a
%  population of size N.

% Reference:
%
% [1] Meng, Xiangrui, Scalable Simple Random Sampling and Stratified
%     Sampling, Proceedings of the 30th International Conference on Machine
%     Learning (2013) Volume: 28, Issue: 1, Pages: pp. 1-9

% Set the failing rate (not having enough samples)
SIGMA = 1e-4;

% A minimal sampling rate for the wait list
minUpper = 1e-10;

gamma1 = -log(SIGMA)./n;
gamma2 = gamma1.*(2/3);

upper = p + gamma1 + sqrt(gamma1.^2 + 2.*gamma1.*p);
lower = p + gamma2 - sqrt(gamma2.^2 + 3.*gamma2.*p);

upper = max(0,upper);
lower = min(1, max(minUpper, lower));
end
