%DATALOADER Create a dataloader for iterating through a Datastore and
%   generating mini-batches of data.
%
%   LOADER = DATALOADER(DS) creates a dataloader, LOADER, given an input
%   Datastore.
%
%   LOADER = DATALOADER(__,'Name1',Value1,'Name2',Value2, ...) specifies the
%   properties of LOADER using optional name-value pairs. Parameters
%   include:
%
%   'MiniBatchSize' - MiniBatchSize that will be used when forming tensors
%                     for each variable present in input Datastore DS.
%
%                     Default: 128
%
%   'CollateFcn'     - A function that can be used to define how tensors are
%                     formed by the DataLoader. The CollateFcn signature is:
%
%                       function outputBatch = CollateFcn(inputCellData)
%
%                     where inputCellData is an MxN cell array where M is
%                     the MiniBatchSize and N is the number of variables
%                     returned by the input Datastore (the number of
%                     columns in the output table/cell array returned  by
%                     read). outputBatch is a 1xN cell  array in which each
%                     each column of inputCellData has been converted into
%                     a tensor.

%   Copyright 2018-2023 The MathWorks, Inc.

classdef DataLoader < handle & matlab.mixin.Copyable
    
    properties (SetAccess = private)
        % Datastore  Datastore this DataLoader will read
        Datastore

        % CollateFcn  Function used to collate observations in order to
        % construct a mini-batch
        CollateFcn
    end

    properties (Dependent)
        % MiniBatchSize  Size of the mini-batch this DataLoader will return
        MiniBatchSize

        % ReturnIncompleteMiniBatch  Indicate what to do with incomplete
        % mini batches
        ReturnIncompleteMiniBatch
    end

    properties (Access = private)
        % MiniBatchReader Strategy used to read the mini-batch
        MiniBatchReader
    end

    methods
        function this = DataLoader(inputDatastore,options)
            arguments
                inputDatastore
                options.CollateFcn (1,1) function_handle = @defaultCollateFcn
                options.MiniBatchSize (1,1) {mustBeInteger, mustBeNonnegative} = 128
                options.ReturnIncompleteMiniBatch = true
            end

            % We'll use a copy of the input datastore
            ds = copy(inputDatastore);

            % Create the reader. It will reset the datastore
            this.MiniBatchReader = iCreateMiniBatchReader(ds, options.MiniBatchSize, options.ReturnIncompleteMiniBatch);

            % Save other properties
            this.CollateFcn = options.CollateFcn;            
            this.Datastore = ds;
        end

        function [miniBatch, miniBatchSize] = nextBatch(this)
            % Return the next mini-batch of data
    
            % Call the strategy to collect a mini-batch from the datastore
            [this.MiniBatchReader, miniBatch, miniBatchSize] = readMiniBatch(this.MiniBatchReader, this.Datastore);
            % If no data, don't apply the collate function. We won't try to
            % use this mini-batch
            if miniBatchSize>0
                miniBatch = this.CollateFcn(miniBatch);
            end
        end
        
        function data = readByIndex(this,indices)
            % readByIndex  Read the given indices 
            % Subset the original dataLoader with the given indices
            subloader = subset(this,indices);
            % Setting the MiniBatchSize will reset the object. Read all the
            % data
            subloader.MiniBatchSize = numel(indices);
            data = nextBatch(subloader);
        end

        function clonedLoader = clone(this)
            % clone  Return a resetted copy of itself
            clonedLoader = copy(this);
            reset(clonedLoader);
        end
        
        function shuffledLoader = shuffle(this)
            % shuffle  Shuffle the underlying Datastore and return a new
            % DataLoader object
            shuffledLoader = copy(this);
            ds = shuffledLoader.Datastore;
            shuffledLoader.Datastore = shuffle(ds);
            reset(shuffledLoader);
        end
        
        function newLoader = partition(this,numPartitions,index)
            % partition  Partition the underlying Datastore and return a
            % new DataLoader object
            newLoader = copy(this);
            ds = newLoader.Datastore;
            newLoader.Datastore = partition(ds,numPartitions,index);
            reset(newLoader);
        end
        
        function newLoader = subset(this,indices)
            % subset  Subset the underlying Datastore and return a new
            % DataLoader object
            newLoader = copy(this);
            ds = newLoader.Datastore;
            newLoader.Datastore = subset(ds,indices);
            reset(newLoader);
        end

        function tf = hasdata(this)
            % hasdata  True if the reader can return some data
            [tf, this.MiniBatchReader] = canReturnData(this.MiniBatchReader, this.Datastore);
        end

        function shuffleUnderlyingDatastore(this)
            % shuffleUnderlyingDatastore  Shuffle the underlying datastore
            ds = this.Datastore;
            if isShuffleable(ds)
                this.Datastore = shuffle(ds);
            else
                error(message('nnet_cnn:internal:cnn:DataLoader:NonShuffleableDatastore'));
            end
        end
        
        function tf = isSubsettable(this)
            % isSubsettable  True if the underlying datastore is
            % subsettable
            ds = this.Datastore;
            tf = isSubsettable(ds);
        end
    end

    % Accessors for Dependent properties
    methods
        function set.MiniBatchSize(this, miniBatchSize)
            % Recreate the reader with new MiniBatchSize. This also resets
            % the datastore
            this.MiniBatchReader = iCreateMiniBatchReader(this.Datastore, miniBatchSize, this.ReturnIncompleteMiniBatch);
        end

        function miniBatchSize = get.MiniBatchSize(this)
            miniBatchSize = this.MiniBatchReader.TargetMiniBatchSize;
        end

        function set.ReturnIncompleteMiniBatch(this, value)
            this.MiniBatchReader.ReturnIncompleteMiniBatch = value;
        end

        function value = get.ReturnIncompleteMiniBatch(this)
            value = this.MiniBatchReader.ReturnIncompleteMiniBatch;
        end
    end

    methods (Access = protected)
        function loaderCopy = copyElement(this)
            loaderCopy = copyElement@matlab.mixin.Copyable(this);

            % Deep copy of the underlying datastore
            ds = this.Datastore;
            loaderCopy.Datastore = copy(ds);
        end
    end

    methods (Access = private)
        function reset(this)
            % Reset the reader & datastore
            reset(this.Datastore);
            this.MiniBatchReader = reset(this.MiniBatchReader);
        end
    end

    methods(Static, Hidden)
        function obj = loadobj(s)
            % Upon loading, reset the datastore. 
            ds = s.Datastore;
            reset(ds);
            % Create a new dataloader
            obj = nnet.internal.cnn.DataLoader(ds, ...
                CollateFcn = s.CollateFcn, MiniBatchSize = s.MiniBatchSize, ...
                ReturnIncompleteMiniBatch = s.ReturnIncompleteMiniBatch);
        end
    end

    methods (Hidden)
        function s = saveobj(this)
            s.Datastore = this.Datastore;
            s.MiniBatchSize = this.MiniBatchSize;
            s.CollateFcn = this.CollateFcn;
            s.ReturnIncompleteMiniBatch = this.ReturnIncompleteMiniBatch;
        end
    end
