"""
speed_controller_node.py — Outdoor adaptive speed controller for SaltyBot.

Sits between upstream cmd_vel sources (person_follower, Nav2) and the
cmd_vel_bridge (PR #46), providing:

  • Adaptive speed profiles   (walk / jog / ride) selected from UWB target
    velocity and GPS ground-truth speed
  • Smooth acceleration curves (trapezoidal ramp) per profile, tuned for
    balance stability at speed — more aggressive jerk at low speed, gentler
    at ride speeds (EUC follow-me at 20-30 km/h)
  • Emergency deceleration    when the target stops suddenly (detected via
    UWB range acceleration and target velocity drop)
  • GPS runaway protection    if GPS ground speed significantly overshoots
    the commanded velocity (downhill, slippage)

Data flow:
  /cmd_vel_raw  (Twist)           — uncapped request from person_follower/Nav2
  /gps/vel      (TwistStamped)    — robot ground-truth speed (GPS, 1 Hz)
  /uwb/ranges   (String JSON)     — UWB range to target (for velocity estimate)
       ↓
  SpeedControllerNode
       ↓
  /cmd_vel      (Twist)           — smoothed, profile-capped output → cmd_vel_bridge

Profiles
────────
  walk : ≤ 1.5 m/s  (person walking  0–2 m/s)
  jog  : ≤ 3.0 m/s  (person jogging  2–5 m/s)
  ride : ≤ 8.0 m/s  (person on EUC   5–15 m/s)

Profile selection uses UWB-derived target velocity with hysteresis to avoid
rapid switching. GPS ground speed provides a secondary confirmation and
runaway guard.
"""

import collections
import json
import math
import time

import rclpy
from rclpy.node import Node
from geometry_msgs.msg import Twist, TwistStamped
from std_msgs.msg import String


# ── Pure control functions ────────────────────────────────────────────────────
# Extracted here so unit tests run without a ROS2 runtime.

# Profile name constants
WALK = "walk"
JOG  = "jog"
RIDE = "ride"

_PROFILE_ORDER = [WALK, JOG, RIDE]


def estimate_target_velocity(prev_range: float, curr_range: float,
                              dt: float, robot_gps_speed: float) -> float:
    """
    Estimate absolute target velocity from UWB range derivative + GPS.

    When the robot follows a target:
      range_rate = d(range)/dt  (positive = target moving away)
      target_vel ≈ robot_gps_speed + range_rate

    This gives a rough scalar speed suitable for profile classification.
    Returns 0 if dt ≤ 0 or range values are invalid.
    """
    if dt <= 0 or prev_range <= 0 or curr_range <= 0:
        return robot_gps_speed   # fall back to robot's own speed
    range_rate = (curr_range - prev_range) / dt
    return max(0.0, robot_gps_speed + range_rate)


def select_profile(target_vel: float, profiles: dict,
                   current_profile: str,
                   hysteresis_up: int = 5,
                   hysteresis_down: int = 10,
                   _counters: dict = None) -> str:
    """
    Select speed profile with up/down hysteresis.

    target_vel    : estimated target velocity (m/s)
    profiles      : dict of profile_name → {target_vel_min, target_vel_max, ...}
    current_profile: currently active profile
    hysteresis_up : ticks before switching to a faster profile
    hysteresis_down: ticks before switching to a slower profile

    Returns the (possibly unchanged) profile name.

    _counters is an internal mutable dict passed by the caller for state
    (avoids making this a method). Pass {} on first call; reuse thereafter.
    """
    if _counters is None:
        _counters = {}

    # Determine which profile best fits the target velocity
    desired = current_profile
    for name in _PROFILE_ORDER:
        p = profiles[name]
        if target_vel <= p["target_vel_max"]:
            desired = name
            break
    else:
        desired = RIDE   # above all max → stay at highest

    if desired == current_profile:
        _counters.clear()
        return current_profile

    key = desired
    _counters[key] = _counters.get(key, 0) + 1

    idx_desired = _PROFILE_ORDER.index(desired)
    idx_current = _PROFILE_ORDER.index(current_profile)
    threshold   = hysteresis_up if idx_desired > idx_current else hysteresis_down

    if _counters[key] >= threshold:
        _counters.clear()
        return desired
    return current_profile


def apply_accel_limit(current: float, target: float,
                      accel_limit: float, decel_limit: float,
                      dt: float) -> float:
    """
    Apply symmetric trapezoidal acceleration / deceleration limit.

    current     : current velocity (m/s)
    target      : desired velocity (m/s)
    accel_limit : max acceleration (m/s²)
    decel_limit : max deceleration (m/s²) — positive value
    dt          : time step (s)

    Returns new velocity, ramped toward target within the limits.
    """
    delta = target - current
    # Decelerating = magnitude is decreasing (current and delta have opposite signs)
    decelerating = (current > 0 and delta < 0) or (current < 0 and delta > 0)
    limit = decel_limit if decelerating else accel_limit
    if delta > 0:
        step = min(delta,  limit * dt)
    else:
        step = max(delta, -limit * dt)
    return current + step


