#!/usr/bin/env python3
"""Motor current protection node for SaltyBot.

Monitors motor current, detects overcurrent conditions (soft 5A, hard 8A),
models thermal state using I^2R model, and applies speed derating.
Triggers emergency stop on hard fault.

Published topics:
  /saltybot/motor_protection (std_msgs/String) - JSON protection status
  /saltybot/speed_limit (std_msgs/Float32) - Thermal derating factor [0, 1]

Subscribed topics:
  /saltybot/motor_current_left (std_msgs/Float32) - Left motor current (A)
  /saltybot/motor_current_right (std_msgs/Float32) - Right motor current (A)
"""

import json
import math
from enum import Enum
from dataclasses import dataclass

import rclpy
from rclpy.node import Node
from rclpy.timer import Timer
from geometry_msgs.msg import Twist
from std_msgs.msg import Float32, String


class ProtectionState(Enum):
    """Motor protection state."""
    NORMAL = 0
    SOFT_FAULT = 1  # Soft overcurrent
    HARD_FAULT = 2  # Hard overcurrent


@dataclass
class ThermalModel:
    """I^2R thermal model parameters."""
    motor_resistance: float = 1.5  # Ohms (motor winding resistance)
    thermal_mass: float = 100.0  # J/°C (thermal capacitance)
    ambient_temp: float = 25.0  # °C
    max_temp: float = 80.0  # °C (safe operating limit)
    cooling_constant: float = 0.05  # 1/s (exponential cooling rate)

    # Derived
    motor_temp: float = 25.0  # Current motor temperature


