"""
adaptive_pid_node.py — Adaptive balance PID controller (Issue #136).

Overview
────────
  Subscribes to IMU and motor-current topics, runs a PID balance loop at
  `control_rate` Hz, and automatically adjusts gains based on estimated
  payload weight.

Weight estimation
─────────────────
  Motor current at steady state (low tilt, no cmd_vel) is sampled into a
  rolling window.  When enough samples accumulate:

      weight_kg = max(0, (I_avg − I_idle) / k_per_kg)

  If the estimate shifts by ≥ weight_change_threshold_kg, gains are
  re-scheduled immediately (subject to blending).

Gain scheduling
───────────────
  Three weight classes (empty / light / heavy) each have a (kp, ki, kd)
  table entry.  Active gains blend exponentially toward the target class at
  rate gain_blend_alpha per tick.

Safety
──────
  • All gains clamped to [kp_min..kp_max], [ki_min..ki_max], [kd_min..kd_max]
  • InstabilityDetector monitors tilt amplitude and sign reversals; on trigger:
      – gains reverted to table defaults (no blending)
      – PID integrator reset
      – WARN logged
  • Watchdog: if /saltybot/imu goes silent > imu_timeout_s → effort = 0

Dynamic reconfigure
───────────────────
  Set `override_enabled=true` plus override_kp / override_ki / override_kd
  via `ros2 param set` to bypass scheduling.  Set `override_enabled=false`
  to return to auto-scheduling.

Calibration service
───────────────────
  Call `/saltybot/calibrate_balance` (std_srvs/Trigger) to start the
  calibration routine.  The robot rocks gently for `calibration_cycles /
  calibration_frequency_hz` seconds, then settles.  The weight estimate is
  updated and logged on completion.

Subscribes
──────────
  /saltybot/imu            sensor_msgs/Imu    — tilt angle (pitch from quaternion)
  /saltybot/motor_current  std_msgs/Float32   — motor current (A)

Publishes
─────────
  /saltybot/balance_effort  std_msgs/Float32  — PID output (balance correction)
  /saltybot/weight_estimate std_msgs/Float32  — current weight estimate (kg)
  /saltybot/pid_state       std_msgs/String   — JSON: gains, class, mode, flags

Services
────────
  /saltybot/calibrate_balance  std_srvs/Trigger

Parameters  (see config/adaptive_pid_params.yaml for defaults)
──────────────────────────────────────────────────────────────
  control_rate, balance_setpoint_rad,
  kp_empty/light/heavy, ki_*, kd_*,
  kp_min/max, ki_min/max, kd_min/max,
  weight_empty_max_kg, weight_light_max_kg,
  idle_current_a, current_per_kg_a,
  weight_window_s, weight_change_threshold_kg,
  steady_tilt_threshold_rad, gain_blend_alpha,
  instability_tilt_threshold_rad, instability_oscillation_count,
  instability_window_size, imu_timeout_s,
  override_enabled, override_kp, override_ki, override_kd,
  calibration_amplitude_rad, calibration_frequency_hz,
  calibration_cycles, calibration_settle_s,
  output_min, output_max, i_clamp
"""

import json
import math
import time

import rclpy
from rclpy.node import Node
from rclpy.qos import HistoryPolicy, QoSProfile, ReliabilityPolicy

from sensor_msgs.msg import Imu
from std_msgs.msg import Float32, String
from std_srvs.srv import Trigger

from saltybot_adaptive_pid.pid_controller import (
    GainScheduler,
    GainSet,
    InstabilityDetector,
    PIDController,
)
from saltybot_adaptive_pid.weight_estimator import (
    CalibrationRoutine,
    CalibrationState,
    WeightEstimator,
)


def _quaternion_to_pitch(qx: float, qy: float, qz: float, qw: float) -> float:
    """Extract pitch (rotation about Y) from a unit quaternion."""
    sinp = 2.0 * (qw * qy - qz * qx)
    return math.asin(max(-1.0, min(1.0, sinp)))


