"""
person_follower_node.py -- Proportional person-following control loop.

Subscribes:
  /uwb/target             (geometry_msgs/PoseStamped)  -- UWB-triangulated position (primary)
  /person/target          (geometry_msgs/PoseStamped)  -- camera-detected position (secondary)
  /local_costmap/costmap  (nav_msgs/OccupancyGrid)     -- Nav2 local obstacle map

Publishes:
  /cmd_vel                (geometry_msgs/Twist)         -- velocity commands

Sensor fusion
-------------
  When both UWB and camera sources are fresh:
      fused = uwb_weight * uwb + (1 - uwb_weight) * camera
  When only UWB is fresh:   use UWB directly.
  When only camera is fresh: use camera directly.
  When neither is fresh:    STOPPING -> SEARCHING.

  UWB is declared stale after uwb_timeout seconds (default 1.0s).
  Camera is declared stale after lost_timeout seconds (default 2.0s).

State machine
-------------
  FOLLOWING  -- at least one source fresh (last detection < timeout ago)
  STOPPING   -- all sources lost, publishing zero until search_timeout
  SEARCHING  -- all sources lost > search_timeout, slow rotation to re-acquire

Safety wiring
-------------
<<<<<<< HEAD
  * cmd_vel bridge (PR #46) applies ramp + deadman + ESP32 BALANCE AUTONOMOUS mode gate --
=======
  * cmd_vel bridge (PR #46) applies ramp + deadman + ESP32-S3 AUTONOMOUS mode gate --
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)
    this node publishes raw /cmd_vel, the bridge handles hardware safety.
  * follow_enabled param (default True) lets the operator disable the controller
    at runtime: ros2 param set /person_follower follow_enabled false
  * Obstacle override: if Nav2 local costmap shows a lethal cell in the forward
    corridor, forward cmd_vel is zeroed; turning is still allowed.

Usage
-----
  ros2 launch saltybot_follower person_follower.launch.py
  ros2 launch saltybot_follower person_follower.launch.py follow_distance:=1.2
"""

import math
import time

import rclpy
from rclpy.node import Node
from geometry_msgs.msg import PoseStamped, Twist
from nav_msgs.msg import OccupancyGrid

# -- State constants -----------------------------------------------------------
_FOLLOWING = "following"
_STOPPING  = "stopping"   # all sources lost < search_timeout, publishing zero
_SEARCHING = "searching"  # all sources lost > search_timeout, slow rotation


# -- Pure helper functions (also exercised by unit tests) ----------------------

def clamp(v: float, lo: float, hi: float) -> float:
    return max(lo, min(hi, v))


def compute_distance(px: float, py: float) -> float:
    return math.sqrt(px * px + py * py)


def compute_bearing(px: float, py: float) -> float:
    """Bearing to point (px, py) in base_link: positive = person to the left."""
    return math.atan2(py, px)


def fuse_targets(
    uwb_x: float,
    uwb_y: float,
    uwb_fresh: bool,
    cam_x: float,
    cam_y: float,
    cam_fresh: bool,
    uwb_weight: float,
) -> tuple:
    """
    Fuse UWB (primary) and camera (secondary) position estimates.

    Returns (fused_x, fused_y, any_fresh).

    Fusion rules:
      Both fresh  -> weighted blend: uwb_weight * uwb + (1-uwb_weight) * camera
      UWB only    -> UWB position, full weight
      Camera only -> camera position, full weight
      Neither     -> (0, 0, False)
    """
    if uwb_fresh and cam_fresh:
        w = clamp(uwb_weight, 0.0, 1.0)
        return (
            w * uwb_x + (1.0 - w) * cam_x,
            w * uwb_y + (1.0 - w) * cam_y,
            True,
        )
    if uwb_fresh:
        return uwb_x, uwb_y, True
    if cam_fresh:
        return cam_x, cam_y, True
    return 0.0, 0.0, False


def compute_follow_cmd(
    px: float,
    py: float,
    follow_distance: float,
    dead_zone: float,
    kp_linear: float,
    kp_angular: float,
    max_linear_vel: float,
    max_angular_vel: float,
    obstacle_ahead: bool = False,
) -> tuple:
    """
    Pure proportional follow controller.

    Returns (linear_x, angular_z) as floats.

    linear_x:
      Proportional to signed distance error (distance - follow_distance).
      Zero inside the dead zone. Clamped to +-max_linear_vel.
      Zeroed (not reversed) if obstacle_ahead and moving forward.

    angular_z:
      Proportional to bearing. Always corrects even when inside the dead zone,
      so the robot stays pointed at the person while holding distance.
      Clamped to +-max_angular_vel.
    """
    distance = compute_distance(px, py)
    bearing  = compute_bearing(px, py)

    dist_error = distance - follow_distance
    if abs(dist_error) > dead_zone:
        linear_x = clamp(kp_linear * dist_error, -max_linear_vel, max_linear_vel)
    else:
        linear_x = 0.0

    if obstacle_ahead and linear_x > 0.0:
        linear_x = 0.0

    angular_z = clamp(kp_angular * bearing, -max_angular_vel, max_angular_vel)

    return linear_x, angular_z