end

function reader = iCreateMiniBatchReader(ds, targetMiniBatchSize, returnIncompleteMiniBatch)

if isa(ds, 'deep.internal.data.InMemoryArrayDatastore')
    % Specialize the DataLoader for in-memory data
    reader = deep.internal.minibatchqueue.InMemoryReader( ...
        ds, targetMiniBatchSize, returnIncompleteMiniBatch);
else
    reader = deep.internal.minibatchqueue.DatastoreReader( ...
        ds, targetMiniBatchSize, returnIncompleteMiniBatch);
end
end

function data = defaultCollateFcn(cellData)
% Convert BatchSize x NumVariables batches into 1xNumVariables cell
% array where each column contains a 4-D array or Tensor object.

% One-hot encoding of categoricals may eventually be a transform.

data = cell([1,size(cellData,2)]);
for col = 1:size(cellData,2)
    if iscategorical(cellData{1,col})
        % Is dummification something that should happen within CollateFcn
        % or should this instead be hidden away as a pure implementation
        % detail within Dispatcher or within low level interface?
        if isvector(cellData{1,col})
            % Vector response for classification.
            responses = [cellData{:,col}];
        else
            % Tensor response for semantic segmentation.
            responses = cat(4,cellData{:,col});
        end
        data{1,col} = nnet.internal.cnn.util.dummify(responses);
    else
        data{1,col} = single(cat(4,cellData{:,col}));
    end
end

end

