classdef exampleHelperTruckPropagator < nav.StatePropagator
    
    properties
        TravelBias = 0;
    end
    
    properties (SetAccess = protected)
        ControlParams
        
        HeuristicCosts
        
        StateValidator
    end
    
    properties (Dependent, SetAccess = protected)
        MaxNumSteps
    end
    
    properties (Hidden)
        StateValidatorInternal %REMOVE
    end
    
    methods
        function obj = exampleHelperTruckPropagator(validator, controlParams, distanceHeuristic)
            
            % Construct base propagator
            ss = validator.StateSpace;
            dt = controlParams.StepSize;
            numControlVariable = 4;
            obj = obj@nav.StatePropagator(ss,dt,numControlVariable);

            % Set up custom validator
            obj.StateValidator = validator;
            
            % Add parameters to object
            obj.ControlParams = controlParams;
            obj.HeuristicCosts = distanceHeuristic;
            obj.StateValidatorInternal = validator; %REMOVE
        end
        
        function maxNumStep = get.MaxNumSteps(obj)
            maxNumStep = obj.ControlParams.MaxNumSteps;
        end
        
        function [iStates, iCommands, numStep] = propagate(obj, q0, uRef, qTgt, maxNumStep)
            % Retrieve propagation step size
            dt = obj.ControlStepSize;
            
            % Propagate system
            [iStates, iCommands] = exampleHelperPropagateTruck(obj.StateSpace.TruckParams, ...
                obj.ControlParams, dt, q0, uRef, qTgt, maxNumStep);
            
            targetReached = exampleHelperTargetReachedFunction(iStates,qTgt);
            
            if any(targetReached)
                i = find(targetReached,1,'first');
                iStates((i+1):end,:) = [];
                iCommands((i+1):end,:) = [];
            end
            
            numStep = ones(size(iStates,1),1);
        end
        
        function [iStates, iCommands, numStep] = propagateWhileValid(obj, q0, u0, qTgt, maxNumStep)
            
            [iStates, iCommands, numStep] = obj.propagate(q0,u0,qTgt,maxNumStep);
            
            % Validate the propagated trajectory
            isValid = obj.StateValidator.isStateValid(iStates);
            
            % Find the end of the free trajectory segment
            lastValid = find(~isValid,1,'first');
            
            % Clear states lying beyond the final state
            iStates(lastValid:end,:) = [];
            iCommands(lastValid:end,:) = [];
            numStep(lastValid:end,:) = [];
        end
        
        function cost = distance(obj, state1, state2)
        %distance Estimates the distance between two states
        %
        %   COST = distance(OBJ, STATE1, STATE2) The truck trailer state
        %   space uses a pre-computed lookup-table to estimate distances
        %   between two states. STATE1 is an N-by-S matrix of initial
        %   states, and STATE2 is a 1-by-S target state. 
        %
        %   For each distance calculation, STATE2 is mapped to STATE1's
        %   local coordinate frame, giving a local distance, (xL, yL). The
        %   distance value is then calculated using a bicubic interpolant.
        %   The 2D cost matrix used in the lookup is determined by STATE2's
        %   velocity. 
        %
        %   If the lookup value falls outside the bounds of the 2D lookup 
        %   region, (xL, yL) is projected to the grid, and an additional
        %   cost is added, proportional to the projection distance.
            
            % Collect needed vars from target state
            xGoal = state2(1);
            yGoal = state2(2);
            params = obj.StateSpace.TruckParams;
            
            % Calculate correct reference frame depending on direction
            if state2(5) > 0
                % Reference frame is truck rear-axle
                state = exampleHelperExpandToFullState(state1, params.L2, params.M);
                costValues = obj.HeuristicCosts.ForwardCost;
            else
                % Reference frame is trailer rear-axle
                state = state1;
                costValues = obj.HeuristicCosts.ReverseCost;
            end
            
            % Extract reference info
            xStart = state(:,1);
            yStart = state(:,2);
            theta = state(:,3);
            
            % Vectorized decomposition of coordinate transform
            t = theta;
            cTh = cos(t);
            sTh = sin(t);
            
            % Transform error to body frame
            xEff0 = cTh*xGoal + sTh*yGoal - cTh.*xStart - sTh.*yStart;
            yEff0 = -sTh*xGoal + cTh*yGoal + sTh.*xStart - cTh.*yStart;
            
            % Grab interpolant points
            xInd = obj.HeuristicCosts.XPoints;
            yInd = obj.HeuristicCosts.YPoints;
            
            % Check if points fall outside range (assumes equal +/- limits)
            th = atan2(yEff0,xEff0);
            
            cost = zeros(numel(th),1);
            
            % Mark
            xOut = abs(xEff0) > xInd(end) & (abs(xEff0) > abs(yEff0));
            yOut = abs(yEff0) > yInd(end) & (abs(xEff0) <= abs(yEff0));
            xEff = xEff0;
            yEff = yEff0;
            xEff(xOut) = xInd(end)*sign(xEff0(xOut));
            yEff(xOut) = xEff(xOut).*tan(th(xOut));
            yEff(yOut) = yInd(end)*sign(yEff0(yOut));
            xEff(yOut) = yEff(yOut)./tan(th(yOut));
            outPts = xOut | yOut;
            cost(outPts) = 10*((xEff0(outPts)-xEff(outPts)).^2+(yEff0(outPts)-yEff(outPts)).^2);
            
            % Evaluate lookup table for points within bounds
            cost = cost + interp2(xInd, yInd, costValues, xEff, yEff, 'cubic');
            m = isnan(cost) | isinf(cost);
            unreachableCost = 10000 + xEff0(m).^2 + yEff0(m).^2;
            cost(m) = unreachableCost;
            
            if rand() < obj.TravelBias
                cost = state1(:,end).^2+cost;
            end
        end
        
        function [u, numStep] = sampleControl(obj, q0, ~, qTgt)
        %sampleControl Capture reference information fow low-level controllers
        %
        %   Rather than generating control inputs that are directly
        %   provided the system's low-level actuators, this helper class 
        %   demonstrates how a state propagator can be designed to capture
        %   high level information in the control signal. This information
        %   can serve as a reference for lower level controllers, or 
        %   capture application-specific information that doesn't fall 
        %   neatly into the "state" or "target" category.
        %
        %   [U, NUMSTEP] = sampleControl(OBJ, Q0, U0, QTGT) Takes in the 
        %   current state, Q0, previous control reference information, U0,
        %   and target state, QTGT. In this example, the propagator stores 
        %   the system's initial xy location in the control command, U, and
        %   always attempts to take the maximum number of steps towards the
        %   target.
            
            % Store the reference information and return max number steps.
            maxStep = obj.MaxNumSteps;
            numStep = randi([ceil(maxStep/10),maxStep]);
            vDir = qTgt(5);
            if vDir == 1
                % Reference point is center of truck's rear-axle
                R = obj.ControlParams.ForwardLookahead;
            else
                % Reference point is center of trailer's rear-axle
                R = obj.ControlParams.ReverseLookahead;
            end
            u = [q0(1:2) vDir R];
        end
        
        function set.TravelBias(obj, bias)
            validateattributes(bias,{'numeric'},{'>=',0,'<=',1},'TravelBias');
        end
        
        function setup(obj)
            obj.StateValidator.UseMEX = coder.target('MATLAB') && exist('exampleHelperCollisionChecker_mex','file') == 3;
        end
        
        function cObj = copy(obj)
            cObj = obj;
        end
    end
end