# bit_star.py
# 版本: Python 3.10+
# 依赖: numpy, heapq, random

import numpy as np
import heapq
import random
from typing import Optional, List, Tuple

class BITStar:
    """
    Batch Informed Trees (BIT*) 算法

    核心思想:
    1. 用隐式RGG表示采样空间
    2. 批采样替代单样本扩展
    3. 边队列按启发式代价排序，优先检测最有希望的边
    """

    def __init__(
        self,
        bounds: List[Tuple[float, float]],
        start: np.ndarray,
        goal: np.ndarray,
        samples_per_batch: int = 500,
        rgg_radius: float = 0.3,
        collision_check_fn=None,
        seed: int = 42
    ):
        self.bounds = bounds
        self.start = start
        self.goal = goal
        self.n = len(start)
        self.samples_per_batch = samples_per_batch
        self.rgg_radius = rgg_radius
        self.collision_check = collision_check_fn or (lambda a, b: False)
        random.seed(seed)
        np.random.seed(seed)

        self.tree_nodes = [self.start]
        self.parent = {tuple(self.start): None}
        self.g = {tuple(self.start): 0.0}

        self.samples = []
        self.queue = []
        self.blacklist = set()
        self.best_solution = None
        self.best_cost = float('inf')

    def euclidean_dist(self, a: np.ndarray, b: np.ndarray) -> float:
        return float(np.linalg.norm(a - b))

    def heuristic(self, x: np.ndarray) -> float:
        return self.euclidean_dist(x, self.goal)

    def radius_for_n(self, n: int, dim: int, volume_free: float) -> float:
        c_sep = 2.0 * ((2.0 + 1.0/dim) * (volume_free / np.pi**0.5)) ** (1.0/dim)
        return c_sep * (np.log(n) / n) ** (1.0 / dim)

    def get_nearby_samples(self, node: np.ndarray, radius: float) -> List[np.ndarray]:
        nearby = []
        for s in self.samples:
            if self.euclidean_dist(node, s) <= radius:
                nearby.append(s)
        return nearby

    def get_nearby_tree_nodes(self, sample: np.ndarray, radius: float) -> List[np.ndarray]:
        nearby = []
        for n in self.tree_nodes:
            if self.euclidean_dist(sample, n) <= radius:
                nearby.append(n)
        return nearby

    def compute_edge_cost(self, parent: np.ndarray, child: np.ndarray) -> float:
        return self.euclidean_dist(parent, child)

    def add_to_queue(self, x: np.ndarray, y: np.ndarray, is_tree: bool):
        if (tuple(x), tuple(y)) in self.blacklist:
            return

        g_x = self.g.get(tuple(x), float('inf'))
        g_y_candidate = g_x + self.compute_edge_cost(x, y)
        h_y = self.heuristic(y)

        if is_tree:
            f_y = g_x + self.euclidean_dist(x, y)
        else:
            f_y = g_x + self.euclidean_dist(x, y) + h_y

        heapq.heappush(self.queue, (f_y, tuple(x), tuple(y), is_tree))

    def new_batch(self):
        new_samples = []
        for _ in range(self.samples_per_batch):
            s = np.array([
                random.uniform(self.bounds[i][0], self.bounds[i][1])
                for i in range(self.n)
            ])
            new_samples.append(s)

        self.samples.extend(new_samples)
        radius = self.radius_for_n(
            len(self.samples) + len(self.tree_nodes),
            self.n, 1.0
        )
        return new_samples, radius

    def update_queue_from_batch(self, new_samples: List[np.ndarray], radius: float):
        for s in new_samples:
            nearby_tree = self.get_nearby_tree_nodes(s, radius)
            for n in nearby_tree:
                self.add_to_queue(n, s, is_tree=True)

        for n in self.tree_nodes:
            nearby_samples = self.get_nearby_samples(n, radius)
            for s in nearby_samples:
                self.add_to_queue(s, n, is_tree=False)

    def update_tree_costs(self, updated_node: np.ndarray):
        stack = [updated_node]
        while stack:
            current = stack.pop()
            current_g = self.g[tuple(current)]

            for i, node in enumerate(self.tree_nodes):
                if self.parent.get(tuple(node)) is not None:
                    parent = self.parent[tuple(node)]
                    if np.allclose(parent, current):
                        new_g = self.g.get(tuple(current), float('inf')) + \
                                self.euclidean_dist(current, node)
                        if new_g < self.g.get(tuple(node), float('inf')):
                            self.g[tuple(node)] = new_g
                            stack.append(node)

    def solve(self, max_time: float = 30.0) -> Optional[List[np.ndarray]]:
        import time
        start_time = time.time()

        new_samples, radius = self.new_batch()
        self.update_queue_from_batch(new_samples, radius)

        while time.time() - start_time < max_time:
            if not self.queue:
                if not self.samples:
                    return self.best_solution
                new_samples, radius = self.new_batch()
                self.update_queue_from_batch(new_samples, radius)
                continue

            f, x_tuple, y_tuple, is_tree = heapq.heappop(self.queue)
            x = np.array(x_tuple)
            y = np.array(y_tuple)

            g_x = self.g.get(x_tuple, float('inf'))

            if is_tree:
                current_g_y = self.g.get(y_tuple, float('inf'))
                if g_x + self.euclidean_dist(x, y) >= current_g_y:
                    continue

            if self.collision_check(x, y):
                self.blacklist.add((x_tuple, y_tuple))
                continue

            if is_tree:
                old_parent = self.parent.get(y_tuple)
                self.parent[y_tuple] = x
                self.g[y_tuple] = g_x + self.euclidean_dist(x, y)

                if old_parent is not None:
                    self.update_tree_costs(y)

                if self.euclidean_dist(y, self.goal) < 0.1:
                    cost = self.g[y_tuple] + self.euclidean_dist(y, self.goal)
                    if cost < self.best_cost:
                        self.best_cost = cost
                        self.best_solution = self._reconstruct_path(y)

                        prune_threshold = self.best_cost * 1.1
                        self.queue = [(f, x_, y_, it) for f, x_, y_, it in self.queue
                                      if f < prune_threshold]
            else:
                g_y = g_x + self.euclidean_dist(x, y)
                if g_y < self.g.get(y_tuple, float('inf')):
                    self.add_to_queue(y, self.goal, is_tree=True)

        return self.best_solution

    def _reconstruct_path(self, goal_near: np.ndarray) -> List[np.ndarray]:
        path = [goal_near]
        current = goal_near
        while self.parent.get(tuple(current)) is not None:
            current = self.parent[tuple(current)]
            path.append(current)
        path.append(self.start)
        path.reverse()
        return path


