"""
Hybrid A*
@author: Huiming Zhou
"""

import os
import sys
import json
import math
import heapq
from heapdict import heapdict
import time
import numpy as np
import matplotlib.pyplot as plt
import scipy.spatial.kdtree as kd
import scipy.ndimage as ndimage
import itertools
from matplotlib.collections import LineCollection
import cvxpy as cp


import Hybrid_A_starPlanner_RS.MotionPlanning.HybridAstarPlanner.astar as astar
import Hybrid_A_starPlanner_RS.MotionPlanning.HybridAstarPlanner.draw as draw
import Hybrid_A_starPlanner_RS.MotionPlanning.CurvesGenerator.reeds_shepp as rs



class C:  # Parameter config
    PI = math.pi

    XY_RESO = 1  # [m]
    YAW_RESO = np.deg2rad(15.0)  # [rad]
    MOVE_STEP = 0.6  # [m] path interporate resolution
    N_STEER = 20.0  # steer command number
    COLLISION_CHECK_STEP = 2  # skip number for collision check
    EXTEND_BOUND = 0  # collision check range extended

    GEAR_COST = 1e6  # switch back penalty cost
    BACKWARD_COST = 500.0  # backward penalty cost
    STEER_CHANGE_COST = 1.0  # steer angle change penalty cost
    STEER_ANGLE_COST = 3.0  # steer angle penalty cost
    H_COST = 15.0  # Heuristic cost penalty cost
    EUC_COST_WEIGHT = 40.0 ##欧式距离启发权重
    TIME_LIMIT = 15

    RF = 7.59  # [m] distance from rear to vehicle front end of vehicle
    RB = 2.1  # [m] distance from rear to vehicle back end of vehicle
    W = 3.65  # [m] width of vehicle
    WD = 0.7 * W  # [m] distance between left-right wheels
    WB = 3.5  # [m] Wheel base
    TR = 0.5  # [m] Tyre radius
    TW = 1  # [m] Tyre width
    MAX_STEER = 0.26  # [rad] maximum steering angle
    MAX_C = 1/13.5

    RS_CALL_RANGE = 50.0 #RS调用起点
# ---- 新增 hmap 降采样倍率 ----
    HMAP_SCALE = 3  # 将 hmap 分辨率缩小为原来的 1/4，再插值回去

    PLOT = True  # 总开关：若 False，则不做任何绘图
    ANIMATE = False  # 若 True，则逐步动画，否则只画最终结果

class Node:
    def __init__(self, xind, yind, yawind, direction,
                 x, y, yaw, directions,
                 steer, cost, pind, ind=None, is_rs=False):
        self.xind = xind; self.yind = yind; self.yawind = yawind
        self.direction = direction
        self.x = x; self.y = y; self.yaw = yaw; self.directions = directions
        self.steer = steer; self.cost = cost; self.pind = pind
        self.ind = ind    # 缓存扁平索引
        self.is_rs = is_rs  # ← 新增：标记是否 RS 扩展节点

class Para:
    def __init__(self, minx, miny, minyaw, maxx, maxy, maxyaw,
                 xw, yw, yaww, xyreso, yawreso, ox, oy, kdtree):
        self.minx, self.miny, self.minyaw = minx, miny, minyaw
        self.maxx, self.maxy, self.maxyaw = maxx, maxy, maxyaw
        self.xw, self.yw, self.yaww = xw, yw, yaww
        self.xyreso, self.yawreso = xyreso, yawreso
        self.ox, self.oy, self.kdtree = ox, oy, kdtree

class Path:
    def __init__(self, x, y, yaw, direction, cost, is_rs_list=None,
                 t=None, v=None, a=None, j=None):
        self.x = x
        self.y = y
        self.yaw = yaw
        self.direction = direction
        self.cost = cost
        self.is_rs_list = is_rs_list

        # 以下四个字段用来存放速度规划结果
        self.t = t      # 时间戳 list, len == len(x)
        self.v = v      # 速度 list
        self.a = a      # 加速度 list
        self.j = j      # jerk list

class QueuePrior:
    def __init__(self):
        self.queue = heapdict()
    def empty(self): return len(self.queue)==0
    def put(self, key, pri): self.queue[key] = pri
    def get(self): return self.queue.popitem()[0]

def calc_parameters(ox, oy, xyreso, yawreso, kdtree):
    minx = round(min(ox)/xyreso); maxx = round(max(ox)/xyreso)
    miny = round(min(oy)/xyreso); maxy = round(max(oy)/xyreso)
    xw = maxx - minx + 1; yw = maxy - miny + 1
    minyaw = round(-C.PI/yawreso)-1; maxyaw = round(C.PI/yawreso)
    yaww = maxyaw - minyaw + 1
    return Para(minx, miny, minyaw, maxx, maxy, maxyaw,
                xw, yw, yaww, xyreso, yawreso, ox, oy, kdtree)