def check_obstacle_in_costmap(
    costmap_data: list,
    width: int,
    height: int,
    resolution: float,
    check_dist: float,
    robot_half_width: float,
    threshold: int,
) -> bool:
    """
    Check for obstacle cells in the robot's forward corridor.

    Assumes a Nav2 rolling-window local costmap (robot ~= centre cell).
    Forward = +col direction (robot +X).

    costmap_data : flat row-major int8 list (from OccupancyGrid.data)
    width, height: costmap dimensions in cells
    resolution   : metres per cell
    check_dist   : how far ahead to look (metres)
    robot_half_width: corridor half-width (metres)
    threshold    : costmap value >= this is treated as obstacle (Nav2 lethal=99)
    """
    if resolution <= 0.0:
        return False

    cx = width  // 2
    cy = height // 2

    n_fwd  = int(check_dist / resolution)
    n_side = int(robot_half_width / resolution)

    for dx in range(1, n_fwd + 1):
        col = cx + dx
        for dy in range(-n_side, n_side + 1):
            row = cy + dy
            if row < 0 or row >= height or col < 0 or col >= width:
                continue
            idx = row * width + col
            if idx < len(costmap_data) and costmap_data[idx] >= threshold:
                return True
    return False


# -- ROS2 Node ----------------------------------------------------------------

