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

class RRTNode:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.parent = None
        self.cost = 0.0

class RRTPlanner:
    """快速扩展随机树路径规划"""
    
    def __init__(self, start, goal, obstacle_list, 
                 rand_area, expand_dis=2.0, goal_sample_rate=0.1, max_iter=500):
        self.start = RRTNode(start[0], start[1])
        self.goal = RRTNode(goal[0], goal[1])
        self.obstacle_list = obstacle_list
        self.min_rand, self.max_rand = rand_area
        self.expand_dis = expand_dis
        self.goal_sample_rate = goal_sample_rate
        self.max_iter = max_iter
        self.node_list = [self.start]
    
    def planning(self) -> Optional[List[RRTNode]]:
        for i in range(self.max_iter):
            # 采样随机点
            if random.random() > self.goal_sample_rate:
                rnd = RRTNode(random.uniform(self.min_rand, self.max_rand),
                              random.uniform(self.min_rand, self.max_rand))
            else:
                rnd = self.goal
            
            # 找最近节点
            nearest_node = self._get_nearest_node(rnd)
            
            # 向随机点扩展
            new_node = self._steer(nearest_node, rnd)
            
            # 检查是否碰撞
            if self._check_collision(new_node):
                self.node_list.append(new_node)
                
                # 检查是否到达目标
                if self._calc_dist_to_goal(new_node) <= self.expand_dis:
                    final_node = self._steer(new_node, self.goal)
                    if self._check_collision(final_node):
                        return self._generate_path(final_node)
        
        return None
    
    def _get_nearest_node(self, node):
        dlist = [(n.x - node.x)**2 + (n.y - node.y)**2 for n in self.node_list]
        return self.node_list[int(np.argmin(dlist))]
    
    def _steer(self, from_node, to_node):
        dx, dy = to_node.x - from_node.x, to_node.y - from_node.y
        d = np.sqrt(dx**2 + dy**2)
        
        if d < self.expand_dis:
            new_node = RRTNode(to_node.x, to_node.y)
        else:
            theta = np.arctan2(dy, dx)
            new_node = RRTNode(from_node.x + self.expand_dis * np.cos(theta),
                               from_node.y + self.expand_dis * np.sin(theta))
        
        new_node.parent = from_node
        new_node.cost = from_node.cost + self.expand_dis
        return new_node
    
    def _check_collision(self, node):
        for (ox, oy, size) in self.obstacle_list:
            dx = ox - node.x
            dy = oy - node.y
            d = np.sqrt(dx**2 + dy**2)
            if d <= size + 0.1:
                return False
        return True
    
    def _calc_dist_to_goal(self, node):
        return np.sqrt((node.x - self.goal.x)**2 + (node.y - self.goal.y)**2)
    
    def _generate_path(self, node):
        path = []
        while node.parent is not None:
            path.append([node.x, node.y])
            node = node.parent
        path.append([self.start.x, self.start.y])
        return path[::-1]