def calc_motion_set():
    s = np.linspace(-C.MAX_STEER, C.MAX_STEER, int(2*C.N_STEER)+1)
    steer = s.tolist()
    direc = [1.0]*len(steer) + [-1.0]*len(steer)
    steer = steer + steer
    return steer, direc

def calc_index(node, P):
    return ((node.yawind - P.minyaw)*P.xw*P.yw +
            (node.yind   - P.miny )*P.xw +
            (node.xind   - P.minx))


def hybrid_astar_planning(sx, sy, syaw,
                          gx, gy, gyaw,
                          ox, oy,
                          xyreso, yawreso,
                          cusp_type,
                          in_path_flag,
                          time_limit=None
                          ):

    # 记录起始时间
    t_start = time.time()
    # 1) 起终点离散化
    sxr, syr = round(sx/xyreso), round(sy/xyreso)
    gxr, gyr = round(gx/xyreso), round(gy/xyreso)
    syawr = round(rs.pi_2_pi(syaw)/yawreso)
    gyawr = round(rs.pi_2_pi(gyaw)/yawreso)

    nstart = Node(sxr, syr, syawr, 1,
                  [sx], [sy], [syaw], [1],
                  0.0, 0.0, -1)
    ngoal  = Node(gxr, gyr, gyawr, 1,
                  [gx], [gy], [gyaw], [1],
                  0.0, 0.0, -1)
    # 缓存目标点的连续坐标，用于欧式距离计算
    goal_x, goal_y = ngoal.x[-1], ngoal.y[-1]

    # 2) KDTree + 参数
    kdtree = kd.KDTree(list(zip(ox, oy)))
    P = calc_parameters(ox, oy, xyreso, yawreso, kdtree)

    # 3) 构造 hmap（粗分辨率 + 线性插值）
    t0 = time.time()
    scale = C.HMAP_SCALE
    if scale > 1:
        h_coarse = np.array(
            astar.calc_holonomic_heuristic_with_obstacle(
                ngoal, P.ox, P.oy, xyreso*scale, C.EXTEND_BOUND))
        hmap = ndimage.zoom(h_coarse, zoom=scale, order=1)
        hmap = hmap[:P.xw, :P.yw]
    else:
        hmap = np.array(
            astar.calc_holonomic_heuristic_with_obstacle(
                ngoal, P.ox, P.oy, xyreso, C.EXTEND_BOUND))
    # print(f"[hmap] shape={hmap.shape}  cost={(time.time()-t0):.3f}s")

    # 4) 扁平化启发式 & 常量
    h_flat = hmap.ravel()
    HW = P.xw * P.yw
    minx, miny = P.minx, P.miny

    # 5) 初始化 open/closed，并把起点加入
    start_id = calc_index(nstart, P)
    nstart.ind = start_id
    # 起点在 h_flat 的索引（仅 x/y，不含 yaw）
    start_xy = (nstart.xind - minx)*P.yw + (nstart.yind - miny)
    open_set, closed_set = {start_id: nstart}, {}
    qp = QueuePrior()
    qp.put(start_id, nstart.cost + C.H_COST * h_flat[start_xy])

    steer_set, direc_set = calc_motion_set()

    # 6) 主循环
    while open_set:
        # —— 超时检测 ——
        if time_limit is not None and time.time() - t_start > time_limit:
            # print(f"[hybrid_astar_planning] time limit {time_limit}s exceeded, abort.")
            return None

        sid = qp.get()
        curr = open_set.pop(sid)
        closed_set[sid] = curr

        # analytic 扩展
        cx, cy = curr.x[-1], curr.y[-1]
        gx, gy = ngoal.x[-1], ngoal.y[-1]
        if math.hypot(cx - gx, cy - gy) <= C.RS_CALL_RANGE:
            update, fpath = update_node_with_analystic_expantion(curr, ngoal, P, cusp_type, in_path_flag)
        else:
            update, fpath = False, None

        if update:
            fnode = fpath
            break

        # 运动原语扩展
        for u, d in zip(steer_set, direc_set):
            nb = calc_next_node(curr, curr.ind, u, d, P)
            if not nb: continue

            nid = calc_index(nb, P)
            nb.ind = nid
            if nid in closed_set: continue

            # 计算对应 x/y 在 h_flat 中的索引
            xyid = (nb.xind - minx)*P.yw + (nb.yind - miny)
            # 新增欧式距离到终点
            dist_to_goal = math.hypot(nb.x[-1] - goal_x, nb.y[-1] - goal_y)
            # 最终 f 值：g + H_COST*h + EUC_COST_WEIGHT*dist
            fcost = (nb.cost + C.H_COST * h_flat[xyid] + C.EUC_COST_WEIGHT * dist_to_goal)

            if nid not in open_set or open_set[nid].cost > nb.cost:
                open_set[nid] = nb
                qp.put(nid, fcost)

    # 7) 回溯路径
    if not update:
        return None
    return extract_path(closed_set, fnode, nstart)