class AdaptivePIDNode(Node):

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

        # ── Parameters ───────────────────────────────────────────────────────
        self._declare_all_parameters()
        p = self._load_params()

        # ── Sub-systems ──────────────────────────────────────────────────────
        self._pid = PIDController(
            kp=p["kp_empty"], ki=p["ki_empty"], kd=p["kd_empty"],
            output_min=p["output_min"], output_max=p["output_max"],
            i_clamp=p["i_clamp"],
        )

        self._scheduler = GainScheduler(
            gain_table={
                "empty": GainSet(p["kp_empty"], p["ki_empty"], p["kd_empty"]),
                "light": GainSet(p["kp_light"], p["ki_light"], p["kd_light"]),
                "heavy": GainSet(p["kp_heavy"], p["ki_heavy"], p["kd_heavy"]),
            },
            weight_thresholds={
                "empty_max": p["weight_empty_max_kg"],
                "light_max": p["weight_light_max_kg"],
            },
            blend_alpha=p["gain_blend_alpha"],
            kp_min=p["kp_min"], kp_max=p["kp_max"],
            ki_min=p["ki_min"], ki_max=p["ki_max"],
            kd_min=p["kd_min"], kd_max=p["kd_max"],
        )

        self._weight_estimator = WeightEstimator(
            idle_current_a=p["idle_current_a"],
            current_per_kg_a=p["current_per_kg_a"],
            window_s=p["weight_window_s"],
            control_rate_hz=p["control_rate"],
            steady_tilt_threshold_rad=p["steady_tilt_threshold_rad"],
            change_threshold_kg=p["weight_change_threshold_kg"],
        )

        self._instability_detector = InstabilityDetector(
            tilt_threshold_rad=p["instability_tilt_threshold_rad"],
            oscillation_count=p["instability_oscillation_count"],
            window_size=p["instability_window_size"],
        )

        self._calibration = CalibrationRoutine(
            amplitude_rad=p["calibration_amplitude_rad"],
            frequency_hz=p["calibration_frequency_hz"],
            rock_cycles=p["calibration_cycles"],
            settle_s=p["calibration_settle_s"],
            settle_tilt_threshold_rad=p["steady_tilt_threshold_rad"],
            idle_current_a=p["idle_current_a"],
            current_per_kg_a=p["current_per_kg_a"],
        )

        # ── State ────────────────────────────────────────────────────────────
        self._tilt_rad: float = 0.0
        self._current_a: float = 0.0
        self._last_imu_t: float = 0.0
        self._weight_kg: float = 0.0
        self._unstable: bool = False
        self._last_ctrl_t: float = time.monotonic()

        # ── QoS ──────────────────────────────────────────────────────────────
        reliable = QoSProfile(
            reliability=ReliabilityPolicy.RELIABLE,
            history=HistoryPolicy.KEEP_LAST,
            depth=10,
        )
        best_effort = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            history=HistoryPolicy.KEEP_LAST,
            depth=1,
        )

        # ── Subscriptions ────────────────────────────────────────────────────
        self.create_subscription(Imu,     "/saltybot/imu",           self._imu_cb,     best_effort)
        self.create_subscription(Float32, "/saltybot/motor_current",  self._current_cb, reliable)

        # ── Publishers ───────────────────────────────────────────────────────
        self._effort_pub  = self.create_publisher(Float32, "/saltybot/balance_effort",  reliable)
        self._weight_pub  = self.create_publisher(Float32, "/saltybot/weight_estimate", reliable)
        self._state_pub   = self.create_publisher(String,  "/saltybot/pid_state",       reliable)

        # ── Service ──────────────────────────────────────────────────────────
        self.create_service(Trigger, "/saltybot/calibrate_balance", self._calibrate_srv)

        # ── Control timer ────────────────────────────────────────────────────
        rate = p["control_rate"]
        self._ctrl_timer = self.create_timer(1.0 / rate, self._control_cb)

        self.get_logger().info(
            f"AdaptivePIDNode ready  "
            f"rate={rate}Hz  "
            f"kp_empty={p['kp_empty']}  ki_empty={p['ki_empty']}  kd_empty={p['kd_empty']}"
        )

    # ── Parameter declaration / load ─────────────────────────────────────────

    def _declare_all_parameters(self) -> None:
        self.declare_parameter("control_rate",                    100.0)
        self.declare_parameter("balance_setpoint_rad",            0.0)
        # Gains per class
        self.declare_parameter("kp_empty",   15.0)
        self.declare_parameter("ki_empty",    0.5)
        self.declare_parameter("kd_empty",    1.5)
        self.declare_parameter("kp_light",   18.0)
        self.declare_parameter("ki_light",    0.6)
        self.declare_parameter("kd_light",    2.0)
        self.declare_parameter("kp_heavy",   22.0)
        self.declare_parameter("ki_heavy",    0.8)
        self.declare_parameter("kd_heavy",    2.5)
        # Safety bounds
        self.declare_parameter("kp_min",      5.0)
        self.declare_parameter("kp_max",     40.0)
        self.declare_parameter("ki_min",      0.0)
        self.declare_parameter("ki_max",      5.0)
        self.declare_parameter("kd_min",      0.0)
        self.declare_parameter("kd_max",     10.0)
        self.declare_parameter("output_min", -50.0)
        self.declare_parameter("output_max",  50.0)
        self.declare_parameter("i_clamp",    10.0)
        # Weight classification
        self.declare_parameter("weight_empty_max_kg",    0.5)
        self.declare_parameter("weight_light_max_kg",    2.0)
        # Weight estimation
        self.declare_parameter("idle_current_a",         0.5)
        self.declare_parameter("current_per_kg_a",       0.3)
        self.declare_parameter("weight_window_s",        3.0)
        self.declare_parameter("weight_change_threshold_kg", 0.3)
        self.declare_parameter("steady_tilt_threshold_rad",  0.05)
        # Gain blending
        self.declare_parameter("gain_blend_alpha",       0.05)
        # Instability detection
        self.declare_parameter("instability_tilt_threshold_rad", 0.4)
        self.declare_parameter("instability_oscillation_count",   6)
        self.declare_parameter("instability_window_size",        100)
        # Watchdog
        self.declare_parameter("imu_timeout_s",          0.5)
        # Dynamic reconfigure / manual override
        self.declare_parameter("override_enabled",       False)
        self.declare_parameter("override_kp",            15.0)
        self.declare_parameter("override_ki",             0.5)
        self.declare_parameter("override_kd",             1.5)
        # Calibration
        self.declare_parameter("calibration_amplitude_rad",  0.05)
        self.declare_parameter("calibration_frequency_hz",   0.5)
        self.declare_parameter("calibration_cycles",         3)
        self.declare_parameter("calibration_settle_s",       2.0)

    def _load_params(self) -> dict:
        g = self.get_parameter
        return {k: g(k).value for k in [
            "control_rate", "balance_setpoint_rad",
            "kp_empty", "ki_empty", "kd_empty",
            "kp_light", "ki_light", "kd_light",
            "kp_heavy", "ki_heavy", "kd_heavy",
            "kp_min", "kp_max", "ki_min", "ki_max", "kd_min", "kd_max",
            "output_min", "output_max", "i_clamp",
            "weight_empty_max_kg", "weight_light_max_kg",
            "idle_current_a", "current_per_kg_a",
            "weight_window_s", "weight_change_threshold_kg",
            "steady_tilt_threshold_rad", "gain_blend_alpha",
            "instability_tilt_threshold_rad", "instability_oscillation_count",
            "instability_window_size", "imu_timeout_s",
            "override_enabled", "override_kp", "override_ki", "override_kd",
            "calibration_amplitude_rad", "calibration_frequency_hz",
            "calibration_cycles", "calibration_settle_s",
        ]}

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

    def _imu_cb(self, msg: Imu) -> None:
        q = msg.orientation
        self._tilt_rad = _quaternion_to_pitch(q.x, q.y, q.z, q.w)
        self._last_imu_t = time.monotonic()

    def _current_cb(self, msg: Float32) -> None:
        self._current_a = msg.data

    def _calibrate_srv(self, request, response) -> Trigger.Response:
        if self._calibration.is_active:
            response.success = False
            response.message = "Calibration already in progress"
        else:
            self._calibration.start()
            self._pid.reset()
            self._weight_estimator.reset()
            response.success = True
            response.message = "Calibration started"
            self.get_logger().info("Calibration started")
        return response

    # ── 100 Hz control loop ──────────────────────────────────────────────────

    def _control_cb(self) -> None:
        now = time.monotonic()
        dt = now - self._last_ctrl_t
        self._last_ctrl_t = now

        p = self._load_params()

        # ── IMU watchdog ─────────────────────────────────────────────────────
        imu_age = (now - self._last_imu_t) if self._last_imu_t > 0.0 else 1e9
        if imu_age > p["imu_timeout_s"]:
            self._publish_effort(0.0, p, reason="imu_timeout")
            return

        tilt = self._tilt_rad
        current = self._current_a

        # ── Calibration ──────────────────────────────────────────────────────
        cal_offset = self._calibration.step(current, tilt, dt)
        if self._calibration.state == CalibrationState.DONE and \
                self._calibration.weight_estimate is not None:
            cal_wt = self._calibration.weight_estimate
            self._weight_kg = cal_wt
            self._weight_estimator.reset()
            self.get_logger().info(
                f"Calibration done: weight={cal_wt:.2f} kg "
                f"class={self._scheduler.classify(cal_wt)}"
            )

        if self._calibration.state == CalibrationState.FAILED:
            self.get_logger().warn("Calibration failed: insufficient settled samples")

        # ── Dynamic reconfigure / override ────────────────────────────────────
        if p["override_enabled"]:
            if not self._scheduler.override_active:
                self._scheduler.override_gains(
                    p["override_kp"], p["override_ki"], p["override_kd"]
                )
                self.get_logger().info(
                    f"Manual override: kp={p['override_kp']} "
                    f"ki={p['override_ki']} kd={p['override_kd']}"
                )
        else:
            if self._scheduler.override_active:
                self._scheduler.clear_override()
                self.get_logger().info("Manual override cleared, resuming scheduling")

        # ── Weight estimation ─────────────────────────────────────────────────
        if not self._calibration.is_active:
            weight_est, weight_changed = self._weight_estimator.update(current, tilt)
            if weight_est is not None:
                if weight_changed:
                    old_class = self._scheduler.active_class
                    self._weight_kg = weight_est
                    new_class = self._scheduler.classify(weight_est)
                    if old_class != new_class:
                        self.get_logger().info(
                            f"Weight changed: {weight_est:.2f} kg "
                            f"({old_class} → {new_class})"
                        )
                    else:
                        self.get_logger().info(
                            f"Weight updated: {weight_est:.2f} kg class={new_class}"
                        )

        # ── Gain scheduling ──────────────────────────────────────────────────
        gains, wt_class, class_changed = self._scheduler.update(self._weight_kg)

        if class_changed:
            self.get_logger().info(
                f"Gains rescheduled → class={wt_class} "
                f"kp={gains.kp:.3f} ki={gains.ki:.3f} kd={gains.kd:.3f}"
            )

        self._pid.update_gains(gains.kp, gains.ki, gains.kd)

        # ── Instability detection ─────────────────────────────────────────────
        unstable = self._instability_detector.update(tilt)
        if unstable and not self._unstable:
            self._unstable = True
            self.get_logger().warn(
                f"Instability detected (tilt={tilt:.3f} rad) — "
                f"reverting gains to defaults"
            )
            self._scheduler.revert_to_defaults(self._weight_kg)
            reverted = self._scheduler.active_gains
            self._pid.update_gains(reverted.kp, reverted.ki, reverted.kd)
            self._pid.reset()
        elif not unstable:
            self._unstable = False

        # ── PID step ─────────────────────────────────────────────────────────
        setpoint = p["balance_setpoint_rad"] + cal_offset
        error = setpoint - tilt
        effort = self._pid.step(error, dt)

        self._publish_effort(effort, p)

        # ── Weight + state publish ────────────────────────────────────────────
        wt_msg = Float32()
        wt_msg.data = float(self._weight_kg)
        self._weight_pub.publish(wt_msg)

        state_msg = String()
        state_msg.data = json.dumps({
            "kp": round(gains.kp, 4),
            "ki": round(gains.ki, 4),
            "kd": round(gains.kd, 4),
            "weight_class":    wt_class,
            "weight_kg":       round(self._weight_kg, 3),
            "override_active": self._scheduler.override_active,
            "unstable":        self._unstable,
            "calibrating":     self._calibration.is_active,
            "cal_state":       self._calibration.state.value,
            "tilt_rad":        round(tilt, 4),
            "effort":          round(effort, 4),
        })
        self._state_pub.publish(state_msg)

    def _publish_effort(self, effort: float, p: dict, reason: str = "") -> None:
        msg = Float32()
        msg.data = float(effort)
        self._effort_pub.publish(msg)

    # ── Lifecycle ─────────────────────────────────────────────────────────────

    def destroy_node(self) -> None:
        # Zero effort on shutdown
        try:
            msg = Float32()
            msg.data = 0.0
            self._effort_pub.publish(msg)
        except Exception:
            pass
        super().destroy_node()


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

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


if __name__ == "__main__":
    main()
