classdef CHOMPUtils < robotics.manip.internal.InternalAccess
%CHOMPUtils Utilities for CHOMP planner
%   UTILS=CHOMPUtils(N,CONFIG0,CONFIGNPLUS1,DT) returns a utility object useful
%   for computing costs pertaining to smoothness and collision avoidance in
%   manpiulatorCHOMP/optimize. N is the number of trajectory points for
%   computing smoothness, CONFIG0 and CONFIGNPLUS1 are column vectors which
%   specify the end points of a trajectory, and DT is the time step of the
%   discretized trajectory. The dimension of the trajectory is the number
%   of rows in CONFIG0 and CONFIGNPLUS1.
%
%   References:
%       [1] N. Ratliff, M. Zucker, J. A. Bagnell and S. Srinivasa, "CHOMP:
%       Gradient optimization techniques for efficient motion planning,"
%       2009 IEEE International Conference on Robotics and Automation,
%       2009, pp. 489-494, doi: 10.1109/ROBOT.2009.5152817.

% Copyright 2022-2023 The MathWorks, Inc.

%#codegen

    properties(SetAccess=immutable)

        %K1 First order backward finite difference matrix.
        %   For notation, refer to equation (1) in [1]
        K1

        %e1 First order backward finite difference terms at boundaries.
        %   For notation, refer to equation (1) in [1]
        e1

        %K2 Second order backward finite difference matrix.
        %   For notation, refer to equation (1) in [1]
        K2

        %e2 Second order backward finite difference terms at boundaries.
        %   For notation, refer to equation (1) in [1]
        e2

        %K3 Third order backward finite difference matrix.
        %   For notation, refer to equation (1) in [1]
        K3

        %e3 Third order backward finite difference terms at boundaries.
        %   For notation, refer to equation (1) in [1]
        e3

    end

    methods
        function obj=CHOMPUtils(n,config0,configNplus1,dt)
        %CHOMPUtils Constructor
            m=size(config0,1);
            obj.K1=kron([diag(ones(1,n))+diag(-1*ones(1,n-1),-1);
                         zeros(1,n-1),-1],eye(m))/dt;
            obj.e1=[-config0;zeros((n-1)*m,1);configNplus1]/dt;
            obj.K2=kron([diag(ones(1,n))+diag(-2*ones(1,n-1),-1)+diag(ones(1,n-2),-2);
                         zeros(1,n-2),1,-2],eye(m))/dt^2;
            obj.e2=[-config0;config0;zeros((n-2)*m,1);configNplus1]/dt^2;
            obj.K3=kron([diag(ones(1,n))+diag(-3*ones(1,n-1),-1)+diag(3*ones(1,n-2),-2)+diag(-1*ones(1,n-3),-3);
                         zeros(1,n-3),-1,3,-3],eye(m))/dt^3;
            obj.e3=[-config0;2*config0;-config0;zeros((n-3)*m,1);configNplus1]/dt^3;
        end

        function [cost,grad]=sumSquaredVelocity(obj,flattraj)
        %sumSquaredVelocity Sum of squared velocities along the trajectory
            [cost,grad]=obj.sumSquaredNDerivativeCost(obj.K1,obj.e1,flattraj);
        end

        function [cost,grad]=sumSquaredAcceleration(obj,flattraj)
        %sumSquaredAcceleration Sum of squared accelerations along the trajectory
            [cost,grad]=obj.sumSquaredNDerivativeCost(obj.K2,obj.e2,flattraj);
        end

        function [cost,grad]=sumSquaredJerk(obj,flattraj)
        %sumSquaredJerk Sum of squared jerks along the trajectory
            [cost,grad]=obj.sumSquaredNDerivativeCost(obj.K3,obj.e3,flattraj);
        end

    end

    methods(Static)
        function [cost,grad]=sumSquaredNDerivativeCost(K,e,flattraj)
        %sumSquaredNDerivativeCost Sum of squared Nth derivative cost
        %   [COST,GRAD]=sumSquaredNDerivativeCost(K,E,FLATTRAJ)
        %   Inputs:
        %       K - Nth finite difference matrix.
        %
        %       E - Boundary condition contributions
        %
        %       FLATTRAJ - trajectory is a NUMCONFIG*M-by-1 column vector where
        %       NUMCONFIG is the number of configurations in a trajectory. And
        %       M is the size of a configuration in the trajectory.
        %
        %   Outputs:
        %       COST - total cost
        %
        %       GRAD - rate of change of cost with respect to trajectory.
        %
        %   Refer to equation (1) in [1] for notation.

            c=K*flattraj+e;
            cost=0.5*(c'*c);
            grad=K'*K*flattraj+K'*e;
        end

        function bodypoints=makeBodyPointsEncodingFromRBTSpheres(treeinternal,...
                                                                 rbtsphtable)
        %makeBodyPointsEncodingFromRBTSpheres Make body points encoding
        %   BODYPOINTS=makeBodyPointsEncodingFromRBTSpheres(TREEINTERNAL,RBTSPHTABLE)
        %
        %   Inputs:
        %       TREEINTERNAL - An instance of robotics.manip.internal.RigidBodyTree
        %
        %       RBTSPHTABLE - A table with row names as bodies and base in
        %
        %       TREEINTERNAL and a column "spheres" which is a cell array
        %       containing matrices of size 4-by-NUMSPHERES, where NUMSPHERES
        %       are the number of spheres on a body/base. The first row
        %       represents the radii of the spheres on the body, and the last
        %       three rows represent their centers.
        %
        %   Outputs:
        %       BODYPOINTS - An encoding of type 5-by-NUMSPHERESONTREE, where
        %
        %       NUMSPHERESONTREE is the total number of spheres on the
        %
        %       TREEINTERNAL. The first row in BODYPOINTS are the radii of the
        %       spheres, rows 2-4 (inclusive) are the co-ordinates of the
        %       center of spheres, and the final row is the index of the body
        %       which has the sphere.

            bodypoints=zeros(5,0);
            coder.varsize('bodypoints',[5,inf],[0,1]);
            rownames=rbtsphtable.Properties.RowNames;
            for i = 1:length(rownames)
                idx=treeinternal.findBodyIndexByName(rownames{i});
                rbsph=rbtsphtable(rownames{i},'spheres').spheres{1};
                sphencoding=zeros(5,size(rbsph,2));
                sphencoding(end,:)=idx;
                sphencoding(1:4,:)=rbsph;
                bodypoints=...
                    [bodypoints,sphencoding];
            end
        end

        function [dist,grad]=distanceToNearestSphereFrom(S1,S2)
        %distanceToNearestSphereFrom Distance of spheres in a set from nearest spheres in another set
        %   [DIST,GRAD]=distanceToNearestSphereFrom(S1,S2)
        %
        %   Inputs:
        %       S1 - An array of spheres of size 4-by-N1. The first row
        %       represents the radii of the spheres, and the last three rows
        %       are their coordinates.
        %
        %       S1 - An array of spheres of size 4-by-M-by-N1. The first row
        %       represents the radii of the spheres, and the last three rows
        %       are their coordinates.
        %
        %   Outputs:
        %       DIST - Distance of sphere in S1 to its nearest sphere in S2.
        %       The size of DIST is 1-by-N1.
        %
        %       GRAD - The rate of change of distance of a sphere in S1 to its
        %       nearest sphere in S2, with respect to its coordinates. The size
        %       of GRAD is 3-by-N1.

            center=@(S)(S(2:end,:,:));
            radius=@(S)(S(1,:,:));
            reshapedS1=reshape(S1,4,1,[]);
            centerS1=center(reshapedS1);
            centerS2=center(S2);
            centerS1MinusCenterS2=centerS1-centerS2;
            distcenters=vecnorm(centerS1MinusCenterS2);
            radiusS1=radius(reshapedS1);
            radiusS2=radius(S2);
            distspheres=distcenters-radiusS1-radiusS2;
            [distanceToNearestSphereInS2,indexOfNearestSphereInS2]=min(distspheres,[],2);
            dist=reshape(distanceToNearestSphereInS2,1,[]);
            grad_=(centerS1-centerS2)./(distspheres+radiusS1+radiusS2);

            % To index selectively in a multi-dimensional array, the subscripts
            % need to be defined first. Given the gradient above, index into
            % the value of gradient for the nearest sphere in S2. The
            % values that we want in grad_ tensor are that of the index of
            % the nearest sphere in S2. Let i1, i2, i3 be the indices along
            % each of the dimension of grad_ (3D tensor).
            %    i1 - We want every value along the first dimension (of size
            %    3).
            %
            %    i2 - We want the value at index of nearest sphere in S2
            %    for every configuration in N1.
            %
            %    i3 - We want the gradient value for every N1.

            N1=size(grad_,3);
            i1=reshape(repmat((1:3)',1,N1),[],1);
            i2=reshape(repmat(indexOfNearestSphereInS2(:)',3,1),[],1);
            i3=reshape(repmat(1:N1,3,1),[],1);
            grad=reshape(grad_(sub2ind(size(grad_),i1,i2,i3)),3,[]);
        end

        function [cost,grad]=costOfSphericalCollisions(S1,epsilon,sphdistfcn)
        %costOfSphericalCollisions Cost of collision between a set of spherical collision objects and another collision model
        %   This function computes a cost and gradient given a distance
        %   function. The distance function returns the distances between
        %   the specified spherical collisions and a distinct set of
        %   collision objects defined in the function.
        %   [COST,GRAD]=costOfSphericalCollisions(S1,EPSILON,SPHDISTFCN)
        %
        %   Inputs:
        %       S1 - An array of spheres of size 4-by-N1. The first row
        %       represents the radii of the spheres, and the last three rows
        %       are their coordinates.
        %
        %       EPSILON - Clearance distance.
        %
        %       SPHDISTFCN - Distance function which accepts S and computes
        %       the nearest distance to another set of collision objects.
        %
        %   Outputs:
        %       COST - Cost of collision for a sphere in S1 to its nearest
        %       sphere in S2.
        %
        %       GRAD - The rate of change of cost of collision of a sphere
        %       in S1 to its nearest sphere in S2, with respect to its
        %       coordinates. The size of GRAD is 3-by-N1.
        %
        %   If "d" is the distance of a sphere in S1 to its nearest sphere
        %   in S2, the cost "c" is defined as such:
        %
        %   if d < 0                            % colliding
        %       c = -d + 0.5*EPSILON
        %   else if d >= 0 && d <= EPSILON      % not colliding but within clearance
        %       c = 0.5/epsilon*(d-EPSILON)^2
        %   else                                % not colliding and outside clearance
        %       c = 0
        %   end
        %
        %   The implementation below is a vectorized version of the above
        %   conditional.

            [g_,J_]=sphdistfcn(S1);
            g=g_;
            gcolliding=(g_<0);
            gcollidingbutepsilon=(g_>=0 & g_<=epsilon);
            gelse=(~gcolliding&~gcollidingbutepsilon);
            g(gcolliding)=-g_(gcolliding)+0.5*epsilon;
            g(gcollidingbutepsilon)=0.5./epsilon*(g_(gcollidingbutepsilon)-epsilon).^2;
            g(gelse)=0;
            cost=g;
            J=J_;
            J(:,gcolliding)=-J_(:,gcolliding);
            J(:,gcollidingbutepsilon)=((g_(gcollidingbutepsilon)-epsilon).*J(:,gcollidingbutepsilon))/epsilon;
            J(:,gelse)=0;
            grad=J;
        end

        function [dist,grad]= distanceFromMeshTSDF(S,tsdf)
        %distanceFromMeshTSDF Distance of spheres in a set from meshtsdf objects in the tsdf set
        %   [DIST,GRAD]=distanceFromMeshTSDF(S,TSDF)
        %
        %   Inputs:
        %       S - An array of spheres of size 4-by-N1. The first row
        %       represents the radii of the spheres, and the last three
        %       rows are their coordinates.
        %
        %       TSDF - A meshTSDF manager object.
        %
        %   Outputs:
        %       DIST - Distance of sphere in S to its nearest point on a
        %       mesh in TSDF as computed by the distance field value at
        %       that point in space.
        %
        %       GRAD - The rate of change of distance of a sphere in S1 to
        %       its nearest point in TSDF, with respect to its coordinates.
        %       The size of GRAD is 3-by-N1.
            pts=S(2:end,:);
            dist=distance(tsdf,pts')'-S(1,:);
            grad=gradient(tsdf,pts')';
        end

        function [cost,grad]=costOfTreeSphericalCollisionsInWorld(...
            S,...
            tform,jac,...
            x0,...
            dt,epsilon,...
            sphdistfcn)
        %costOfTreeSphericalCollisionsInWorld Computes the cost of a rigid body's sphere's collisions in world
        %   [COST,GRAD]=costOfTreeSphericalCollisionsInWorld(S,TFORM,JAC,X0,DT,EPSILON,SPHDISTFCN)
        %
        %   Inputs:
        %       S - An array of spheres of size 4-by-NUMCONFIG. The first
        %       row represents the radii of the spheres, and the last three
        %       rows are their coordinates. NUMCONFIG is the number of
        %       configurations of spheres in S.
        %
        %       TFORM - Homogeneous transformation of sphere S1's local
        %       frame to world frame. The size of TFORM is
        %       4-by-4-by-NUMCONFIG, where NUMCONFIG are the number of
        %       configurations of the sphere in S.
        %
        %       JAC - Velocity of the local frame expressed in the world
        %       frame with respect to the joints on the rigid body tree.
        %       JAC is a 6-by-CONFIGSIZE-by-NUMCONFIG page, where
        %       CONFIGSIZE is the size of a joint configuration of the
        %       rigid body tree, and NUMCONFIG are the number of
        %       configurations of spheres in S. The first three rows of JAC
        %       are the angular velocities and the last three are linear
        %       velocities.
        %
        %       X0 - Initial position of sphere in S.
        %
        %       DT - Time step between consecutive configurations of S.
        %
        %       EPSILON - Collision clearance.
        %
        %       SPHDISTFCN - Distance function which accepts S and computes
        %       the distance to the world obstacles.
        %
        %   Outputs:
        %       COST - Scalar cost representing the collision cost
        %       accumulated by sphere in S as it transitions over all
        %       configurations.
        %
        %       GRAD - Rate of change of COST with respect to joint
        %       configurations. GRAD is of size NUMCONFIG*CONFIGSIZE-by-1.
        %
        %   See equation (4) in [1] for notational correspondence.

            center=@(S)(S(2:end,:));

            % Transform the spheres' centers to world frame.
            % For notation refer to equation (4) in [1]
            x=robotics.manip.internal.CHOMPUtils.transformPoints(tform,center(S));
            xdot=diff([x0,x],1,2)/dt;
            xddot=diff([repmat(x0,1,2),x],2,2)/dt^2;
            normxdot=vecnorm(xdot,2,1);
            xdothat=normalize(xdot,'norm');
            xdothat(isnan(xdothat))=0;
            xdothat3=reshape(xdothat,3,1,[]);
            kappa=pagemtimes(repmat(eye(3),1,1,size(x,2))...
                             -pagemtimes(xdothat3,pagetranspose(xdothat3)),...
                             reshape(xddot,3,1,[]));
            kappa(isnan(kappa))=0;
            omega=jac(1:3,:,:);
            posxsph=S(2:end);
            Jrel=cross(omega,repmat(posxsph,1,size(omega,2),size(x,2)));
            J=jac(4:6,:,:)+Jrel;
            [c,gradc]=...
                robotics.manip.internal.CHOMPUtils.costOfSphericalCollisions(...
                [repmat(S(1),1,size(x,2));x],epsilon,sphdistfcn);
            gradadd__=repmat(eye(3),1,1,size(x,2))-pagemtimes(xdothat3,pagetranspose(xdothat3));
            gradadd_=pagemtimes(gradadd__,reshape(gradc,3,1,[]));
            grad_=pagemtimes(pagemtimes(pagetranspose(J),reshape(normxdot,1,1,[])),...
                             gradadd_-pagemtimes(reshape(c,1,1,[]),kappa));
            cost=dot(c,normxdot);
            grad=grad_(:);
            grad(isnan(grad))=0;
        end

        function y=transformPoints(tform,x)
        %transformPoints Transform points given homogeneous transform
        % Transform the spheres' centers to world frame.
            y__=pagemtimes(tform,[x;1]);
            y_=reshape(y__,4,[]);

            % For notation refer to equation (4) in [1]
            y=y_(1:3,:);
        end

        function selfcollisionmask=maskSelfAndParentCollisionIndices(bodyindices,rbt,bidx)
        %maskSelfAndParentCollisionIndices Mask body indices that are self, parent body, and its children.
        %   Input:
        %       BODYINDICES - A row vector of size 1-by-N where each value
        %       is the index of a body on the rigid body tree RBT.
        %
        %       RBT - A rigid body tree.
        %
        %       BIDX - A body index on the rigid body tree whose parent
        %       body, itself, and its children, need to be masked.
        %
        %   Output:
        %       SELFCOLLISIONMASK - A logical mask over the input
        %       BODYINDICES. A value of true at an array index in
        %       SELFCOLLISIONMASK implies that the corresponding value at
        %       that array index in BODYINDICES should be masked or not be
        %       included for self-collision costs.
            if(bidx==0)
                body=rbt.TreeInternal.Base;
            else
                body=rbt.TreeInternal.Bodies{bidx};
            end
            selfcollisionmask=bodyindices~=bidx & ...
                bodyindices~=body.ParentIndex & ...
                ~ismember(bodyindices,find(body.ChildrenIndices));

        end

        function selfcollisionmask=maskSelfAndAdjacentCollisionIndices(bodyindices,bidx)
        %maskSelfAndAdjacentCollisionIndices Mask body indices that are self and its adjacent values.
        %   Input:
        %       BODYINDICES - A row vector of size 1-by-N where each value
        %       is the index of a body on the rigid body tree RBT.
        %
        %       RBT - A rigid body tree.
        %
        %       BIDX - A body index on the rigid body tree whose adjacent
        %       bodies, and itself, needs to be masked.
        %
        %   Output:
        %       SELFCOLLISIONMASK - A logical mask over the input
        %       BODYINDICES. A value of true at an array index in
        %       SELFCOLLISIONMASK implies that the corresponding value at
        %       that array index in BODYINDICES should be masked or not be
        %       included for self-collision costs.
            selfcollisionmask=...
                bodyindices~=bidx & ...
                bodyindices~=bidx+1 & ...
                bodyindices~=bidx-1;
        end
    end
end