def extract_path(closed, ngoal, nstart):
    rx, ry, ryaw, direc, is_rs = [], [], [], [], []
    cost = 0.0
    node = ngoal

    while True:
        # 逆序累加当前节点的所有点和它的 is_rs 标记
        rx += node.x[::-1]
        ry += node.y[::-1]
        ryaw += node.yaw[::-1]
        direc += node.directions[::-1]
        # node.is_rs 是一个标记，把它复制 len(node.x) 次
        is_rs += [node.is_rs] * len(node.x)

        cost += node.cost
        if is_same_grid(node, nstart):
            break
        node = closed[node.pind]

    # 正序
    rx = rx[::-1]
    ry = ry[::-1]
    ryaw = ryaw[::-1]
    direc = direc[::-1]
    is_rs = is_rs[::-1]

    # 修正第一个方向
    direc[0] = direc[1]

    path = Path(rx, ry, ryaw, direc, cost, is_rs)
    return path

def calc_next_node(n_curr, c_id, u, d, P):
    step = C.XY_RESO * 2

    nlist = math.ceil(step / C.MOVE_STEP)
    xlist = [n_curr.x[-1] + d * C.MOVE_STEP * math.cos(n_curr.yaw[-1])]
    ylist = [n_curr.y[-1] + d * C.MOVE_STEP * math.sin(n_curr.yaw[-1])]
    yawlist = [rs.pi_2_pi(n_curr.yaw[-1] + d * C.MOVE_STEP / C.WB * math.tan(u))]

    for i in range(nlist - 1):
        xlist.append(xlist[i] + d * C.MOVE_STEP * math.cos(yawlist[i]))
        ylist.append(ylist[i] + d * C.MOVE_STEP * math.sin(yawlist[i]))
        yawlist.append(rs.pi_2_pi(yawlist[i] + d * C.MOVE_STEP / C.WB * math.tan(u)))

    xind = round(xlist[-1] / P.xyreso)
    yind = round(ylist[-1] / P.xyreso)
    yawind = round(yawlist[-1] / P.yawreso)

    if not is_index_ok(xind, yind, xlist, ylist, yawlist, P):
        return None

    cost = 0.0

    if d > 0:
        direction = 1
        cost += abs(step)
    else:
        direction = -1
        cost += abs(step) * C.BACKWARD_COST

    if direction != n_curr.direction:  # switch back penalty
        cost += C.GEAR_COST

    cost += C.STEER_ANGLE_COST * abs(u)  # steer angle penalyty
    cost += C.STEER_CHANGE_COST * abs(n_curr.steer - u)  # steer change penalty
    cost = n_curr.cost + cost

    directions = [direction for _ in range(len(xlist))]

    node = Node(
        xind, yind, yawind, direction,
        xlist, ylist, yawlist, directions,
        u, cost, c_id,
        ind=None,
        is_rs=False  # ← 普通扩展
    )

    return node


def is_index_ok(xind, yind, xlist, ylist, yawlist, P):
    # （1）原有的地图边界检查
    if xind <= P.minx or xind >= P.maxx or \
       yind <= P.miny or yind >= P.maxy:
        return False

    # （3）原有的碰撞检查
    ind = range(0, len(xlist), C.COLLISION_CHECK_STEP)
    nodex = [xlist[k] for k in ind]
    nodey = [ylist[k] for k in ind]
    nodeyaw = [yawlist[k] for k in ind]
    if is_collision(nodex, nodey, nodeyaw, P):
        return False

    return True


def update_node_with_analystic_expantion(n_curr, ngoal, P, cusp_type, in_path_flag):
    path = analystic_expantion(n_curr, ngoal, P, cusp_type, in_path_flag)  # rs path: n -> ngoal

    if not path:
        return False, None

    fx = path.x[1:-1]
    fy = path.y[1:-1]
    fyaw = path.yaw[1:-1]
    fd = path.directions[1:-1]

    cost = calc_rs_path_cost(path)
    fcost = n_curr.cost + cost
    fpind = calc_index(n_curr, P)
    fsteer = 0.0

    fpath = Node(
        n_curr.xind, n_curr.yind, n_curr.yawind, n_curr.direction,
        fx, fy, fyaw, fd,
        fsteer, fcost, fpind,
        ind=None,
        is_rs=True  # ← RS 扩展
    )

    return True, fpath

def angle_difference(rad1, rad2):
    # 将角度归一化到 [-pi, pi) 区间
    rad1_normalized = (rad1 + np.pi) % (2 * np.pi) - np.pi
    rad2_normalized = (rad2 + np.pi) % (2 * np.pi) - np.pi

    # 计算差异并取模 2pi，然后取最小的绝对值差异
    diff = (rad1_normalized - rad2_normalized) % (2 * np.pi)
    return min(diff, 2 * np.pi - diff)