def simple_collision(a, b):
    p1 = np.array([1.5, 1.5])
    r = 0.3
    mid = (a + b) / 2.0
    dist = np.linalg.norm(mid - p1)
    return dist < r


planner = BITStar(
    bounds=[(0, 5), (0, 5)],
    start=np.array([0.5, 0.5]),
    goal=np.array([4.5, 4.5]),
    samples_per_batch=300,
    collision_check_fn=simple_collision
)

path = planner.solve(max_time=10.0)
if path:
    print(f"BIT*找到解! 路径长度: {len(path)}, 代价: {planner.best_cost:.3f}")
else:
    print("未找到解")

# abit_star.py
# 基于BIT*的ABIT*扩展

class ABITStar(BITStar):
    """
    Advanced BIT* (ABIT*) 算法

    核心改进: 双代价函数膨胀
    - g_s: 向前膨胀因子，控制从起点的搜索范围
    - r_s: 向后膨胀因子，控制从目标的可达范围

    当g_s和r_s都很小时，只有最有希望的边才会进入队列
    """

    def __init__(self, *args,
                 g_inflation: float = 1.0,
                 r_inflation: float = 1.0,
                 **kwargs):
        super().__init__(*args, **kwargs)
        self.g_inflation = g_inflation
        self.r_inflation = r_inflation

    def should_prune_edge(self, x: np.ndarray, y: np.ndarray) -> bool:
        g_x = self.g.get(tuple(x), float('inf'))
        h_x_y = self.heuristic(x)
        h_y = self.heuristic(y)

        lower_bound = g_x * self.g_inflation + \
                      (1 - self.g_inflation) * h_x_y + \
                      self.euclidean_dist(x, y) * self.r_inflation + \
                      (1 - self.r_inflation) * h_y

        return lower_bound >= self.best_cost

    def update_queue_from_batch(self, new_samples: List, radius: float):
        for s in new_samples:
            nearby_tree = self.get_nearby_tree_nodes(s, radius)
            for n in nearby_tree:
                if not self.should_prune_edge(n, s):
                    self.add_to_queue(n, s, is_tree=True)

        for n in self.tree_nodes:
            nearby_samples = self.get_nearby_samples(n, radius)
            for s in nearby_samples:
                if not self.should_prune_edge(s, n):
                    self.add_to_queue(s, n, is_tree=False)