class PersonFollowerNode(Node):

    def __init__(self):
        super().__init__("person_follower")

        # -- Parameters --------------------------------------------------------
        self.declare_parameter("follow_distance",      1.5)
        self.declare_parameter("dead_zone",            0.3)
        self.declare_parameter("kp_linear",            0.6)
        self.declare_parameter("kp_angular",           1.2)
        self.declare_parameter("max_linear_vel",       0.5)
        self.declare_parameter("max_angular_vel",      1.0)
        self.declare_parameter("lost_timeout",         2.0)
        self.declare_parameter("search_timeout",       5.0)
        self.declare_parameter("search_angular_vel",   0.3)
        self.declare_parameter("obstacle_check_dist",  0.5)
        self.declare_parameter("obstacle_threshold",   70)
        self.declare_parameter("robot_width",          0.3)
        self.declare_parameter("control_rate",         20.0)
        self.declare_parameter("follow_enabled",       True)
        # UWB fusion
        self.declare_parameter("uwb_weight",           0.7)
        self.declare_parameter("uwb_timeout",          1.0)

        self._p = {}
        self._reload_params()

        # -- State -------------------------------------------------------------
        self._state          = _STOPPING
        self._last_target_t  = 0.0    # camera; monotonic; 0 = never received
        self._person_x       = 0.0
        self._person_y       = 0.0
        self._last_uwb_t     = 0.0    # UWB; monotonic; 0 = never received
        self._uwb_x          = 0.0
        self._uwb_y          = 0.0
        self._obstacle_ahead = False

        # -- Subscriptions -----------------------------------------------------
        self.create_subscription(
            PoseStamped, "/uwb/target",    self._uwb_cb,    10)
        self.create_subscription(
            PoseStamped, "/person/target", self._target_cb, 10)
        self.create_subscription(
            OccupancyGrid, "/local_costmap/costmap", self._costmap_cb, 1)

        # -- Publisher ---------------------------------------------------------
        self._cmd_pub = self.create_publisher(Twist, "/cmd_vel", 10)

        # -- Control timer -----------------------------------------------------
        self._timer = self.create_timer(
            1.0 / self._p["control_rate"], self._control_cb)

        self.get_logger().info(
            "PersonFollower ready  "
            f"follow_dist={self._p['follow_distance']}m "
            f"dead_zone=+/-{self._p['dead_zone']}m "
            f"max_vel={self._p['max_linear_vel']}m/s "
            f"uwb_weight={self._p['uwb_weight']} "
            f"lost_timeout={self._p['lost_timeout']}s"
        )

    # -- Parameter helper ------------------------------------------------------

    def _reload_params(self):
        self._p = {
            "follow_distance":     self.get_parameter("follow_distance").value,
            "dead_zone":           self.get_parameter("dead_zone").value,
            "kp_linear":           self.get_parameter("kp_linear").value,
            "kp_angular":          self.get_parameter("kp_angular").value,
            "max_linear_vel":      self.get_parameter("max_linear_vel").value,
            "max_angular_vel":     self.get_parameter("max_angular_vel").value,
            "lost_timeout":        self.get_parameter("lost_timeout").value,
            "search_timeout":      self.get_parameter("search_timeout").value,
            "search_angular_vel":  self.get_parameter("search_angular_vel").value,
            "obstacle_check_dist": self.get_parameter("obstacle_check_dist").value,
            "obstacle_threshold":  self.get_parameter("obstacle_threshold").value,
            "robot_width":         self.get_parameter("robot_width").value,
            "control_rate":        self.get_parameter("control_rate").value,
            "follow_enabled":      self.get_parameter("follow_enabled").value,
            "uwb_weight":          self.get_parameter("uwb_weight").value,
            "uwb_timeout":         self.get_parameter("uwb_timeout").value,
        }

    # -- Callbacks -------------------------------------------------------------

    def _uwb_cb(self, msg):
        """UWB-triangulated person position (primary source)."""
        self._uwb_x      = msg.pose.position.x
        self._uwb_y      = msg.pose.position.y
        self._last_uwb_t = time.monotonic()
        self._state      = _FOLLOWING

    def _target_cb(self, msg):
        """Camera-detected person position (secondary source)."""
        self._person_x      = msg.pose.position.x
        self._person_y      = msg.pose.position.y
        self._last_target_t = time.monotonic()
        self._state         = _FOLLOWING

    def _costmap_cb(self, msg):
        self._obstacle_ahead = check_obstacle_in_costmap(
            costmap_data=list(msg.data),
            width=msg.info.width,
            height=msg.info.height,
            resolution=msg.info.resolution,
            check_dist=self._p["obstacle_check_dist"],
            robot_half_width=self._p["robot_width"] / 2.0,
            threshold=self._p["obstacle_threshold"],
        )

    # -- Control loop ----------------------------------------------------------

    def _control_cb(self):
        self._reload_params()

        twist = Twist()

        if not self._p["follow_enabled"]:
            self._cmd_pub.publish(twist)
            return

        now = time.monotonic()

        uwb_fresh = (
            self._last_uwb_t > 0.0 and
            (now - self._last_uwb_t) <= self._p["uwb_timeout"]
        )
        cam_fresh = (
            self._last_target_t > 0.0 and
            (now - self._last_target_t) <= self._p["lost_timeout"]
        )

        any_fresh = uwb_fresh or cam_fresh
        if self._state == _FOLLOWING:
            if not any_fresh:
                uwb_stale = (now - self._last_uwb_t) if self._last_uwb_t > 0.0 else 1e9
                cam_stale = (now - self._last_target_t) if self._last_target_t > 0.0 else 1e9
                dt_lost   = min(uwb_stale, cam_stale)
                if dt_lost > self._p["search_timeout"]:
                    self._state = _SEARCHING
                    self.get_logger().warn("Person lost -- entering SEARCHING state")
                else:
                    self._state = _STOPPING
                    self.get_logger().warn("Person lost -- stopping")

        linear_x, angular_z = 0.0, 0.0

        if self._state == _FOLLOWING:
            fx, fy, fresh = fuse_targets(
                uwb_x=self._uwb_x,
                uwb_y=self._uwb_y,
                uwb_fresh=uwb_fresh,
                cam_x=self._person_x,
                cam_y=self._person_y,
                cam_fresh=cam_fresh,
                uwb_weight=self._p["uwb_weight"],
            )
            if fresh:
                linear_x, angular_z = compute_follow_cmd(
                    px=fx,
                    py=fy,
                    follow_distance=self._p["follow_distance"],
                    dead_zone=self._p["dead_zone"],
                    kp_linear=self._p["kp_linear"],
                    kp_angular=self._p["kp_angular"],
                    max_linear_vel=self._p["max_linear_vel"],
                    max_angular_vel=self._p["max_angular_vel"],
                    obstacle_ahead=self._obstacle_ahead,
                )
                if self._obstacle_ahead and linear_x == 0.0:
                    self.get_logger().warn(
                        "Obstacle in path -- forward motion suppressed",
                        throttle_duration_sec=2.0,
                    )

        elif self._state == _SEARCHING:
            angular_z = float(self._p["search_angular_vel"])

        twist.linear.x  = linear_x
        twist.angular.z = angular_z
        self._cmd_pub.publish(twist)


# -- Entry point ---------------------------------------------------------------

def main(args=None):
    rclpy.init(args=args)
    node = PersonFollowerNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.try_shutdown()


if __name__ == "__main__":
    main()