def analystic_expantion(node, ngoal, P, cusp_type, in_path_flag):
    sx, sy, syaw = node.x[-1], node.y[-1], node.yaw[-1]
    gx, gy, gyaw = ngoal.x[-1], ngoal.y[-1], ngoal.yaw[-1]

    maxc = math.tan(C.MAX_STEER) / C.WB
    paths = rs.calc_all_paths(
        sx, sy, syaw, gx, gy, gyaw, maxc, step_size=C.MOVE_STEP)

    if not paths:
        return None

    pq = QueuePrior()
    for path in paths:
        # 如果对人字尖朝向有限制，那么先检测人字尖朝向是否符合要求
        if in_path_flag and cusp_type != "both":
            cusp_idx = 0
            for i in range(len(path.directions) - 1):
                if path.directions[i] * path.directions[i + 1] < 0.0:
                    cusp_idx = i
                    break
            end_x, end_y, end_yaw = path.x[-1], path.y[-1], path.yaw[-1]
            cusp_x, cusp_y = path.x[cusp_idx], path.y[cusp_idx]
            # 计算终点指向人字尖的向量角度，(-pi, pi],atan2会自动处理x坐标相同的问题
            end_toward_cusp_yaw = math.atan2((cusp_y-end_y), (cusp_x-end_x))
            # 终点左前方45度,并统一到(-pi, pi]
            end_left_front_yaw = (end_yaw + np.pi/4 + np.pi) % (2 * np.pi) - np.pi
            end_left_front_yaw = np.pi if end_left_front_yaw==-np.pi else end_left_front_yaw
            # 终点右前方45度,并统一到(-pi, pi]
            end_right_front_yaw = (end_yaw - np.pi/4 + np.pi) % (2 * np.pi) - np.pi
            end_right_front_yaw = np.pi if end_right_front_yaw==-np.pi else end_right_front_yaw

            if angle_difference(end_left_front_yaw, end_toward_cusp_yaw) < angle_difference(end_right_front_yaw, end_toward_cusp_yaw):
                actual_cusp_type = "left"
            else:
                actual_cusp_type = "right"
            if actual_cusp_type != cusp_type:
                continue
        cost = calc_rs_path_cost(path)
        pq.put(path, cost)


    while not pq.empty():
        path = pq.get()
        ind = range(0, len(path.x), C.COLLISION_CHECK_STEP)
        pathx = [path.x[k] for k in ind]
        pathy = [path.y[k] for k in ind]
        pathyaw = [path.yaw[k] for k in ind]

        # 碰撞过滤
        if is_collision(pathx, pathy, pathyaw, P):
            continue

        # 既不碰撞，也都在可行驶区，接受这条 RS 解
        return path

    return None


def is_collision(x, y, yaw, P):
    for ix, iy, iyaw in zip(x, y, yaw):
        d = 1
        dl = (C.RF - C.RB) / 2.0
        r = (C.RF + C.RB) / 2.0 + d

        cx = ix + dl * math.cos(iyaw)
        cy = iy + dl * math.sin(iyaw)

        ids = P.kdtree.query_ball_point([cx, cy], r)

        if not ids:
            continue

        for i in ids:
            xo = P.ox[i] - cx
            yo = P.oy[i] - cy
            dx = xo * math.cos(iyaw) + yo * math.sin(iyaw)
            dy = -xo * math.sin(iyaw) + yo * math.cos(iyaw)

            if abs(dx) < r and abs(dy) < C.W / 2 + d:
                return True

    return False


def calc_rs_path_cost(rspath):
    cost = 0.0

    for lr in rspath.lengths:
        if lr >= 0:
            cost += 1
        else:
            cost += abs(lr) * C.BACKWARD_COST

    for i in range(len(rspath.directions) - 1):
        if rspath.directions[i] * rspath.directions[i + 1] < 0.0:
            cost += C.GEAR_COST

    for ctype in rspath.ctypes:
        if ctype != "S":
            cost += C.STEER_ANGLE_COST * abs(C.MAX_STEER)

    nctypes = len(rspath.ctypes)
    ulist = [0.0 for _ in range(nctypes)]

    for i in range(nctypes):
        if rspath.ctypes[i] == "R":
            ulist[i] = -C.MAX_STEER
        elif rspath.ctypes[i] == "WB":
            ulist[i] = C.MAX_STEER

    for i in range(nctypes - 1):
        cost += C.STEER_CHANGE_COST * abs(ulist[i + 1] - ulist[i])

    return cost


def calc_hybrid_cost(node, hmap, P):
    cost = node.cost + \
           C.H_COST * hmap[node.xind - P.minx][node.yind - P.miny]

    return cost