# lqr_cbf_rrt_star.py
import numpy as np

class LQR_CBF_RRT:
    """
    LQR-CBF-RRT* 算法

    核心思想:
    1. 用LQR预计算反馈增益矩阵
    2. 用CBF在规划时过滤不安全动作
    3. 交叉熵方法(CEM)做重要性采样，避免均匀采样
    """

    def __init__(self, A, B, Q, R, safety_gain=1.0):
        self.A = A  # 系统矩阵 [n, n]
        self.B = B  # 输入矩阵 [n, m]
        self.Q = Q  # 状态惩罚矩阵
        self.R = R  # 控制惩罚矩阵
        self.safety_gain = safety_gain

        self.K_cached = {}
        self.P_cached = {}

    def solve_lqr(self, x_goal: np.ndarray) -> np.ndarray:
        """
        求解LQR反馈增益: K = (B^T P B + R)^-1 B^T P A

        返回:
            K: 反馈增益矩阵 [m, n]
        """
        from scipy.linalg import solve_continuous_are
        P = solve_continuous_are(self.A, self.B, self.Q, self.R)
        K = np.linalg.solve(self.B.T @ P @ self.B + self.R, self.B.T @ P @ self.A)
        return K

    def compute_cbfaction(self, x: np.ndarray, x_goal: np.ndarray) -> np.ndarray:
        """
        CBF动作过滤 + LQR基础动作

        返回:
            u_safe: 通过CBF安全过滤的控制输入
        """
        K = self.solve_lqr(x_goal)
        u_lqr = -K @ (x - x_goal)

        h_x = self.safety_margin(x)
        if h_x < 0:
            grad_h = self.safety_gradient(x)
            u_safe = u_lqr + self.safety_gain * grad_h
        else:
            u_safe = u_lqr

        return np.clip(u_safe, -1.0, 1.0)

    def safety_margin(self, x: np.ndarray) -> float:
        obs = np.array([1.5, 1.5])
        r_obs = 0.3
        dist = np.linalg.norm(x - obs)
        return dist - r_obs

    def safety_gradient(self, x: np.ndarray) -> np.ndarray:
        obs = np.array([1.5, 1.5])
        dist = np.linalg.norm(x - obs) + 1e-8
        return (x - obs) / dist

    def cross_entropy_sampling(
        self,
        mean: np.ndarray,
        cov: np.ndarray,
        n_samples: int,
        percentile: float = 0.1
    ) -> tuple:
        """
        交叉熵方法(CEM)重要性采样

        用当前最好的样本更新采样分布均值和方差
        比均匀采样更可能采样到高质量解
        """
        samples = np.random.multivariate_normal(mean, cov, n_samples)
        costs = [self.path_cost(s) for s in samples]
        elite_idx = np.argsort(costs)[:int(n_samples * percentile)]

        new_mean = np.mean(samples[elite_idx], axis=0)
        new_cov = np.cov(samples[elite_idx].T)

        return new_mean, new_cov

    def path_cost(self, sample: np.ndarray) -> float:
        return float(np.linalg.norm(sample))