classdef BidirectionalNeuralPlanner < nav.algs.internal.InternalAccess
%This class is for internal use only. It may be removed in the future.

%BidirectionalNeuralPlanner Create bidirectional neural planner
% The BidirectionalNeuralPlanner utilizes a trained motion planning network
% to iteratively generate path from the start and goal states in both
% forward and backward directions. % Note that the states generated by
% the motion planning network may not always be valid, which means that
% further re-planning may be necessary.
%
% In each iteration, the planner attempts to connect the forward and
% reverse paths using the greedy connect heuristic like the one used in
% RRT-connect.
%
% BidirectionalNeuralPlanner properties:
%   StateValidator          - State validator for planner
%   MotionPlanningNetwork   - Motion planning networks object
%   MaxLearnedStates        - Maximum number of learned states
%
% BidirectionalNeuralPlanner methods:
%   plan                    - Plan path from start to goal states

% Copyright 2023 The MathWorks, Inc.

%#codegen


    properties(SetAccess=protected)
        MotionPlanningNetwork
        StateValidator
    end

    properties(Access=public)
        MaxLearnedStates (1,1) {mustBeNumeric, mustBeInteger, mustBeNonnegative} = 50
    end


    properties(Access=?nav.algs.internal.InternalAccess)
        StateSpace
        LearnedStates
        NumLearnedStates
    end

    properties(Dependent, GetAccess=?nav.algs.internal.InternalAccess)
        NumStateVariables
    end

    properties(Abstract,Hidden)
        Environment
    end

    methods
        function obj = BidirectionalNeuralPlanner(stateValidator, mpnet)
            arguments
                stateValidator (1,1) {mustBeA(stateValidator, 'nav.StateValidator')}
                mpnet (1,1) {mustBeA(mpnet, 'nav.internal.MPNET')}
            end

            obj.StateValidator = stateValidator;
            obj.MotionPlanningNetwork = mpnet;
            obj.StateSpace = stateValidator.StateSpace;
            obj.LearnedStates = nan(obj.MaxLearnedStates, obj.NumStateVariables);
            obj.NumLearnedStates = 0; % Iterator to track learned states
                                      % for successive calls to plan method
        end
    end


    methods(Access=?nav.algs.internal.InternalAccess)

        function pathStates = planNeuralPath(obj, startState, goalState)
        % planNeuralPath Plan path between startState and goalState
        % using bidirectional neural planner


        % Initial output path
            pathStates = zeros(0, obj.StateSpace.NumStateVariables);

            % Initialize bidirectional paths, pathA and pathB. pathA is the
            % forward path from the start state and pathB is the reverse
            % path from the goal state
            bufferSize = ceil(obj.MaxLearnedStates/2);
            pathA = nan(bufferSize, obj.NumStateVariables);
            pathB = nan(bufferSize, obj.NumStateVariables);
            iA = 1;
            iB = 1;
            pathA(iA,:) = startState;
            pathB(iB,:) = goalState;

            % During the planning, pathA and pathB are swapped at each
            % iteration. So we keep track of whether pathA is forward path
            % or not. This information will be useful while concatenating
            % pathA and pathB at the end of planning.
            pathAIsForwardPath = true;

            % Generate bidirectional path using MPNet
            while obj.NumLearnedStates < obj.MaxLearnedStates
                % Stop when neural sample count exceeds MaxLearnedStates

                if obj.StateValidator.isMotionValid(pathA(iA,:), pathB(iB,:))
                    % Stop when bidirectional paths meet

                    pathA = pathA(~isnan(pathA(:,1)), :);
                    pathB = pathB(~isnan(pathB(:,1)), :);

                    if ~pathAIsForwardPath
                        % Make sure pathA is the forward path before
                        % concatenating them
                        [pathA, pathB] = deal(pathB, pathA);
                    end

                    % Concatenate forward (pathA) and reverse (pathB) paths
                    pathStates = [pathA; flipud(pathB)];
                    break
                end

                % Predict next state using MPNet
                stateNext = obj.MotionPlanningNetwork.predict(...
                    pathA(iA,:), pathB(iB,:), obj.Environment);

                % Add predicted state to the pathA
                iA = iA + 1;
                pathA(iA,:) = stateNext;

                % Increment neural sample count
                obj.NumLearnedStates = obj.NumLearnedStates + 1;

                % Store the learned states for debugging information
                obj.LearnedStates(obj.NumLearnedStates,:) = stateNext;

                % Swap forward and reverse paths
                [iA, iB] = deal(iB, iA);
                [pathA, pathB] = deal(pathB, pathA);
                pathAIsForwardPath = ~pathAIsForwardPath;
            end
        end
    end

    methods
        function numStateVariables = get.NumStateVariables(obj)
            numStateVariables = obj.StateValidator.StateSpace.NumStateVariables;
        end
    end
end