def calc_motion_set():
    s = np.arange(C.MAX_STEER / C.N_STEER,
                  C.MAX_STEER, C.MAX_STEER / C.N_STEER)

    steer = list(s) + [0.0] + list(-s)
    direc = [1.0 for _ in range(len(steer))]
    steer = steer

    return steer, direc


def is_same_grid(node1, node2):
    if node1.xind != node2.xind or \
            node1.yind != node2.yind or \
            node1.yawind != node2.yawind:
        return False

    return True


def calc_index(node, P):
    ind = (node.yawind - P.minyaw) * P.xw * P.yw + \
          (node.yind - P.miny) * P.xw + \
          (node.xind - P.minx)

    return ind


def calc_parameters(ox, oy, xyreso, yawreso, kdtree):
    minx = round(min(ox) / xyreso)
    miny = round(min(oy) / xyreso)
    maxx = round(max(ox) / xyreso)
    maxy = round(max(oy) / xyreso)

    xw, yw = maxx - minx, maxy - miny

    minyaw = round(-C.PI / yawreso) - 1
    maxyaw = round(C.PI / yawreso)
    yaww = maxyaw - minyaw

    return Para(minx, miny, minyaw, maxx, maxy, maxyaw,
                xw, yw, yaww, xyreso, yawreso, ox, oy, kdtree)


def draw_car(x, y, yaw, steer, color='black'):
    car = np.array([[-C.RB, -C.RB, C.RF, C.RF, -C.RB],
                    [C.W / 2, -C.W / 2, -C.W / 2, C.W / 2, C.W / 2]])

    wheel = np.array([[-C.TR, -C.TR, C.TR, C.TR, -C.TR],
                      [C.TW / 4, -C.TW / 4, -C.TW / 4, C.TW / 4, C.TW / 4]])

    rlWheel = wheel.copy()
    rrWheel = wheel.copy()
    frWheel = wheel.copy()
    flWheel = wheel.copy()

    Rot1 = np.array([[math.cos(yaw), -math.sin(yaw)],
                     [math.sin(yaw), math.cos(yaw)]])

    Rot2 = np.array([[math.cos(steer), math.sin(steer)],
                     [-math.sin(steer), math.cos(steer)]])

    frWheel = np.dot(Rot2, frWheel)
    flWheel = np.dot(Rot2, flWheel)

    frWheel += np.array([[C.WB], [-C.WD / 2]])
    flWheel += np.array([[C.WB], [C.WD / 2]])
    rrWheel[1, :] -= C.WD / 2
    rlWheel[1, :] += C.WD / 2

    frWheel = np.dot(Rot1, frWheel)
    flWheel = np.dot(Rot1, flWheel)

    rrWheel = np.dot(Rot1, rrWheel)
    rlWheel = np.dot(Rot1, rlWheel)
    car = np.dot(Rot1, car)

    frWheel += np.array([[x], [y]])
    flWheel += np.array([[x], [y]])
    rrWheel += np.array([[x], [y]])
    rlWheel += np.array([[x], [y]])
    car += np.array([[x], [y]])

    plt.plot(car[0, :], car[1, :], color)
    plt.plot(frWheel[0, :], frWheel[1, :], color)
    plt.plot(rrWheel[0, :], rrWheel[1, :], color)
    plt.plot(flWheel[0, :], flWheel[1, :], color)
    plt.plot(rlWheel[0, :], rlWheel[1, :], color)
    draw.Arrow(x, y, yaw, C.WB * 0.8, color)


def read_grid_map(json_file_path):
    """读取地图数据并生成障碍物字典"""
    obstacles = {}
    # 读取JSON文件
    with open(json_file_path, 'r') as file:
        data = json.load(file)

        # 提取地图信息
        width = data['metadata']['dimensions']['width']
        height = data['metadata']['dimensions']['height']

        # 初始化地图数据数组
        map_data = np.zeros((width, height), dtype=int)

        # 填充地图数据
        for index in data['data']:
            # 计算二维数组的坐标
            x = index % width  # 余数，列位置
            y = index // width
            map_data[x, y] = 1  # 将障碍物标记为1
            obstacles[(x, y)] = True

    return map_data, height, width, obstacles


