#!/usr/bin/env python3
"""Lightweight grid A*/Theta* planner for random circular obstacles."""

import heapq
import math

import rclpy
from geometry_msgs.msg import PoseStamped
from nav_msgs.msg import Odometry, Path as NavPath
from rclpy.duration import Duration
from rclpy.node import Node
from rclpy.time import Time
from std_msgs.msg import Float32MultiArray, String
from tf2_ros import Buffer, TransformException, TransformListener


def yaw_from_quaternion(q):
    return math.atan2(
        2.0 * (q.w * q.z + q.x * q.y),
        1.0 - 2.0 * (q.y * q.y + q.z * q.z),
    )


def normalize_angle(angle):
    return math.atan2(math.sin(angle), math.cos(angle))


def make_pose(frame_id, stamp, x, y, yaw):
    pose = PoseStamped()
    pose.header.frame_id = frame_id
    pose.header.stamp = stamp
    pose.pose.position.x = float(x)
    pose.pose.position.y = float(y)
    pose.pose.position.z = 0.0
    pose.pose.orientation.z = math.sin(0.5 * yaw)
    pose.pose.orientation.w = math.cos(0.5 * yaw)
    return pose


class GridAstarThetaPlanner(Node):
    def __init__(self):
        super().__init__("grid_astar_theta_planner")

        self.declare_parameter("map_frame", "map")
        self.declare_parameter("odom_topic", "/odom")
        self.declare_parameter("goal_pose_topic", "/goal_pose")
        self.declare_parameter("obstacles_topic", "/random_obstacles")
        self.declare_parameter("plan_topic", "/plan")
        self.declare_parameter("status_topic", "/grid_planner_status")
        self.declare_parameter("map_min_x", -5.0)
        self.declare_parameter("map_min_y", -5.0)
        self.declare_parameter("map_max_x", 5.0)
        self.declare_parameter("map_max_y", 5.0)
        self.declare_parameter("resolution", 0.10)
        self.declare_parameter("obstacle_format", "xyr")
        self.declare_parameter("default_obstacle_radius", 0.15)
        self.declare_parameter("robot_radius", 0.14)
        self.declare_parameter("localization_error", 0.10)
        self.declare_parameter("latency_margin", 0.05)
        self.declare_parameter("safety_margin", 0.05)
        self.declare_parameter("unknown_obstacle_radius", 0.15)
        self.declare_parameter("use_theta_star", True)
        self.declare_parameter("allow_diagonal", True)
        self.declare_parameter("plan_on_obstacle_update", True)
        self.declare_parameter("max_iterations", 50000)
        self.declare_parameter("path_resolution", 0.10)
        self.declare_parameter("smooth_iterations", 2)
        self.declare_parameter("max_odom_age", 0.80)
        self.declare_parameter("tf_timeout", 0.05)

        self.map_frame = self.get_parameter("map_frame").value
        self.map_min_x = float(self.get_parameter("map_min_x").value)
        self.map_min_y = float(self.get_parameter("map_min_y").value)
        self.map_max_x = float(self.get_parameter("map_max_x").value)
        self.map_max_y = float(self.get_parameter("map_max_y").value)
        self.resolution = float(self.get_parameter("resolution").value)
        self.obstacle_format = str(self.get_parameter("obstacle_format").value)
        self.default_obstacle_radius = float(
            self.get_parameter("default_obstacle_radius").value
        )
        self.robot_radius = float(self.get_parameter("robot_radius").value)
        self.localization_error = float(
            self.get_parameter("localization_error").value
        )
        self.latency_margin = float(self.get_parameter("latency_margin").value)
        self.safety_margin = float(self.get_parameter("safety_margin").value)
        self.unknown_obstacle_radius = float(
            self.get_parameter("unknown_obstacle_radius").value
        )
        self.use_theta_star = bool(self.get_parameter("use_theta_star").value)
        self.allow_diagonal = bool(self.get_parameter("allow_diagonal").value)
        self.plan_on_obstacle_update = bool(
            self.get_parameter("plan_on_obstacle_update").value
        )
        self.max_iterations = int(self.get_parameter("max_iterations").value)
        self.path_resolution = float(self.get_parameter("path_resolution").value)
        self.smooth_iterations = int(self.get_parameter("smooth_iterations").value)
        self.max_odom_age = float(self.get_parameter("max_odom_age").value)
        self.tf_timeout = float(self.get_parameter("tf_timeout").value)

        if self.resolution <= 0.0:
            raise RuntimeError("resolution must be positive")
        if self.map_max_x <= self.map_min_x or self.map_max_y <= self.map_min_y:
            raise RuntimeError("invalid map bounds")

        self.width = int(math.ceil((self.map_max_x - self.map_min_x) / self.resolution))
        self.height = int(math.ceil((self.map_max_y - self.map_min_y) / self.resolution))
        if self.width <= 1 or self.height <= 1:
            raise RuntimeError("map bounds are too small for the selected resolution")

        self.tf_buffer = Buffer()
        self.tf_listener = TransformListener(self.tf_buffer, self)
        self.latest_odom = None
        self.goal_pose = None
        self.obstacles = []
        self.last_tf_warn_ns = 0

        self.plan_pub = self.create_publisher(
            NavPath, self.get_parameter("plan_topic").value, 10
        )
        self.status_pub = self.create_publisher(
            String, self.get_parameter("status_topic").value, 10
        )

        self.odom_sub = self.create_subscription(
            Odometry,
            self.get_parameter("odom_topic").value,
            self._odom_callback,
            10,
        )
        self.goal_sub = self.create_subscription(
            PoseStamped,
            self.get_parameter("goal_pose_topic").value,
            self._goal_callback,
            10,
        )
        self.obstacles_sub = self.create_subscription(
            Float32MultiArray,
            self.get_parameter("obstacles_topic").value,
            self._obstacles_callback,
            10,
        )

        self._publish_status(
            f"ready size={self.width}x{self.height} resolution={self.resolution:.3f}"
        )

    def _odom_callback(self, msg):
        self.latest_odom = msg

    def _goal_callback(self, msg):
        transformed = self._transform_pose_msg(msg)
        if transformed is None:
            self._publish_status("goal_rejected transform_failed")
            return
        self.goal_pose = transformed
        self.plan()

    def _obstacles_callback(self, msg):
        self.obstacles = self._parse_obstacles(msg.data)
        if self.plan_on_obstacle_update and self.goal_pose is not None:
            self.plan()

    def _parse_obstacles(self, values):
        stride = 3 if self.obstacle_format == "xyr" else 2
        obstacles = []
        if stride <= 0:
            return obstacles
        count = len(values) // stride
        for index in range(count):
            base = index * stride
            x = float(values[base])
            y = float(values[base + 1])
            radius = self.default_obstacle_radius
            if stride == 3:
                radius = max(0.0, float(values[base + 2]))
            obstacles.append((x, y, radius))
        return obstacles

    def plan(self):
        start_pose = self._current_pose()
        if start_pose is None:
            self._publish_status("plan_rejected no_recent_odom")
            return False
        if self.goal_pose is None:
            self._publish_status("plan_rejected no_goal")
            return False

        start = self._world_to_grid(start_pose[0], start_pose[1])
        goal = self._world_to_grid(self.goal_pose[0], self.goal_pose[1])
        if start is None or goal is None:
            self._publish_status(f"plan_rejected out_of_bounds start={start} goal={goal}")
            return False

        occupied = self._build_occupancy(start, goal)
        if occupied[start[1]][start[0]] or occupied[goal[1]][goal[0]]:
            self._publish_status("plan_rejected start_or_goal_occupied")
            return False

        path_cells = self._search(start, goal, occupied)
        if not path_cells:
            self._publish_status("plan_failed no_path")
            return False

        points = [self._grid_to_world(cell[0], cell[1]) for cell in path_cells]
        points[0] = (start_pose[0], start_pose[1])
        points[-1] = (self.goal_pose[0], self.goal_pose[1])
        if self.smooth_iterations > 0:
            points = self._smooth_by_line_of_sight(points, occupied)
        points = self._densify(points)
        self._publish_path(points)
        self._publish_status(
            f"plan_ok cells={len(path_cells)} points={len(points)} obstacles={len(self.obstacles)}"
        )
        return True

    def _build_occupancy(self, start, goal):
        occupied = [[False for _ in range(self.width)] for _ in range(self.height)]
        base_padding = self.robot_radius + self.localization_error + self.latency_margin + self.safety_margin
        for obs_x, obs_y, obs_radius in self.obstacles:
            radius = max(obs_radius, self.unknown_obstacle_radius) + base_padding
            min_cell = self._world_to_grid(obs_x - radius, obs_y - radius, clamp=True)
            max_cell = self._world_to_grid(obs_x + radius, obs_y + radius, clamp=True)
            radius_sq = radius * radius
            for gy in range(min_cell[1], max_cell[1] + 1):
                for gx in range(min_cell[0], max_cell[0] + 1):
                    wx, wy = self._grid_to_world(gx, gy)
                    if (wx - obs_x) * (wx - obs_x) + (wy - obs_y) * (wy - obs_y) <= radius_sq:
                        occupied[gy][gx] = True
        occupied[start[1]][start[0]] = False
        occupied[goal[1]][goal[0]] = False
        return occupied

    def _search(self, start, goal, occupied):
        neighbors = self._neighbors8 if self.allow_diagonal else self._neighbors4
        open_heap = []
        heapq.heappush(open_heap, (0.0, start))
        came_from = {start: start}
        cost_so_far = {start: 0.0}
        iterations = 0

        while open_heap and iterations < self.max_iterations:
            iterations += 1
            _, current = heapq.heappop(open_heap)
            if current == goal:
                return self._reconstruct_path(came_from, start, goal)

            for neighbor, move_cost in neighbors(current, occupied):
                if occupied[neighbor[1]][neighbor[0]]:
                    continue

                parent = came_from[current]
                if (
                    self.use_theta_star
                    and parent != current
                    and self._line_of_sight(parent, neighbor, occupied)
                ):
                    candidate_parent = parent
                    new_cost = cost_so_far[parent] + self._grid_distance(parent, neighbor)
                else:
                    candidate_parent = current
                    new_cost = cost_so_far[current] + move_cost

                if new_cost < cost_so_far.get(neighbor, float("inf")):
                    cost_so_far[neighbor] = new_cost
                    came_from[neighbor] = candidate_parent
                    priority = new_cost + self._grid_distance(neighbor, goal)
                    heapq.heappush(open_heap, (priority, neighbor))

        return []

    def _neighbors4(self, cell, occupied):
        del occupied
        x, y = cell
        for dx, dy in ((1, 0), (-1, 0), (0, 1), (0, -1)):
            nx, ny = x + dx, y + dy
            if 0 <= nx < self.width and 0 <= ny < self.height:
                yield (nx, ny), 1.0

    def _neighbors8(self, cell, occupied):
        x, y = cell
        for dx in (-1, 0, 1):
            for dy in (-1, 0, 1):
                if dx == 0 and dy == 0:
                    continue
                nx, ny = x + dx, y + dy
                if 0 <= nx < self.width and 0 <= ny < self.height:
                    if dx != 0 and dy != 0:
                        # Avoid cutting through the corner of two inflated cells.
                        if occupied[y][nx] or occupied[ny][x]:
                            continue
                    yield (nx, ny), math.hypot(dx, dy)

    def _reconstruct_path(self, came_from, start, goal):
        current = goal
        path = [current]
        while current != start:
            current = came_from[current]
            path.append(current)
        path.reverse()
        return path

    def _line_of_sight(self, start, goal, occupied):
        x0, y0 = start
        x1, y1 = goal
        dx = abs(x1 - x0)
        dy = abs(y1 - y0)
        sx = 1 if x0 < x1 else -1
        sy = 1 if y0 < y1 else -1
        err = dx - dy
        x, y = x0, y0

        while True:
            if x < 0 or x >= self.width or y < 0 or y >= self.height:
                return False
            if occupied[y][x]:
                return False
            if x == x1 and y == y1:
                return True
            e2 = 2 * err
            if e2 > -dy:
                err -= dy
                x += sx
            if e2 < dx:
                err += dx
                y += sy

    def _smooth_by_line_of_sight(self, points, occupied):
        if len(points) <= 2:
            return points
        cells = [self._world_to_grid(x, y, clamp=True) for x, y in points]
        for _ in range(self.smooth_iterations):
            smoothed = [points[0]]
            anchor_index = 0
            test_index = 2
            while test_index < len(points):
                if not self._line_of_sight(cells[anchor_index], cells[test_index], occupied):
                    smoothed.append(points[test_index - 1])
                    anchor_index = test_index - 1
                test_index += 1
            smoothed.append(points[-1])
            points = smoothed
            cells = [self._world_to_grid(x, y, clamp=True) for x, y in points]
        return points

    def _densify(self, points):
        if len(points) <= 1:
            return points
        dense = [points[0]]
        step = max(self.resolution, self.path_resolution)
        for target in points[1:]:
            sx, sy = dense[-1]
            tx, ty = target
            distance = math.hypot(tx - sx, ty - sy)
            if distance < 1e-6:
                continue
            steps = max(1, int(math.ceil(distance / step)))
            for idx in range(1, steps + 1):
                ratio = idx / steps
                dense.append((sx + (tx - sx) * ratio, sy + (ty - sy) * ratio))
        return dense

    def _publish_path(self, points):
        msg = NavPath()
        msg.header.frame_id = self.map_frame
        msg.header.stamp = self.get_clock().now().to_msg()
        for index, (x, y) in enumerate(points):
            if index + 1 < len(points):
                nx, ny = points[index + 1]
                yaw = math.atan2(ny - y, nx - x)
            elif index > 0:
                px, py = points[index - 1]
                yaw = math.atan2(y - py, x - px)
            else:
                yaw = 0.0
            msg.poses.append(make_pose(self.map_frame, msg.header.stamp, x, y, yaw))
        self.plan_pub.publish(msg)

    def _current_pose(self):
        if self.latest_odom is None:
            return None
        msg = self.latest_odom
        if msg.header.stamp.sec != 0 or msg.header.stamp.nanosec != 0:
            age = (self.get_clock().now() - Time.from_msg(msg.header.stamp)).nanoseconds * 1e-9
            if age > self.max_odom_age:
                return None

        x = msg.pose.pose.position.x
        y = msg.pose.pose.position.y
        yaw = yaw_from_quaternion(msg.pose.pose.orientation)
        source_frame = msg.header.frame_id or self.map_frame
        if source_frame != self.map_frame:
            return self._transform_pose_2d(x, y, yaw, source_frame)
        return x, y, yaw

    def _transform_pose_msg(self, msg):
        source_frame = msg.header.frame_id or self.map_frame
        x = msg.pose.position.x
        y = msg.pose.position.y
        yaw = yaw_from_quaternion(msg.pose.orientation)
        if source_frame != self.map_frame:
            return self._transform_pose_2d(x, y, yaw, source_frame)
        return x, y, yaw

    def _transform_pose_2d(self, x, y, yaw, source_frame):
        try:
            transform = self.tf_buffer.lookup_transform(
                self.map_frame,
                source_frame,
                Time(),
                timeout=Duration(seconds=self.tf_timeout),
            )
        except TransformException as exc:
            now_ns = self.get_clock().now().nanoseconds
            if now_ns - self.last_tf_warn_ns > 2_000_000_000:
                self.last_tf_warn_ns = now_ns
                self.get_logger().warn(
                    f"TF unavailable {source_frame}->{self.map_frame}: {exc}"
                )
            return None

        t = transform.transform.translation
        q = transform.transform.rotation
        transform_yaw = yaw_from_quaternion(q)
        cos_yaw = math.cos(transform_yaw)
        sin_yaw = math.sin(transform_yaw)
        target_x = t.x + cos_yaw * x - sin_yaw * y
        target_y = t.y + sin_yaw * x + cos_yaw * y
        target_yaw = normalize_angle(transform_yaw + yaw)
        return target_x, target_y, target_yaw

    def _world_to_grid(self, x, y, clamp=False):
        gx = int(math.floor((x - self.map_min_x) / self.resolution))
        gy = int(math.floor((y - self.map_min_y) / self.resolution))
        if clamp:
            gx = min(self.width - 1, max(0, gx))
            gy = min(self.height - 1, max(0, gy))
            return gx, gy
        if gx < 0 or gx >= self.width or gy < 0 or gy >= self.height:
            return None
        return gx, gy

    def _grid_to_world(self, gx, gy):
        return (
            self.map_min_x + (gx + 0.5) * self.resolution,
            self.map_min_y + (gy + 0.5) * self.resolution,
        )

    @staticmethod
    def _grid_distance(a, b):
        return math.hypot(a[0] - b[0], a[1] - b[1])

    def _publish_status(self, text):
        msg = String()
        msg.data = text
        self.status_pub.publish(msg)
        self.get_logger().info(text)


def main(args=None):
    rclpy.init(args=args)
    node = None
    try:
        node = GridAstarThetaPlanner()
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        if node is not None:
            node.destroy_node()
        rclpy.shutdown()


if __name__ == "__main__":
    main()
