#!/usr/bin/env python3
"""Adaptive PID gain scheduler for SaltyBot.

Monitors speed scale and terrain roughness, adjusts PID gains dynamically.
Higher P on rough terrain, lower D at low speed.

Subscribed topics:
  /saltybot/speed_scale (std_msgs/Float32) - Speed reduction factor (0.0-1.0)
  /saltybot/terrain_roughness (std_msgs/Float32) - Terrain roughness (0.0-1.0)

Published topics:
  /saltybot/pid_gains (std_msgs/Float32MultiArray) - [Kp, Ki, Kd] gains
"""

import rclpy
from rclpy.node import Node
from std_msgs.msg import Float32, Float32MultiArray
import numpy as np


class PIDSchedulerNode(Node):
    """ROS2 node for adaptive PID gain scheduling."""

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

        # Base PID gains
        self.declare_parameter("base_kp", 1.0)
        self.declare_parameter("base_ki", 0.1)
        self.declare_parameter("base_kd", 0.05)

        # Gain scheduling parameters
        self.declare_parameter("kp_terrain_scale", 0.5)  # P gain increase on rough terrain
        self.declare_parameter("kd_speed_scale", -0.3)   # D gain decrease at low speed
        self.declare_parameter("ki_scale", 0.1)          # I gain smoothing

        self.base_kp = self.get_parameter("base_kp").value
        self.base_ki = self.get_parameter("base_ki").value
        self.base_kd = self.get_parameter("base_kd").value

        self.kp_terrain_scale = self.get_parameter("kp_terrain_scale").value
        self.kd_speed_scale = self.get_parameter("kd_speed_scale").value
        self.ki_scale = self.get_parameter("ki_scale").value

        # Current sensor inputs
        self.speed_scale = 1.0
        self.terrain_roughness = 0.0

        # Current gains
        self.current_kp = self.base_kp
        self.current_ki = self.base_ki
        self.current_kd = self.base_kd

        # Subscriptions
        self.sub_speed = self.create_subscription(
            Float32, "/saltybot/speed_scale", self._on_speed_scale, 10
        )
        self.sub_terrain = self.create_subscription(
            Float32, "/saltybot/terrain_roughness", self._on_terrain_roughness, 10
        )

        # Publisher for PID gains
        self.pub_gains = self.create_publisher(Float32MultiArray, "/saltybot/pid_gains", 10)

        self.get_logger().info(
            f"PID gain scheduler initialized. Base gains: "
            f"Kp={self.base_kp}, Ki={self.base_ki}, Kd={self.base_kd}"
        )

    def _on_speed_scale(self, msg: Float32) -> None:
        """Update speed scale and recalculate gains."""
        self.speed_scale = np.clip(msg.data, 0.0, 1.0)
        self._update_gains()

    def _on_terrain_roughness(self, msg: Float32) -> None:
        """Update terrain roughness and recalculate gains."""
        self.terrain_roughness = np.clip(msg.data, 0.0, 1.0)
        self._update_gains()

    def _update_gains(self) -> None:
        """Compute adaptive PID gains based on speed and terrain."""
        # P gain increases with terrain roughness (better response on rough surfaces)
        self.current_kp = self.base_kp * (1.0 + self.kp_terrain_scale * self.terrain_roughness)

        # D gain decreases at low speed (avoid oscillation at low speeds)
        # Low speed_scale means robot is moving slow, so reduce D damping
        speed_factor = 1.0 + self.kd_speed_scale * (1.0 - self.speed_scale)
        self.current_kd = self.base_kd * max(0.1, speed_factor)  # Don't let D go negative

        # I gain scales smoothly with both factors
        terrain_factor = 1.0 + self.ki_scale * self.terrain_roughness
        speed_damping = 1.0 - 0.3 * (1.0 - self.speed_scale)  # Reduce I at low speed
        self.current_ki = self.base_ki * terrain_factor * speed_damping

        # Publish gains
        self._publish_gains()

    def _publish_gains(self) -> None:
        """Publish current PID gains as Float32MultiArray."""
        msg = Float32MultiArray()
        msg.data = [self.current_kp, self.current_ki, self.current_kd]

        self.pub_gains.publish(msg)

        self.get_logger().debug(
            f"PID gains updated: Kp={self.current_kp:.4f}, "
            f"Ki={self.current_ki:.4f}, Kd={self.current_kd:.4f} "
            f"(speed={self.speed_scale:.2f}, terrain={self.terrain_roughness:.2f})"
        )


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


if __name__ == "__main__":
    main()