def resample_path(path, N, preserve_shifts=True):
    """
    将原始 path 重采样为恰好 N 个点。
    如果 preserve_shifts=True，则强制保留每个 direction 切换处的原始点。
    path 必须带有 .x, .y, .yaw, .direction, .is_rs_list 属性。
    """
    # 原始路径数据
    x = np.array(path.x)
    y = np.array(path.y)
    yaw = np.array(path.yaw)
    direction = np.array(path.direction)
    is_rs = np.array(path.is_rs_list, dtype=bool)

    # 1) 计算累积弧长 s
    dx = np.diff(x)
    dy = np.diff(y)
    ds = np.hypot(dx, dy)
    s = np.concatenate(([0.], np.cumsum(ds)))  # len = M

    s_total = s[-1]
    if s_total == 0:
        raise ValueError("路径长度为0，无法重采样。")

    # 2) 找到所有 gear-shift 处的 s_shift
    #    切换发生在 i->i+1，取 i+1 处的 s 值
    idx = np.where(direction[:-1] != direction[1:])[0] + 1
    s_shift = s[idx] if preserve_shifts else np.empty(0)
    M_shift = len(s_shift)

    if preserve_shifts and M_shift >= N:
        raise ValueError(f"N={N} 太小，无法同时保留 {M_shift} 个切换点。")

    # 3) 在剩余弧长上等距取点
    N_uniform = N - M_shift if preserve_shifts else N
    s_uniform = np.linspace(0, s_total, N_uniform)

    # 4) 合并、去重、排序
    if preserve_shifts:
        s_new = np.concatenate((s_uniform, s_shift))
    else:
        s_new = s_uniform
    s_new = np.unique(s_new)
    # 如果去重后少于 N，补齐端点附近的点
    if len(s_new) < N:
        extra = np.linspace(0, s_total, N - len(s_new))
        s_new = np.unique(np.concatenate((s_new, extra)))
    # 最后再截断到前 N 个
    s_new = np.sort(s_new)[:N]

    # 5) 插值 x,y
    x_new = np.interp(s_new, s, x)
    y_new = np.interp(s_new, s, y)

    # 6) 插值 yaw：先 unwrap 再 wrap
    yaw_un = np.unwrap(yaw)
    yaw_i  = np.interp(s_new, s, yaw_un)
    yaw_new = (yaw_i + np.pi) % (2*np.pi) - np.pi

    # 7) 离散量 direction / is_rs 最近邻
    #    查找每个 s_new 对应的原始索引
    idx_nn = np.searchsorted(s, s_new, side='right') - 1
    idx_nn = np.clip(idx_nn, 0, len(direction)-1)
    dir_new   = direction[idx_nn].tolist()
    is_rs_new = is_rs[idx_nn].tolist()

    # 8) 检查：如果最终点数不是 N，就抛错
    final_len = len(x_new)
    if final_len != N or len(yaw_new) != N or len(dir_new) != N or len(is_rs_new) != N:
        raise ValueError(
            f"重采样失败：期望点数 {N}，但得到 {final_len} (dir={len(dir_new)}, is_rs={len(is_rs_new)})"
        )

    # 8) 返回新的 Path（复用原 cost）
    return Path(
        x_new.tolist(),
        y_new.tolist(),
        yaw_new.tolist(),
        dir_new,
        path.cost,
        is_rs_new
    )

import numpy as np