def limit_velocity(linear_x: float, angular_z: float,
                   max_linear: float, max_angular: float) -> tuple:
    """Hard-clamp linear and angular velocity to profile limits."""
    lin = max(-max_linear,  min(max_linear,  linear_x))
    ang = max(-max_angular, min(max_angular, angular_z))
    return lin, ang


def check_emergency_stop(target_vel: float, target_vel_prev: float,
                         dt: float,
                         stop_vel_threshold: float,
                         hard_decel_threshold: float) -> bool:
    """
    Detect sudden target stop requiring emergency deceleration.

    Triggers if:
      (a) target velocity drops below stop_vel_threshold, OR
      (b) target deceleration exceeds hard_decel_threshold (m/s²)

    Returns True if emergency stop should be engaged.
    """
    if target_vel < stop_vel_threshold:
        return True
    if dt > 0:
        target_decel = (target_vel_prev - target_vel) / dt   # positive = decelerating
        if target_decel > hard_decel_threshold:
            return True
    return False


def check_gps_runaway(gps_speed: float, cmd_vel: float,
                      runaway_factor: float, profile_max: float) -> bool:
    """
    Detect GPS runaway: robot moving significantly faster than commanded.

    Triggers if gps_speed > max(cmd_vel, 0.5) * runaway_factor
    AND gps_speed > profile_max (not a sensor glitch on a stationary robot).
    """
    reference = max(abs(cmd_vel), 0.5)
    return gps_speed > reference * runaway_factor and gps_speed > profile_max * 0.5


def emergency_decel_step(current_vel: float, decel_limit: float, dt: float) -> float:
    """Decelerate toward zero at emergency_decel_limit. Returns new velocity."""
    if current_vel > 0:
        return max(0.0, current_vel - decel_limit * dt)
    if current_vel < 0:
        return min(0.0, current_vel + decel_limit * dt)
    return 0.0


# ── Default profiles (overridable via YAML) ──────────────────────────────────

DEFAULT_PROFILES = {
    WALK: {
        "max_linear_vel":   1.5,
        "max_angular_vel":  1.2,
        "accel_limit":      0.6,    # m/s²
        "decel_limit":      1.0,    # m/s²
        "target_vel_min":   0.0,
        "target_vel_max":   2.0,
    },
    JOG: {
        "max_linear_vel":   3.0,
        "max_angular_vel":  1.5,
        "accel_limit":      0.5,
        "decel_limit":      0.8,
        "target_vel_min":   2.0,
        "target_vel_max":   5.0,
    },
    RIDE: {
        "max_linear_vel":   8.0,
        "max_angular_vel":  1.5,    # conservative for balance at speed
        "accel_limit":      0.3,    # very gentle — balance needs time at 8 m/s
        "decel_limit":      0.5,
        "target_vel_min":   5.0,
        "target_vel_max":   15.0,
    },
}


# ── ROS2 Node ─────────────────────────────────────────────────────────────────

