classdef textDatastore < matlab.io.Datastore & ...
        matlab.io.datastore.MiniBatchable
    
    properties
        ClassNames
        Datastore
        EmbeddingDimension
        LabelName
        MiniBatchSize
        NumClasses
    end
    
    properties(SetAccess = protected)
        NumObservations
    end
    
    properties(Access = private)
        CurrentFileIndex
    end
    
    methods
        function ds = textDatastore(filename,textName,labelName,emb)
            % ds = textDatastore(filename,textName,labelName,emb) creates a
            % text datastore.
            %
            % Input Arguments:
            %     filename  - CSV file
            %     textName  - Column of table containing text data
            %     labelName - Column of table containing labels
            %     emb       - Word embedding
            
            % Set default mini-batch size.
            miniBatchSize = 128;
            ds.MiniBatchSize = miniBatchSize;
            ds.CurrentFileIndex = 1;
            
            % Create underlying tabular text datastore.
            variableNames = [textName labelName];
            ttds = tabularTextDatastore(filename, ...
                'TextType','string', ...
                'SelectedVariableNames',variableNames, ...
                'ReadSize',miniBatchSize);
            ds.LabelName = labelName;
            
            % Read class names from data and count observations.
            classNames = [];
            numObservations = 0;
            while hasdata(ttds)
                data = read(ttds);
                labels = data{:,2};
                numObservations = numObservations + size(data,1);
                classNames = unique([classNames labels']);
            end
            ds.ClassNames = classNames;
            ds.NumClasses = numel(classNames);
            ds.NumObservations = numObservations;
            
            % Create transformed datastore.
            tds = transform(ttds,@(data) preprocessTable(data,emb,classNames));
            ds.Datastore = tds;
            ds.EmbeddingDimension = emb.Dimension;
        end
        
        function tf = hasdata(ds)
            % tf = hasdata(ds) returns true if more data is available.
            
            miniBatchSize = ds.MiniBatchSize;
            currentFileIndex = ds.CurrentFileIndex;
            numObservations = ds.NumObservations;
            tf = (currentFileIndex + miniBatchSize - 1) <= numObservations;
        end
        
        function [data,info] = read(ds)
            % [data,info] = read(ds) reads one mini-batch of data.
            
            tds = ds.Datastore;
            miniBatchSize = ds.MiniBatchSize;
            
            % Read data.
            tds.UnderlyingDatastore.ReadSize = miniBatchSize;
            ds.CurrentFileIndex = ds.CurrentFileIndex + miniBatchSize;
            
            [data,info] = read(tds);
        end
        
        function labels = readLabels(ds)
            % labels = readLabels(ds) reads all labels from the datastore.
            
            labelName = ds.LabelName;
            ttds = copy(ds.Datastore.UnderlyingDatastore);
            reset(ttds);
            
            ttds.SelectedVariableNames = labelName;
            tbl = readall(ttds);
            labels = categorical(tbl.(labelName));
        end
        
        function reset(ds)
            % reset(ds) resets the datastore to the start of the data.
            
            tds = ds.Datastore;
            reset(tds);
            ds.CurrentFileIndex = 1;
        end
    end
    
    methods (Hidden = true)
        function frac = progress(ds)
            % frac = progress(ds) returns the percentage of observations
            % read in the datastore.
            
            numObservations = ds.NumObservations;
            currentFileIndex = ds.CurrentFileIndex;
            frac = (currentFileIndex - 1) / numObservations;
        end
    end
end

function dataPreprocessed = preprocessTable(data,emb,classNames)

% Preprocess documents.
textData = data{:,1};
documents = tokenizedDocument(textData);
documents = lower(documents);
documents = erasePunctuation(documents);

% Convert documents to sequences of word vectors.
predictors = doc2sequence(emb,documents);

% Read labels.
labels = data{:,2};
responses = categorical(labels,classNames);

% Return data as table.
dataPreprocessed = table(predictors,responses);

end