def speed_profile_with_asymmetric_limits(
    path,
    v_max_global: float,
    a_lat_max: float,
    a_acc_max: float,   # 最大加速能力 [m/s²]
    a_dec_max: float,   # 最大制动能力 [m/s²]（取正值）
    j_max: float = None, # 新增：最大 jerk [m/s³]，不限制时传 None
    gamma: float = 1e-2  # 提升速度的线性权重，越大速度提升越多
):
    import numpy as np
    import cvxpy as cp

    # 1) 提取数据
    x = np.array(path.x);   y = np.array(path.y)
    yaw = np.array(path.yaw)
    direction = np.array(path.direction, dtype=int)
    N = len(x)
    assert N > 2, "点太少"

    # 2) 等弧长 Δs
    ds = np.hypot(x[1]-x[0], y[1]-y[0])

    # 3) 曲率限速
    kappa = np.zeros(N)
    kappa[1:-1] = np.abs((yaw[2:] - yaw[:-2]) / (2*ds))
    kappa[0], kappa[-1] = kappa[1], kappa[-2]
    v_max_curve = np.sqrt(a_lat_max/np.maximum(kappa, 1e-6))
    v_ub = np.minimum(v_max_curve, v_max_global)

    # 4) 换档点
    shift_idx = np.where(direction[:-1]!=direction[1:])[0] + 1
    shift_idx = shift_idx[(shift_idx>0)&(shift_idx<N-1)]

    # 5) 优化变量
    v = cp.Variable(N)

    # 6) 目标：jerk 二次项 - γ·sum(v)
    dv2 = v[2:] - 2*v[1:-1] + v[:-2]  # 离散二阶差分
    obj = cp.Minimize(cp.sum_squares(dv2) - gamma*cp.sum(v))

    # 7) 约束集合
    cons = []
    # 7.1 速度上下限
    cons += [v >= 0, v <= v_ub]
    # 7.2 起终点、换档点 v=0
    cons += [v[0]==0, v[-1]==0]
    for idx in shift_idx:
        cons += [v[idx]==0]
    # 7.3 区分加速 & 制动
    for i in range(1, N):
        cons += [v[i] - v[i-1] <= a_acc_max * ds]
        cons += [v[i-1] - v[i] <= a_dec_max * ds]

    # 7.4 （可选）jerk 上下限
    if j_max is not None:
        # j ≈ dv2 / ds^2，因此 |dv2| ≤ j_max * ds^2
        cons += [ dv2 <=  j_max * ds**2 ]
        cons += [ dv2 >= -j_max * ds**2 ]

    # 8) 求解
    prob = cp.Problem(obj, cons)
    prob.solve(solver=cp.OSQP)

    # 9) 导出结果
    v_opt = np.maximum(v.value, 0)
    v_safe = np.maximum(v_opt, 1e-3)
    v_front = v_safe[:-1]
    v_latter = v_safe[1:]
    v_used_for_calc_t = (v_front + v_latter)/2
    dt = ds / v_used_for_calc_t
    # 避免某些情况点与点之间的时间间隔过大的问题，限制点与点之间时间间隔最大为3s
    dt = np.minimum(dt, 3)
    t = np.concatenate(([0.], np.cumsum(dt[:])))
    a = np.zeros_like(v_opt)
    a[1:] = (v_opt[1:] - v_opt[:-1]) / dt[:]
    j = np.zeros_like(v_opt)
    j[2:] = (a[2:]     - a[1:-1])  / dt[1:]
    # # —9— 打印统计信息
    # print(f"[speed_profile] speed: min={v_opt.min():.3f}, max={v_opt.max():.3f}")
    # print(f"[speed_profile] accel: min={a.min():.3f}, max={a.max():.3f}, |a|_max={np.abs(a).max():.3f}")
    # print(f"[speed_profile] jerk:  min={j.min():.3f}, max={j.max():.3f}, |j|_max={np.abs(j).max():.3f}")

    path.t = t.tolist()
    path.v = v_opt.tolist()
    path.a = a.tolist()
    path.j = j.tolist()
    return path

def hybrid_Astar_with_speed(
    sx, sy, syaw,
    gx, gy, gyaw,
    ox, oy,
    xyreso=C.XY_RESO,
    yawreso=C.YAW_RESO,
    time_limit=C.TIME_LIMIT,
    # 重采样参数
    N_resample=128,
    preserve_shifts=True,
    # 速度规划参数
    v_max_global=8.0,
    a_lat_max=0.8,
    a_acc_max=0.8,
    a_dec_max=2.0,
    j_max=2.0,
    gamma=1e-2
):
    """
    一步完成：
      1) 调用 hybrid_astar_planning 获得初步轨迹
      2) 调用 resample_path 重采样到 N_resample 点
      3) 调用 speed_profile_with_asymmetric_limits 做速度规划

    返回：带时间、速度、加速度、jerk 的 Path 对象，失败返回 None。
    """
    # 1) Hybrid A* 规划
    path = hybrid_astar_planning(
        sx, sy, syaw,
        gx, gy, gyaw,
        ox, oy,
        xyreso, yawreso,
        time_limit=time_limit
    )
    if path is None:
        # print("[plan_with_speed] Hybrid A* 规划失败")
        return None

    # 2) 重采样
    try:
        path = resample_path(path, N_resample, preserve_shifts=preserve_shifts)
    except Exception as e:
        # print(f"[plan_with_speed] 重采样失败: {e}")
        return None

    # 3) 速度规划
    path = speed_profile_with_asymmetric_limits(
        path,
        v_max_global=v_max_global,
        a_lat_max=a_lat_max,
        a_acc_max=a_acc_max,
        a_dec_max=a_dec_max,
        j_max=j_max,
        gamma=gamma
    )
    return path

