#!/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 PD gains. 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 time
from enum import Enum
from typing import Optional, Tuple, List

import rclpy
from rclpy.node import Node
from rclpy.service import Service
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


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)  # Relay on-off magnitude
        self.declare_parameter("tilt_safety_limit", 25.0)  # degrees
        self.declare_parameter("max_tune_duration", 120.0)  # seconds
        self.declare_parameter("oscillation_count", 5)  # Number of cycles to measure
        self.declare_parameter("settle_time", 2.0)  # Settle before starting relay

        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

        # State tracking
        self.tune_state = TuneState.IDLE
        self.last_imu_tilt = 0.0
        self.last_feedback = 0.0
        self.feedback_history: List[Tuple[float, float]] = []  # (time, feedback)
        self.peaks: List[Tuple[float, float]] = []  # (time, peak_value)
        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"
        )

    def _on_imu(self, msg: Imu) -> None:
        """Extract tilt angle from IMU."""
        # Approximate tilt from acceleration (pitch angle)
        # pitch = atan2(ax, sqrt(ay^2 + az^2))
        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:
        """Record balance feedback for oscillation measurement."""
        now = self.get_clock().now().nanoseconds / 1e9
        self.last_feedback = msg.data

        # During tuning, record history
        if self.tune_state in (TuneState.MEASURING,):
            self.feedback_history.append((now, msg.data))

    def _handle_autotune_request(
        self, request: Trigger.Request, response: Trigger.Response
    ) -> Trigger.Response:
        """Handle service request to start auto-tuning."""
        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

        # Start tuning process
        self._run_autotuning()

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

    def _run_autotuning(self) -> None:
        """Execute the complete auto-tuning sequence."""
        # 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: Analysis
        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 gains
        kp, kd = self._compute_ziegler_nichols(ku, tu)

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

    def _settle(self) -> bool:
        """Allow system to settle before starting relay."""
        settle_start = self.get_clock().now().nanoseconds / 1e9
        relay_off = Twist()  # Zero command

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

            # Safety check
            if abs(self.last_imu_tilt) > self.tilt_safety_limit:
                self.get_logger().error(
                    f"Safety abort: Tilt {self.last_imu_tilt:.1f}° > {self.tilt_safety_limit}°"
                )
                return False

            # Timeout check
            if elapsed > self.max_tune_duration:
                self.get_logger().error("Settling timeout")
                return False

            # Continue settling
            if elapsed > self.settle_time:
                return True

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

    def _relay_oscillation(self) -> bool:
        """Apply relay control to induce oscillation."""
        relay_on = Twist()
        relay_on.linear.x = self.relay_magnitude

        relay_off = Twist()

        oscillation_start = self.get_clock().now().nanoseconds / 1e9
        relay_state = True  # True = relay ON, False = relay OFF
        last_switch = oscillation_start
        cycles_measured = 0
        last_peak_value = None

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

            # Safety checks
            if abs(self.last_imu_tilt) > self.tilt_safety_limit:
                self.get_logger().error(
                    f"Safety abort: Tilt {self.last_imu_tilt:.1f}° > {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
            if relay_state and self.last_feedback < 0:
                relay_state = False
                self.peaks.append((now, self.last_feedback))
                cycles_measured += 1
                last_switch = now
                self.get_logger().debug(
                    f"Cycle {cycles_measured}: Peak at {now - oscillation_start:.2f}s = {self.last_feedback:.4f}"
                )

            elif not relay_state and self.last_feedback > 0:
                relay_state = True
                self.peaks.append((now, self.last_feedback))
                cycles_measured += 1
                last_switch = now
                self.get_logger().debug(
                    f"Cycle {cycles_measured}: Peak at {now - oscillation_start:.2f}s = {self.last_feedback:.4f}"
                )

            # Publish relay command
            cmd = relay_on if relay_state else relay_off
            self.pub_cmd_vel.publish(cmd)

            # Report progress
            self._publish_info(
                {
                    "status": "measuring",
                    "cycles": cycles_measured,
                    "feedback": self.last_feedback,
                    "tilt": self.last_imu_tilt,
                }
            )

            time.sleep(0.01)

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

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

        # Ku = 4*relay_magnitude / (pi * |peak|)
        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 * (time between peaks) - average period
        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]:
        """Compute Ziegler-Nichols PD gains.

        For PD controller:
          Kp = 0.6 * Ku
          Kd = 0.075 * Ku * Tu (approximately Kp*Tu/8)
        """
        kp = 0.6 * ku
        kd = 0.075 * ku * tu
        return kp, kd

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


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()
