#!/usr/bin/env python3
"""PID auto-tuner for SaltyBot using Astrom-Hagglund relay feedback method.

Service-triggered relay oscillation measures ultimate gain (Ku) and ultimate period
(Tu), then computes Ziegler-Nichols PID gains (Kp, Ki, Kd). After tuning the computed
gains are written to the FC via JLink binary protocol (JLINK_CMD_PID_SET) and
immediately saved to flash (JLINK_CMD_PID_SAVE). Aborts on >25deg tilt for safety.

Service:
  /saltybot/autotune_pid (std_srvs/Trigger) - Start auto-tuning process

Published topics:
  /saltybot/autotune_info (std_msgs/String) - JSON with tuning progress/results
  /saltybot/autotune_cmd_vel (geometry_msgs/Twist) - Relay control during tuning

Subscribed topics:
  /imu/data (sensor_msgs/Imu) - IMU for tilt safety check
  /saltybot/balance_feedback (std_msgs/Float32) - Balance angle feedback for oscillation
"""

import json
import math
import struct
import time
from enum import Enum
from typing import Optional, Tuple, List

import rclpy
from rclpy.node import Node
from geometry_msgs.msg import Twist
from sensor_msgs.msg import Imu
from std_msgs.msg import Float32, String
from std_srvs.srv import Trigger

try:
    import serial
    _SERIAL_AVAILABLE = True
except ImportError:
    _SERIAL_AVAILABLE = False


# ── JLink binary protocol constants (must match jlink.h) ─────────────────────
JLINK_STX = 0x02
JLINK_ETX = 0x03
JLINK_CMD_PID_SET  = 0x05
JLINK_CMD_PID_SAVE = 0x0A


def _crc16_xmodem(data: bytes) -> int:
    """CRC16-XModem (poly 0x1021, init 0x0000) — matches FC implementation."""
    crc = 0x0000
    for byte in data:
        crc ^= byte << 8
        for _ in range(8):
            if crc & 0x8000:
                crc = (crc << 1) ^ 0x1021
            else:
                crc <<= 1
        crc &= 0xFFFF
    return crc


def _build_jlink_frame(cmd: int, payload: bytes) -> bytes:
    """Build a JLink binary frame: [STX][LEN][CMD][PAYLOAD][CRC_hi][CRC_lo][ETX]."""
    data = bytes([cmd]) + payload
    length = len(data)
    crc = _crc16_xmodem(data)
    return bytes([JLINK_STX, length]) + data + bytes([crc >> 8, crc & 0xFF, JLINK_ETX])


class TuneState(Enum):
    """State machine for auto-tuning process."""
    IDLE = 0
    TUNING = 1
    MEASURING = 2
    COMPLETE = 3
    ABORTED = 4


