#!/usr/bin/env python3
"""IMU calibration node for SaltyBot."""

import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Imu
from std_srvs.srv import Trigger
from std_msgs.msg import Header
import numpy as np
from collections import deque


class IMUCalibrationNode(Node):
    """ROS2 node for IMU gyro + accel calibration."""

    def __init__(self):
        super().__init__("imu_calibration")
        self.declare_parameter("calibration_samples", 100)
        self.declare_parameter("auto_calibrate", False)

        self.calibration_samples = self.get_parameter("calibration_samples").value
        self.auto_calibrate = self.get_parameter("auto_calibrate").value

        self.gyro_bias = np.array([0.0, 0.0, 0.0])
        self.accel_bias = np.array([0.0, 0.0, 0.0])
        self.is_calibrated = False

        self.gyro_samples = deque(maxlen=self.calibration_samples)
        self.accel_samples = deque(maxlen=self.calibration_samples)
        self.calibrating = False

        self.sub_imu = self.create_subscription(Imu, "/imu", self._on_imu_raw, 10)
        self.pub_calibrated = self.create_publisher(Imu, "/imu/calibrated", 10)
        self.srv_calibrate = self.create_service(
            Trigger, "/saltybot/calibrate_imu", self._on_calibrate_service
        )

        self.get_logger().info(
            f"IMU calibration node initialized. Samples: {self.calibration_samples}. Auto: {self.auto_calibrate}"
        )

        if self.auto_calibrate:
            self.calibrating = True
            self.get_logger().info("Starting auto-calibration...")

    def _on_imu_raw(self, msg: Imu) -> None:
        if self.calibrating:
            gyro = np.array([msg.angular_velocity.x, msg.angular_velocity.y, msg.angular_velocity.z])
            accel = np.array([msg.linear_acceleration.x, msg.linear_acceleration.y, msg.linear_acceleration.z])
            self.gyro_samples.append(gyro)
            self.accel_samples.append(accel)

            if len(self.gyro_samples) == self.calibration_samples:
                self._compute_calibration()
        else:
            self._publish_calibrated(msg)

    def _compute_calibration(self) -> None:
        if len(self.gyro_samples) == 0 or len(self.accel_samples) == 0:
            return

        gyro_data = np.array(list(self.gyro_samples))
        accel_data = np.array(list(self.accel_samples))

        self.gyro_bias = np.mean(gyro_data, axis=0)
        self.accel_bias = np.mean(accel_data, axis=0)

        self.is_calibrated = True
        self.calibrating = False

        self.get_logger().info(
            f"Calibration complete. Gyro: {self.gyro_bias}. Accel: {self.accel_bias}"
        )

        self.gyro_samples.clear()
        self.accel_samples.clear()

    def _on_calibrate_service(self, request, response) -> Trigger.Response:
        if self.calibrating:
            response.success = False
            response.message = "Calibration already in progress"
            return response

        self.get_logger().info("Calibration service called")
        self.calibrating = True
        self.gyro_samples.clear()
        self.accel_samples.clear()

        response.success = True
        response.message = f"Calibration started, collecting {self.calibration_samples} samples"
        return response

    def _publish_calibrated(self, msg: Imu) -> None:
        calibrated = Imu()
        calibrated.header = Header(frame_id=msg.header.frame_id, stamp=msg.header.stamp)

        calibrated.angular_velocity.x = msg.angular_velocity.x - self.gyro_bias[0]
        calibrated.angular_velocity.y = msg.angular_velocity.y - self.gyro_bias[1]
        calibrated.angular_velocity.z = msg.angular_velocity.z - self.gyro_bias[2]

        calibrated.linear_acceleration.x = msg.linear_acceleration.x - self.accel_bias[0]
        calibrated.linear_acceleration.y = msg.linear_acceleration.y - self.accel_bias[1]
        calibrated.linear_acceleration.z = msg.linear_acceleration.z - self.accel_bias[2]

        calibrated.angular_velocity_covariance = msg.angular_velocity_covariance
        calibrated.linear_acceleration_covariance = msg.linear_acceleration_covariance
        calibrated.orientation_covariance = msg.orientation_covariance
        calibrated.orientation = msg.orientation

        self.pub_calibrated.publish(calibrated)


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


if __name__ == "__main__":
    main()
