import math
import scipy.integrate
import scipy.optimize
import numpy as np
# import matplotlib.pyplot as plt
from collections import namedtuple
from roboticstoolbox.mobile import *
def solvepath(x, q0=[0, 0, 0], **kwargs):
    # x[:4] is 4 coeffs of curvature polynomial
    # x[4] is total path length
    # q0 is initial state of the vehicle
    maxcurvature = 0
    def dotfunc(s, q, poly):
        # q = (x, y, θ)
        # qdot = (cosθ, sinθ, ϰ)
        k = poly[0] * s**3 + poly[1] * s**2 + poly[2] * s + poly[3]
        # k = ((poly[0] * s + poly[1]) * s  + poly[2]) * s + poly[3]
        # save maximum curvature for this path solution
        nonlocal maxcurvature
        maxcurvature = max(maxcurvature, abs(k))
        theta = q[2]
        return math.cos(theta), math.sin(theta), k
    cpoly = x[:4]
    s_f = x[4]
    sol = scipy.integrate.solve_ivp(dotfunc, [0, s_f], q0, args=(cpoly,), **kwargs)
    return sol.y, maxcurvature
def xcurvature(x):
    # inequality constraint function, must be non-negative
    _, maxcurvature = solvepath(x, q0=(0, 0, 0))
    return maxcurvature
def costfunc(x, start, goal):
    # final cost of path from start with params
    # p[0:4] is polynomial: k0, a, b, c
    # p[4] is s_f
    # integrate the path for this curvature polynomial and length
    # path is 3xN
    path, _ = solvepath(x, q0=start)
    # cost is configuration error at end of path
    e = np.linalg.norm(path[:, -1] - np.r_[goal])
    return e
[docs]class CurvaturePolyPlanner(PlannerBase):
    def __init__(self, curvature=None):
        """
        Continuous curvature path planner
        :param curvature: , defaults to None
        :type curvature: _type_, optional
        ==================   ========================
        Feature              Capability
        ==================   ========================
        Plan                 :math:`\SE{2}`
        Obstacle avoidance   No
        Curvature            Continuous
        Motion               Forwards only
        ==================   ========================
        Example:
        .. runblock:: pycon
            >>> from roboticstoolbox import DubinsPlanner
            >>> from math import pi
            >>> planner = CurvaturePolyPlanner()
            >>> path, status = planner.query(start=(0, 0, pi/2), goal=(1, 0, pi/2))
            >>> print(path[:5,:])
            >>> print(status)
        """
        super().__init__(ndims=3)
        self.curvature = curvature
[docs]    def query(self, start, goal):
        r"""
        Find a path betwee two configurations
        :param start: start configuration :math:`(x, y, \theta)`
        :type start: array_like(3), optional
        :param goal: goal configuration :math:`(x, y, \theta)`
        :type goal: array_like(3), optional
        :return: path and status
        :rtype: ndarray(N,3), namedtuple
        The path comprises points equally spaced at a distance of ``stepsize``.
        The returned status value has elements:
        +-------------------+-----------------------------------------------+
        | Element           |  Description                                  |
        +-------------------+-----------------------------------------------+
        | ``length``        | total path length                             |
        +-------------------+-----------------------------------------------+
        | ``maxcurvature``  | maximum curvature on path                     |
        +-------------------+-----------------------------------------------+
        | ``poly``          | curvature polynomial coefficients             |
        +-------------------+-----------------------------------------------+
        """
        goal = np.r_[goal]
        start = np.r_[start]
        self._start = start
        self._goal = goal
        # initial estimate of path length is Euclidean distance
        d = np.linalg.norm(goal[:2] - start[:2])
        # state vector is kappa_0, a, b, c, s_f
        if self.curvature is not None:
            nlcontraints = (
                scipy.optimize.NonlinearConstraint(xcurvature, 0, self.curvature),
            )
        else:
            nlcontraints = ()
        sol = scipy.optimize.minimize(
            costfunc,
            [0, 0, 0, 0, d],
            constraints=nlcontraints,
            bounds=[(None, None), (None, None), (None, None), (None, None), (d, None)],
            args=(start, goal),
        )
        print(sol)
        path, maxcurvature = solvepath(
            sol.x, q0=start, dense_output=True, max_step=1e-2
        )
        status = namedtuple("CurvaturePolyStatus", ["length", "maxcurvature", "poly"])(
            sol.x[4], maxcurvature, sol.x[:4]
        )
        return path.T, status  
if __name__ == "__main__":
    from math import pi
    # start = (1, 1, pi/4)
    # goal = (-3, -3, -pi/4)
    # start = (0, 0, -pi/4)
    # goal = (1, 2, pi/4)
    start = (0, 0, pi / 2)
    goal = (1, 0, pi / 2)
    planner = CurvaturePolyPlanner()
    path, status = planner.query(start, goal)
    print("start", path[0, :])
    print("goal", path[-1, :])
    print(status)
    planner.plot(path, block=True)
    ## attempt polynomial scaling, doesnt seem to work
    # sf = status.s_f
    # c = status.poly
    # print(c)
    # print(solvepath(np.r_[c, sf], start))
    # for i in range(4):
    #     c[i] /= sf ** (3-i)
    # print(solvepath(np.r_[c, 1], start))
    # print(c)