classdef dlCHOMPDatastore < matlab.io.Datastore & ...
        matlab.io.datastore.Partitionable & ...
        matlab.io.datastore.Shuffleable
%
%   Copyright 2023-2024 The MathWorks, Inc.

    properties (Access = private)
        %CurrentFileIndex Pointer to the current file index being read
        CurrentFileIndex double

    end

    properties (SetAccess = private)
        %FileSet Manages the AlternateFileSystemRoots for the datastore
        FileSet matlab.io.datastore.DsFileSet
    end

    properties(Dependent)
        %AlternateFileSystemRoots Use to support paths across file systems
        %   Used to support saving, loading, and processing of
        %   datastore on different file system machines or clusters.
        AlternateFileSystemRoots
    end

    methods

        function myds = dlCHOMPDatastore(location,sampleIndices)
        %dlCHOMPDatastore constructor

        % Obtain filenames from location, fetching only the paths
        % at provided indices, then reset the file index pointer
            filePaths = myds.getFilenames(location, sampleIndices);
            myds.FileSet = matlab.io.datastore.DsFileSet(filePaths);
            myds.CurrentFileIndex = 1;
            reset(myds);
        end

        function tf = hasdata(myds)
        %HASDATA Returns true if more data is available.
        %   Return logical scalar indicating availability of data.
        %   This method should be called before calling read.
        %   hasdata is used in conjunction with read to
        %   read all the data within the datastore.
            tf = hasfile(myds.FileSet);
        end

        function [data,info] = read(myds)
        %READ Read data and information about the extracted data.
        %   Return the data extracted from the datastore in the
        %   appropriate form for this datastore. Also return
        %   information about where the data was extracted from in
        %   the datastore. Both the outputs are required to be
        %   returned from the read method, and can be of any type.
        %   info is a struct with information
        %   about the chunk of data read. data represents the
        %   underlying class of tall, if tall is created on top of
        %   this datastore.

        % Read data and information about the extracted data.
            if ~hasdata(myds)
                error(message('robotics:robotmanip:dlchomp:NoDataAvailableToRead'));
            end

            fileInfoTable = nextfile(myds.FileSet);
            data = myds.parseFile(fileInfoTable);
            info.Size = size(data);
            info.FileName = fileInfoTable.FileName;
            info.Offset = fileInfoTable.Offset;

            % Update CurrentFileIndex for tracking progress
            if fileInfoTable.Offset + fileInfoTable.SplitSize >= ...
                    fileInfoTable.FileSize
                myds.CurrentFileIndex = myds.CurrentFileIndex + 1 ;
            end
        end

        function reset(myds)
        %RESET Reset to the start of the data.
        %   Reset the datastore to the state where no data has been
        %   read from it.
            reset(myds.FileSet);
            myds.CurrentFileIndex = 1;
        end

        function altRoots = get.AlternateFileSystemRoots(myds)
        %get.AlternateFileSystemRoots Getter for AlternateFileSystemRoots property
            altRoots = myds.FileSet.AlternateFileSystemRoots;
        end

        function set.AlternateFileSystemRoots(myds,altRoots)
        %set.AlternateFileSystemRoots Setter for AlternateFileSystemRoots property
            myds.FileSet.AlternateFileSystemRoots = altRoots;
            reset(myds);
        end

        function subds = partition(myds,n,ii)
        %partition  Definition needed to support parallel processing
        %   Defines how data can be partitioned for parallel processing
            subds = copy(myds);
            subds.FileSet = partition(myds.FileSet,n,ii);
            reset(subds);
        end

        function dsNew = shuffle(ds)
        %shuffle Shuffle definition needed to support shuffling
        %   DSNEW = shuffle(DS) shuffles the files and the
        %   corresponding labels in the datastore DS.
            dsNew = copy(ds);
            numObservations = dsNew.FileSet.NumFiles;
            idx = randperm(numObservations);
            dsNew.FileSet = subset(dsNew.FileSet,idx);
        end
    end

    methods (Hidden = true)
        function frac = progress(myds)
        %progress Progress tracker of data being read
            if hasdata(myds)
                frac = (myds.CurrentFileIndex-1)/...
                       myds.FileSet.NumFiles;
            else
                frac = 1;
            end
        end
    end

    methods (Access = protected)

        function dscopy = copyElement(ds)
        %copyElement Allows methods such as readall and preview to remain stateless
            dscopy = copyElement@matlab.mixin.Copyable(ds);
            dscopy.FileSet = copy(ds.FileSet);
        end

        function sortedFilepaths = sortFilepathsByIndex(~,filepaths)
        %sortedFilepaths Sorts file paths by the numeric index postfix in their
        %name
            [~,filenames,~] = fileparts(filepaths);
            fileIndices = regexp(filenames,"[0-9]+","match");
            fileIndicesLastMatch = cellfun(@(x) x(end),fileIndices);
            positions = str2double(string(fileIndicesLastMatch));
            [~,sortedIdxs] = sort(positions);
            sortedFilepaths = filepaths(sortedIdxs);
        end


        function sampleFilePaths = getFilenames(myds,location,sampleIndices)
        %sampleFilePaths Function to get filenames from folder location and sample indices
        % Get file having a JSON extension, which correspond to
        % the data sample files
            folderContent = dir(fullfile(location,"*.json"));

            % Get full path to all matching files by concatenating
            % folder names to file names, sort them, and return desired files
            listOfAllFilePaths = fullfile({folderContent.folder},{folderContent.name});
            listOfAllSortedFilePaths = myds.sortFilepathsByIndex(listOfAllFilePaths);
            sampleFilePaths = listOfAllSortedFilePaths(sampleIndices);
        end

        function n = maxpartitions(myds)
        %maxpartitions Function needed to support parallel processing
            n = maxpartitions(myds.FileSet);
        end
    end

    methods (Access=private,Static)
        function data = parseFile(fileInfoTable)
        %parseFile Used to read a given JSON file and generate a data sample
        %   JSON file schema must match dlCHOMP.DATA_SAMPLE_SCHEMA
        %
        %   See also: dlCHOMP.generateSamples
            data = cell(1,3);
            jsonText = fileread(fileInfoTable.FileName);
            dataSample = jsondecode(jsonText);

            % Obtain inputs 1 and 2 for dlCHOMP.Network from current data
            % sample struct
            trajectory = dataSample.(dlCHOMP.DATA_SAMPLE_SCHEMA.TrajectoryFieldName);
            data{1} = [trajectory(1,:),trajectory(end,:)];
            data{2} = (dataSample.(dlCHOMP.DATA_SAMPLE_SCHEMA.EncodingFieldName))';

            % Obtain outputs expected of dlCHOMP.Network from current
            % data sample struct
            data{3} = dlCHOMPDatastore.convertCHOMPOutputToNetworkOutputFormat( ...
                trajectory);
        end
    end

    methods(Static,Access={?matlab.unittest.TestCase})
        function reshapedTraj = convertCHOMPOutputToNetworkOutputFormat(...
            chompOutputTraj)
        %convertCHOMPOutputToNetworkOutputFormat Reshapes CHOMP optimized
        %output trajectory into expected network output format to train
        %DLCHOMP
        %   Reshapes a M-by-N optimized trajectory output by CHOMP into a
        %   1-by-P trajectory output by the network, where M is the number of
        %   waypoints including start and goal, N is the number of movable
        %   joints of the robot, and P is equal to (M-2) times N. M is equal
        %   to the NUMWAYPOINTS property of the DLCHOMP object

        % Network does not output start and goal which is a part of
        % CHOMP output trajectory
            reshapedTraj = reshape(chompOutputTraj(2:end-1,:)',1,[]);
        end
    end
end

% LocalWords:  HASDATA hasdata robotmanip dlchomp waypoints NUMWAYPOINTS