def HA_path(map_path, test_start_list, test_goal_list, in_path_flags, path_fig_dir, cusp_type):
    print("start!") 

    # 读取地图数据
    _, height, width, obstacles = read_grid_map(map_path)
    print("height:", height)
    print("width:", width)
    print("maxc:", math.tan(C.MAX_STEER) / C.WB)
    print("max_steer:", np.arctan(C.MAX_C*C.WB))

    # 生成障碍物点列表
    ox, oy = [], []
    for (x, y) in obstacles.keys():
        ox.append(x)
        oy.append(y)
    #增加地图边界
    for i in range(width):
        ox.append(i)
        oy.append(0)
        ox.append(i)
        oy.append(height-1)
    for i in range(height):
        ox.append(0)
        oy.append(i)
        ox.append(width-1)
        oy.append(i)

    for idx, (start, goal) in enumerate(zip(test_start_list, test_goal_list)):
        print(f"Test case {idx} ")

        sx, sy, syaw0 = start
        gx, gy, gyaw0 = goal

        t0 = time.time()
        path = hybrid_astar_planning(
            sx, sy, syaw0, gx, gy, gyaw0,
            ox, oy, C.XY_RESO, C.YAW_RESO,cusp_type=cusp_type,
            in_path_flag = in_path_flags[idx], time_limit = C.TIME_LIMIT
        )
        path = resample_path(path, N=128, preserve_shifts=True)

        t1 = time.time()

        # 2) 做速度规划（传入物理限幅参数）
        path = speed_profile_with_asymmetric_limits(
            path,
            v_max_global=10.0,
            a_lat_max=6.0,
            a_acc_max=2.0,
            a_dec_max=2.0,
            j_max=2.0,
            gamma=5e-1
        )

        print(f"")

        print("running T:", t1 - t0)

        if not path:
            print("  Searching failed!")
            continue

        x, y, yaw, direction, t = path.x, path.y, path.yaw, path.direction, path.t

        if not C.PLOT:
            continue

        # 2) 外部可视化
        if C.ANIMATE:
            # 动画模式：逐步显示
            for k in range(len(x)):
                plt.cla()
                # 绘制障碍物
                plt.plot(ox, oy, "sk", markersize=2)
                # 绘制全路径
                plt.plot(x, y, "o", markersize=1.5)

                # 计算当前舵角
                if k < len(x) - 2:
                    dy = (yaw[k + 1] - yaw[k]) / C.MOVE_STEP
                    steer = rs.pi_2_pi(math.atan2(-C.WB * dy, direction[k]))
                else:
                    steer = 0.0

                # 画终点与车辆
                draw_car(gx, gy, gyaw0, 0.0, 'dimgray')
                draw_car(x[k], y[k], yaw[k], steer)
                plt.title(f"Hybrid A* - Case {idx + 1}")
                plt.axis("equal")
                plt.pause(0.0001)

            plt.show()

        else:
            # 假设已有以下变量：
            # ox, oy      障碍物坐标列表
            # x, y, yaw   路径点坐标与航向（弧度制）
            # path.v      速度列表

            # 1) 静态图 + 速度色带
            fig, ax = plt.subplots(figsize=(8, 8))
            ax.plot(ox, oy, "sk", markersize=1)

            v = np.array(path.v)  # len == len(x)
            points = np.vstack([x, y]).T  # shape (N,2)
            segs = np.concatenate([points[:-1, None, :],
                                   points[1:, None, :]],
                                  axis=1)  # (N-1, 2, 2)

            lc = LineCollection(segs,
                                cmap="viridis",
                                norm=plt.Normalize(v.min(), v.max()),
                                linewidth=1.5)
            lc.set_array(v[:-1])
            ax.add_collection(lc)

            cbar = fig.colorbar(lc, ax=ax)
            cbar.set_label("Speed [m/s]")

            # # 2) 航向箭头 (quiver) 其实是车头朝向
            # # 每隔 step 个点画一个箭头，避免过密
            # step = max(1, len(x) // 20)
            # # 箭头方向由 yaw 给出，长度 scale 可调
            # ax.quiver(
            #     x[::step],  # 起点 x
            #     y[::step],  # 起点 y
            #     np.cos(yaw[::step]),  # 箭头在 x 方向的分量
            #     np.sin(yaw[::step]),  # 箭头在 y 方向的分量
            #     color="red",
            #     width=0.005,
            #     scale=20,
            #     label="heading"
            # )

            draw_car(gx, gy, gyaw0, 0.0, color="dimgray")
            draw_car(x[-1], y[-1], yaw[-1], 0.0, color="black")

            ax.set_title(f"Hybrid A* — idx {idx}, path_time:{path.t[-1]} ")
            ax.set_aspect('equal', 'datalim')
            ax.set_xlabel("X [m]")
            ax.set_ylabel("Y [m]")

            savefig_path = os.path.join(path_fig_dir, f"hybrid_a_star_case_{idx}_yaw.png")
            plt.savefig(savefig_path, dpi=300)
            # plt.show()
        print("Done!")


if __name__ == '__main__':
    map_path = os.path.join(os.path.dirname(__file__), 'MVTP_2.json')
    path_fig_dir = os.path.join(os.path.dirname(__file__), 'path_figs')
    # 限制人字尖点方向, left为左前方、right为右前方、both为不限制
    cusp_type = "left"
    # 定义测试起点和终点
    test_start_list = [
        [240.0,60.0,np.deg2rad(180.0)],
        [52.0,60.0,np.deg2rad(0.0)]
    ]

    test_goal_list = [
        [52.0,60.0,np.deg2rad(0.0)],
        [250.0, 50.0, np.deg2rad(0.0)]
    ]
    in_path_flags = [True, False]

    HA_path(map_path, test_start_list, test_goal_list, in_path_flags, path_fig_dir, cusp_type)