class PIDAutotuneNode(Node):
    """ROS2 node for automatic PID tuning via relay feedback."""

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

        # Parameters
        self.declare_parameter("relay_magnitude", 0.3)
        self.declare_parameter("tilt_safety_limit", 25.0)
        self.declare_parameter("max_tune_duration", 120.0)
        self.declare_parameter("oscillation_count", 5)
        self.declare_parameter("settle_time", 2.0)
        self.declare_parameter("jlink_serial_port", "/dev/ttyTHS0")
        self.declare_parameter("jlink_baud_rate", 921600)

        self.relay_magnitude   = self.get_parameter("relay_magnitude").value
        self.tilt_safety_limit = self.get_parameter("tilt_safety_limit").value
        self.max_tune_duration = self.get_parameter("max_tune_duration").value
        self.oscillation_count = self.get_parameter("oscillation_count").value
        self.settle_time       = self.get_parameter("settle_time").value
        self._jlink_port       = self.get_parameter("jlink_serial_port").value
        self._jlink_baud       = self.get_parameter("jlink_baud_rate").value

        # State
        self.tune_state = TuneState.IDLE
        self.last_imu_tilt = 0.0
        self.last_feedback = 0.0
        self.feedback_history: List[Tuple[float, float]] = []
        self.peaks: List[Tuple[float, float]] = []
        self.start_time = None

        # Subscriptions
        self.create_subscription(Imu, "/imu/data", self._on_imu, 10)
        self.create_subscription(
            Float32, "/saltybot/balance_feedback", self._on_feedback, 10
        )

        # Publications
        self.pub_cmd_vel = self.create_publisher(
            Twist, "/saltybot/autotune_cmd_vel", 10
        )
        self.pub_info = self.create_publisher(String, "/saltybot/autotune_info", 10)

        # Service
        self.srv_autotune = self.create_service(
            Trigger, "/saltybot/autotune_pid", self._handle_autotune_request
        )

        self.get_logger().info(
            "PID auto-tuner initialized. Service: /saltybot/autotune_pid\n"
            f"Relay magnitude: {self.relay_magnitude}, "
            f"Tilt safety limit: {self.tilt_safety_limit}°, "
            f"Max duration: {self.max_tune_duration}s, "
            f"JLink port: {self._jlink_port}"
        )

    # ── IMU / feedback ────────────────────────────────────────────────────────

    def _on_imu(self, msg: Imu) -> None:
        accel = msg.linear_acceleration
        self.last_imu_tilt = math.degrees(
            math.atan2(accel.x, math.sqrt(accel.y**2 + accel.z**2))
        )

    def _on_feedback(self, msg: Float32) -> None:
        now = self.get_clock().now().nanoseconds / 1e9
        self.last_feedback = msg.data
        if self.tune_state == TuneState.MEASURING:
            self.feedback_history.append((now, msg.data))

    # ── Service handler ───────────────────────────────────────────────────────

    def _handle_autotune_request(
        self, request: Trigger.Request, response: Trigger.Response
    ) -> Trigger.Response:
        if self.tune_state != TuneState.IDLE:
            response.success = False
            response.message = f"Already tuning (state: {self.tune_state.name})"
            return response

        self.get_logger().info("Starting PID auto-tune...")
        self.tune_state = TuneState.TUNING
        self.feedback_history = []
        self.peaks = []
        self.start_time = self.get_clock().now().nanoseconds / 1e9

        self._run_autotuning()

        response.success = self.tune_state == TuneState.COMPLETE
        response.message = f"Tuning {self.tune_state.name}"
        return response

    # ── Tuning sequence ───────────────────────────────────────────────────────

    def _run_autotuning(self) -> None:
        # Phase 1: Settle
        self.get_logger().info("Phase 1: Settling...")
        if not self._settle():
            self.tune_state = TuneState.ABORTED
            self._publish_info(
                {"status": "aborted", "reason": "Failed to settle before tuning"}
            )
            return

        # Phase 2: Relay oscillation
        self.get_logger().info("Phase 2: Starting relay oscillation...")
        self.tune_state = TuneState.MEASURING
        if not self._relay_oscillation():
            self.tune_state = TuneState.ABORTED
            self._publish_info(
                {"status": "aborted", "reason": "Safety limit exceeded or timeout"}
            )
            return

        # Phase 3: Analyse
        self.get_logger().info("Phase 3: Analyzing measurements...")
        ku, tu = self._analyze_oscillation()
        if ku is None or tu is None:
            self.tune_state = TuneState.ABORTED
            self._publish_info(
                {"status": "aborted", "reason": "Insufficient oscillation data"}
            )
            return

        # Phase 4: Compute ZN gains
        kp, ki, kd = self._compute_ziegler_nichols(ku, tu)

        # Phase 5: Push gains to FC and save to flash
        saved = self._send_pid_to_fc(kp, ki, kd)

        self.tune_state = TuneState.COMPLETE
        self._publish_info(
            {
                "status": "complete",
                "ku": ku,
                "tu": tu,
                "kp": kp,
                "ki": ki,
                "kd": kd,
                "flash_saved": saved,
            }
        )
        self.get_logger().info(
            f"Auto-tune complete: Ku={ku:.4f}, Tu={tu:.2f}s | "
            f"Kp={kp:.4f}, Ki={ki:.4f}, Kd={kd:.4f} | flash_saved={saved}"
        )

    def _settle(self) -> bool:
        settle_start = self.get_clock().now().nanoseconds / 1e9
        relay_off = Twist()

        while True:
            now = self.get_clock().now().nanoseconds / 1e9
            elapsed = now - settle_start

            if abs(self.last_imu_tilt) > self.tilt_safety_limit:
                self.get_logger().error(
                    f"Safety abort: Tilt {self.last_imu_tilt:.1f}° > "
                    f"{self.tilt_safety_limit}°"
                )
                return False
            if elapsed > self.max_tune_duration:
                self.get_logger().error("Settling timeout")
                return False
            if elapsed > self.settle_time:
                return True

            self.pub_cmd_vel.publish(relay_off)
            time.sleep(0.01)

    def _relay_oscillation(self) -> bool:
        relay_on = Twist()
        relay_on.linear.x = self.relay_magnitude
        relay_off = Twist()

        oscillation_start = self.get_clock().now().nanoseconds / 1e9
        relay_state = True
        cycles_measured = 0

        while cycles_measured < self.oscillation_count:
            now = self.get_clock().now().nanoseconds / 1e9
            elapsed = now - oscillation_start

            if abs(self.last_imu_tilt) > self.tilt_safety_limit:
                self.get_logger().error(
                    f"Safety abort: Tilt {self.last_imu_tilt:.1f}° > "
                    f"{self.tilt_safety_limit}°"
                )
                return False
            if elapsed > self.max_tune_duration:
                self.get_logger().error("Relay oscillation timeout")
                return False

            # Switch relay at zero crossings of balance feedback
            if relay_state and self.last_feedback < 0:
                relay_state = False
                self.peaks.append((now, self.last_feedback))
                cycles_measured += 1
                self.get_logger().debug(
                    f"Cycle {cycles_measured}: peak={self.last_feedback:.4f} "
                    f"t={elapsed:.2f}s"
                )
            elif not relay_state and self.last_feedback > 0:
                relay_state = True
                self.peaks.append((now, self.last_feedback))
                cycles_measured += 1
                self.get_logger().debug(
                    f"Cycle {cycles_measured}: peak={self.last_feedback:.4f} "
                    f"t={elapsed:.2f}s"
                )

            self.pub_cmd_vel.publish(relay_on if relay_state else relay_off)
            self._publish_info(
                {
                    "status": "measuring",
                    "cycles": cycles_measured,
                    "feedback": self.last_feedback,
                    "tilt": self.last_imu_tilt,
                }
            )
            time.sleep(0.01)

        self.pub_cmd_vel.publish(relay_off)
        return len(self.peaks) >= self.oscillation_count

    def _analyze_oscillation(
        self,
    ) -> Tuple[Optional[float], Optional[float]]:
        if len(self.peaks) < 3:
            self.get_logger().error(f"Insufficient peaks: {len(self.peaks)}")
            return None, None

        # Ku = 4·d / (π·|a|)  where d = relay magnitude, a = average peak amplitude
        peak_values = [abs(p[1]) for p in self.peaks]
        avg_peak = sum(peak_values) / len(peak_values)
        ku = (4.0 * self.relay_magnitude) / (math.pi * avg_peak)

        # Tu = 2 × mean half-period (time between consecutive zero crossings)
        peak_times = [p[0] for p in self.peaks]
        time_diffs = [
            peak_times[i + 1] - peak_times[i] for i in range(len(peak_times) - 1)
        ]
        tu = 2.0 * (sum(time_diffs) / len(time_diffs))

        self.get_logger().info(
            f"Analysis: avg_peak={avg_peak:.4f}, Ku={ku:.4f}, Tu={tu:.2f}s"
        )
        return ku, tu

    @staticmethod
    def _compute_ziegler_nichols(
        ku: float, tu: float
    ) -> Tuple[float, float, float]:
        """Ziegler-Nichols PID gains from relay feedback critical params.

        Standard ZN PID table:
          Kp = 0.60 · Ku
          Ki = 1.20 · Ku / Tu     (= 2·Kp / Tu)
          Kd = 0.075 · Ku · Tu    (= Kp·Tu / 8)
        """
        kp = 0.60 * ku
        ki = 1.20 * ku / tu
        kd = 0.075 * ku * tu
        return kp, ki, kd

    # ── JLink communication ───────────────────────────────────────────────────

    def _send_pid_to_fc(self, kp: float, ki: float, kd: float) -> bool:
        """Send JLINK_CMD_PID_SET then JLINK_CMD_PID_SAVE over JLink serial.

        Returns True if both frames were written without serial error.
        """
        if not _SERIAL_AVAILABLE:
            self.get_logger().warn(
                "pyserial not available — PID gains not sent to FC. "
                "Install with: pip install pyserial"
            )
            return False

        try:
            payload_set = struct.pack("<fff", kp, ki, kd)  # 12 bytes, little-endian
            frame_set  = _build_jlink_frame(JLINK_CMD_PID_SET, payload_set)
            frame_save = _build_jlink_frame(JLINK_CMD_PID_SAVE, b"")

            with serial.Serial(
                port=self._jlink_port,
                baudrate=self._jlink_baud,
                timeout=1.0,
                write_timeout=0.5,
            ) as ser:
                ser.write(frame_set)
                time.sleep(0.05)   # allow FC to process PID_SET
                ser.write(frame_save)
<<<<<<< HEAD
                # Flash erase takes ~1s on ESP32; wait for it
=======
                # Flash erase takes ~1s on ESP32-S3; wait for it
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)
                time.sleep(1.5)

            self.get_logger().info(
                f"Sent PID_SET+PID_SAVE to FC via {self._jlink_port}: "
                f"Kp={kp:.4f}, Ki={ki:.4f}, Kd={kd:.4f}"
            )
            return True

        except Exception as exc:
            self.get_logger().error(
                f"Failed to send PID to FC on {self._jlink_port}: {exc}"
            )
            return False

    # ── Helpers ───────────────────────────────────────────────────────────────

    def _publish_info(self, data: dict) -> None:
        info = {
            "timestamp": self.get_clock().now().nanoseconds / 1e9,
            "state": self.tune_state.name,
            **data,
        }
        self.pub_info.publish(String(data=json.dumps(info)))


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


if __name__ == "__main__":
    main()