class SpeedControllerNode(Node):

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

        # ── Parameters ────────────────────────────────────────────────────────
        self.declare_parameter("control_rate",           50.0)
        self.declare_parameter("input_topic",            "/cmd_vel_raw")
        self.declare_parameter("output_topic",           "/cmd_vel")
        self.declare_parameter("initial_profile",        WALK)

        # Profile selection
        self.declare_parameter("hysteresis_up_ticks",   5)
        self.declare_parameter("hysteresis_down_ticks", 15)

        # Emergency stop
        self.declare_parameter("emergency_decel_mps2",       2.0)
        self.declare_parameter("stop_vel_threshold",         0.4)    # m/s
        self.declare_parameter("hard_stop_decel_threshold",  3.0)    # m/s²

        # GPS runaway
        self.declare_parameter("gps_runaway_factor",    1.5)   # 1.5× commanded
        self.declare_parameter("gps_timeout",           3.0)   # s before ignoring GPS

        # UWB range
        self.declare_parameter("uwb_range_field",       "range")
        self.declare_parameter("uwb_timeout",           2.0)

        # Per-profile overrides (flat params, merged into DEFAULT_PROFILES)
        for profile in _PROFILE_ORDER:
            self.declare_parameter(f"{profile}.max_linear_vel",
                                   DEFAULT_PROFILES[profile]["max_linear_vel"])
            self.declare_parameter(f"{profile}.max_angular_vel",
                                   DEFAULT_PROFILES[profile]["max_angular_vel"])
            self.declare_parameter(f"{profile}.accel_limit",
                                   DEFAULT_PROFILES[profile]["accel_limit"])
            self.declare_parameter(f"{profile}.decel_limit",
                                   DEFAULT_PROFILES[profile]["decel_limit"])
            self.declare_parameter(f"{profile}.target_vel_max",
                                   DEFAULT_PROFILES[profile]["target_vel_max"])

        self._profiles = self._load_profiles()
        self._p = {
            "control_rate":              self.get_parameter("control_rate").value,
            "hysteresis_up":             self.get_parameter("hysteresis_up_ticks").value,
            "hysteresis_down":           self.get_parameter("hysteresis_down_ticks").value,
            "emergency_decel":           self.get_parameter("emergency_decel_mps2").value,
            "stop_vel_threshold":        self.get_parameter("stop_vel_threshold").value,
            "hard_stop_decel_threshold": self.get_parameter("hard_stop_decel_threshold").value,
            "gps_runaway_factor":        self.get_parameter("gps_runaway_factor").value,
            "gps_timeout":               self.get_parameter("gps_timeout").value,
            "uwb_range_field":           self.get_parameter("uwb_range_field").value,
            "uwb_timeout":               self.get_parameter("uwb_timeout").value,
        }
        dt = 1.0 / self._p["control_rate"]

        # ── State ─────────────────────────────────────────────────────────────
        self._profile        = self.get_parameter("initial_profile").value
        self._profile_counters: dict = {}
        self._current_linear = 0.0   # smoothed output linear.x
        self._current_angular= 0.0   # smoothed output angular.z
        self._request_linear = 0.0   # latest /cmd_vel_raw linear.x
        self._request_angular= 0.0   # latest /cmd_vel_raw angular.z
        self._gps_speed      = 0.0   # magnitude from /gps/vel
        self._gps_stamp      = 0.0   # monotonic time of last GPS
        self._uwb_range      = 0.0   # latest range (m)
        self._uwb_range_prev = 0.0
        self._uwb_stamp      = 0.0
        self._uwb_stamp_prev = 0.0
        self._target_vel     = 0.0
        self._target_vel_prev= 0.0
        self._emergency      = False
        self._dt             = dt

        # History for runaway filter (last 3 GPS samples)
        self._gps_history: collections.deque = collections.deque(maxlen=3)

        # ── Subscriptions ─────────────────────────────────────────────────────
        self.create_subscription(
            Twist, self.get_parameter("input_topic").value,
            self._cmd_raw_cb, 10)
        self.create_subscription(
            TwistStamped, "/gps/vel", self._gps_cb, 10)
        self.create_subscription(
            String, "/uwb/ranges", self._uwb_cb, 10)

        # ── Publisher ──────────────────────────────────────────────────────────
        self._cmd_pub = self.create_publisher(
            Twist, self.get_parameter("output_topic").value, 10)
        self._profile_pub = self.create_publisher(
            String, "/speed_controller/profile", 10)

        # ── Control timer ──────────────────────────────────────────────────────
        self.create_timer(dt, self._control_cb)

        self.get_logger().info(
            f"SpeedController ready  profile={self._profile} "
            f"rate={self._p['control_rate']}Hz  "
            f"walk≤{self._profiles[WALK]['max_linear_vel']}m/s  "
            f"jog≤{self._profiles[JOG]['max_linear_vel']}m/s  "
            f"ride≤{self._profiles[RIDE]['max_linear_vel']}m/s")

    def _load_profiles(self) -> dict:
        """Merge ROS2 params into DEFAULT_PROFILES."""
        import copy
        profiles = copy.deepcopy(DEFAULT_PROFILES)
        for name in _PROFILE_ORDER:
            for field in ("max_linear_vel", "max_angular_vel", "accel_limit",
                          "decel_limit", "target_vel_max"):
                val = self.get_parameter(f"{name}.{field}").value
                profiles[name][field] = val
        return profiles

    # ── Callbacks ─────────────────────────────────────────────────────────────

    def _cmd_raw_cb(self, msg: Twist):
        self._request_linear  = msg.linear.x
        self._request_angular = msg.angular.z

    def _gps_cb(self, msg: TwistStamped):
        vx = msg.twist.linear.x
        vy = msg.twist.linear.y
        self._gps_speed = math.sqrt(vx * vx + vy * vy)
        self._gps_stamp = time.monotonic()
        self._gps_history.append(self._gps_speed)

    def _uwb_cb(self, msg: String):
        """Parse JSON range message, update range estimate."""
        try:
            data = json.loads(msg.data)
            field = self._p["uwb_range_field"]
            r = float(data.get(field, data.get("distance", 0.0)))
            if r > 0:
                self._uwb_range_prev = self._uwb_range
                self._uwb_stamp_prev = self._uwb_stamp
                self._uwb_range      = r
                self._uwb_stamp      = time.monotonic()
        except (json.JSONDecodeError, ValueError, KeyError):
            pass

    # ── Control loop ──────────────────────────────────────────────────────────

    def _control_cb(self):
        now = time.monotonic()
        dt  = self._dt

        # ── 1. Estimate target velocity from UWB range derivative + GPS ──────
        uwb_fresh = (now - self._uwb_stamp) < self._p["uwb_timeout"]
        gps_fresh = (now - self._gps_stamp) < self._p["gps_timeout"]

        if uwb_fresh and self._uwb_stamp_prev > 0:
            uwb_dt = self._uwb_stamp - self._uwb_stamp_prev
            robot_speed = self._gps_speed if gps_fresh else self._current_linear
            self._target_vel_prev = self._target_vel
            self._target_vel = estimate_target_velocity(
                self._uwb_range_prev, self._uwb_range, uwb_dt, robot_speed)
        elif gps_fresh:
            # No UWB — use GPS speed as proxy for target velocity
            self._target_vel_prev = self._target_vel
            self._target_vel = self._gps_speed
        # else: hold previous estimate

        # ── 2. Profile selection ──────────────────────────────────────────────
        self._profile = select_profile(
            self._target_vel, self._profiles, self._profile,
            self._p["hysteresis_up"], self._p["hysteresis_down"],
            self._profile_counters,
        )
        prof = self._profiles[self._profile]

        # ── 3. Emergency stop detection ───────────────────────────────────────
        emerg = check_emergency_stop(
            self._target_vel, self._target_vel_prev, dt,
            self._p["stop_vel_threshold"],
            self._p["hard_stop_decel_threshold"],
        )
        # GPS runaway check (use smoothed GPS)
        if gps_fresh:
            gps_smooth = sum(self._gps_history) / len(self._gps_history)
            if check_gps_runaway(
                    gps_smooth, self._current_linear,
                    self._p["gps_runaway_factor"],
                    prof["max_linear_vel"]):
                emerg = True
                self.get_logger().warn(
                    f"GPS runaway: gps={gps_smooth:.2f} "
                    f"cmd={self._current_linear:.2f} m/s — emergency decel",
                    throttle_duration_sec=2.0)

        # Latch emergency; clear when speed reaches zero
        if emerg and not self._emergency:
            self.get_logger().warn(
                f"Emergency stop: target_vel={self._target_vel:.2f} m/s "
                f"(profile={self._profile})")
        self._emergency = emerg or (
            self._emergency and abs(self._current_linear) > 0.05)

        # ── 4. Compute target command ─────────────────────────────────────────
        if self._emergency:
            # Override request: decelerate to zero at emergency rate
            target_lin = 0.0
            target_ang = 0.0
            new_lin = emergency_decel_step(
                self._current_linear, self._p["emergency_decel"], dt)
            new_ang = emergency_decel_step(
                self._current_angular, self._p["emergency_decel"] * 2.0, dt)
        else:
            # Apply profile caps to request
            target_lin, target_ang = limit_velocity(
                self._request_linear, self._request_angular,
                prof["max_linear_vel"], prof["max_angular_vel"])
            # Apply smooth acceleration / deceleration
            new_lin = apply_accel_limit(
                self._current_linear,  target_lin,
                prof["accel_limit"], prof["decel_limit"], dt)
            new_ang = apply_accel_limit(
                self._current_angular, target_ang,
                prof["accel_limit"] * 2.0,   # angular can be snappier
                prof["decel_limit"] * 2.0, dt)

        self._current_linear  = new_lin
        self._current_angular = new_ang

        # ── 5. Publish ────────────────────────────────────────────────────────
        out = Twist()
        out.linear.x  = self._current_linear
        out.angular.z = self._current_angular
        self._cmd_pub.publish(out)

        # Profile status (for dashboard / logging)
        status = String()
        status.data = json.dumps({
            "profile":      self._profile,
            "max_vel":      prof["max_linear_vel"],
            "target_vel":   round(self._target_vel, 2),
            "gps_speed":    round(self._gps_speed, 2),
            "cmd_linear":   round(new_lin, 3),
            "emergency":    self._emergency,
        })
        self._profile_pub.publish(status)


# ── Entry point ───────────────────────────────────────────────────────────────

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


if __name__ == "__main__":
    main()