class MotorProtectionNode(Node):
    """ROS2 node for motor current protection."""

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

        # Parameters
        self.declare_parameter("soft_current_limit", 5.0)  # Amps
        self.declare_parameter("hard_current_limit", 8.0)  # Amps
        self.declare_parameter("frequency", 50)  # Hz

        self.soft_limit = self.get_parameter("soft_current_limit").value
        self.hard_limit = self.get_parameter("hard_current_limit").value
        frequency = self.get_parameter("frequency").value

        # Current measurements
        self.current_left = 0.0
        self.current_right = 0.0
        self.current_max = 0.0  # Maximum of the two motors

        # Thermal model
        self.thermal = ThermalModel()

        # Protection state
        self.state = ProtectionState.NORMAL
        self.soft_fault_duration = 0.0  # Time in soft fault state

        # Subscriptions
        self.create_subscription(
            Float32, "/saltybot/motor_current_left", self._on_current_left, 10
        )
        self.create_subscription(
            Float32, "/saltybot/motor_current_right", self._on_current_right, 10
        )

        # Publications
        self.pub_protection = self.create_publisher(
            String, "/saltybot/motor_protection", 10
        )
        self.pub_speed_limit = self.create_publisher(
            Float32, "/saltybot/speed_limit", 10
        )

        # Timer for monitoring at 50Hz
        period = 1.0 / frequency
        self.timer: Timer = self.create_timer(period, self._timer_callback)

        self.get_logger().info(
            f"Motor protection initialized at {frequency}Hz. "
            f"Soft limit: {self.soft_limit}A, Hard limit: {self.hard_limit}A"
        )

    def _on_current_left(self, msg: Float32) -> None:
        """Update left motor current."""
        self.current_left = max(0.0, msg.data)  # Ensure non-negative

    def _on_current_right(self, msg: Float32) -> None:
        """Update right motor current."""
        self.current_right = max(0.0, msg.data)  # Ensure non-negative

    def _timer_callback(self) -> None:
        """Monitor motor current and thermal state at 50Hz."""
        dt = 0.02  # 50Hz = 20ms

        # Get maximum current (worst case)
        self.current_max = max(self.current_left, self.current_right)

        # Update thermal state
        self._update_thermal_state(dt)

        # Check overcurrent conditions
        self._check_overcurrent()

        # Calculate speed derate factor
        derate_factor = self._calculate_derate()

        # Publish status
        self._publish_protection_status()
        self._publish_speed_limit(derate_factor)

    def _update_thermal_state(self, dt: float) -> None:
        """Update motor temperature using I^2R thermal model.

        Model: dT/dt = (P_loss - P_cool) / C_thermal
        where:
          P_loss = I^2 * R (Joule heating)
          P_cool = h * (T - T_ambient) (convective cooling)
          C_thermal = thermal mass / dt
        """
        # Heat generation: P_loss = I^2 * R (both motors)
        p_loss_left = self.current_left**2 * self.thermal.motor_resistance
        p_loss_right = self.current_right**2 * self.thermal.motor_resistance
        p_loss_total = p_loss_left + p_loss_right  # Combined heat

        # Passive cooling (exponential decay)
        temp_diff = self.thermal.motor_temp - self.thermal.ambient_temp
        p_cool = self.thermal.cooling_constant * temp_diff

        # Temperature change
        dtemp_dt = (p_loss_total - p_cool) / self.thermal.thermal_mass

        # Update temperature (with saturation)
        new_temp = self.thermal.motor_temp + dtemp_dt * dt
        self.thermal.motor_temp = max(self.thermal.ambient_temp, new_temp)

        self.get_logger().debug(
            f"Thermal: T={self.thermal.motor_temp:.1f}°C, "
            f"P_loss={p_loss_total:.1f}W, P_cool={p_cool:.1f}W"
        )

    def _check_overcurrent(self) -> None:
        """Check for overcurrent conditions and update protection state."""
        if self.current_max >= self.hard_limit:
            # Hard fault: immediate stop
            if self.state != ProtectionState.HARD_FAULT:
                self.get_logger().error(
                    f"HARD OVERCURRENT FAULT: {self.current_max:.2f}A >= {self.hard_limit}A"
                )
            self.state = ProtectionState.HARD_FAULT
            self.soft_fault_duration = 0.0

        elif self.current_max >= self.soft_limit:
            # Soft fault: warning and derating
            if self.state == ProtectionState.NORMAL:
                self.get_logger().warn(
                    f"Soft overcurrent: {self.current_max:.2f}A >= {self.soft_limit}A"
                )
            self.state = ProtectionState.SOFT_FAULT
            self.soft_fault_duration += 0.02  # 50Hz = 20ms

        else:
            # Back to normal
            if self.state != ProtectionState.NORMAL:
                self.get_logger().info("Overcurrent cleared, resuming normal operation")
            self.state = ProtectionState.NORMAL
            self.soft_fault_duration = 0.0

    def _calculate_derate(self) -> float:
        """Calculate speed derating factor based on thermal state.

        Returns:
          float: Speed derate factor in [0, 1]
        """
        # Hard fault: complete stop
        if self.state == ProtectionState.HARD_FAULT:
            return 0.0

        # Soft fault: reduce speed based on current and temperature
        soft_derate = 1.0
        if self.state == ProtectionState.SOFT_FAULT:
            # Reduce speed proportionally to current above soft limit
            current_excess = (self.current_max - self.soft_limit) / (
                self.hard_limit - self.soft_limit
            )
            soft_derate = max(0.1, 1.0 - current_excess * 0.5)  # Min 10% speed

        # Thermal derating: reduce speed as temperature increases
        temp_range = self.thermal.max_temp - self.thermal.ambient_temp
        if temp_range > 0:
            temp_ratio = (
                self.thermal.motor_temp - self.thermal.ambient_temp
            ) / temp_range

            # Linear derating: full speed at ambient, zero at max
            if temp_ratio > 1.0:
                thermal_derate = 0.0
            elif temp_ratio > 0.8:
                # Aggressive derating near limit
                thermal_derate = (1.0 - temp_ratio) / 0.2
            else:
                thermal_derate = 1.0
        else:
            thermal_derate = 1.0

        # Combined derate: use minimum of both
        derate = min(soft_derate, thermal_derate)
        return max(0.0, min(1.0, derate))  # Clamp to [0, 1]

    def _publish_protection_status(self) -> None:
        """Publish motor protection status as JSON."""
        status = {
            "timestamp": self.get_clock().now().nanoseconds / 1e9,
            "state": self.state.name,
            "current_left": float(self.current_left),
            "current_right": float(self.current_right),
            "current_max": float(self.current_max),
            "motor_temp": float(self.thermal.motor_temp),
            "soft_limit": self.soft_limit,
            "hard_limit": self.hard_limit,
            "soft_fault_duration": float(self.soft_fault_duration),
        }

        msg = String(data=json.dumps(status))
        self.pub_protection.publish(msg)

    def _publish_speed_limit(self, derate: float) -> None:
        """Publish thermal derate speed limit."""
        msg = Float32(data=derate)
        self.pub_speed_limit.publish(msg)


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


if __name__ == "__main__":
    main()
