#!/usr/bin/env python3

import math
from dataclasses import dataclass, fields

import rclpy
from rclpy.experimental import EventsExecutor
from rclpy.node import Node
from rclpy.parameter import Parameter
from rclpy.qos import qos_profile_sensor_data

from geometry_msgs.msg import Twist
from sensor_msgs.msg import Imu
from tf_transformations import euler_from_quaternion


@dataclass
class HeadingParams:
    linear_hold_deadband: float = 0.05
    kp: float = 15.0
    max_correction: float = 1.0
    yaw_deadband: float = 0.005
    angular_hold_deadband: float = 0.2


class HeadingController(Node):
    def __init__(self) -> None:
        super().__init__("heading_controller")

        self.params = HeadingParams()
        self._declare_parameters()
        self._load_parameters()
        self.add_post_set_parameters_callback(self.post_set_parameters_callback)

        self.current_yaw: float | None = None
        self.heading_ref: float | None = None
        self.last_cmd_time: rclpy.time.Time | None = None

        self.imu_sub = self.create_subscription(
            Imu,
            "imu/data",
            self.imu_callback,
            qos_profile_sensor_data,
        )
        self.cmd_sub = self.create_subscription(
            Twist,
            "~/cmd_vel",
            self.cmd_callback,
            1,
        )
        self.cmd_pub = self.create_publisher(Twist, "cmd_vel", 1)

    def _declare_parameters(self) -> None:
        for field in fields(HeadingParams):
            self.declare_parameter(field.name, getattr(self.params, field.name))

    def _load_parameters(self) -> None:
        for field in fields(HeadingParams):
            setattr(self.params, field.name, float(self.get_parameter(field.name).value))

    def post_set_parameters_callback(self, params: list[Parameter]) -> None:
        for param in params:
            if hasattr(self.params, param.name):
                setattr(self.params, param.name, float(param.value))

    def imu_callback(self, msg: Imu) -> None:
        q = msg.orientation
        _, _, self.current_yaw = euler_from_quaternion((q.x, q.y, q.z, q.w))

    def cmd_callback(self, msg: Twist) -> None:
        now = self.get_clock().now()
        if self.last_cmd_time is None:
            dt = 0.0
        else:
            # Time delta in seconds
            dt = (now - self.last_cmd_time).nanoseconds * 1e-9
        self.last_cmd_time = now

        if self.current_yaw is None:
            self.cmd_pub.publish(msg)
            return

        heading_hold_request = (
            abs(msg.angular.z) <= self.params.angular_hold_deadband
            and abs(msg.linear.x) >= self.params.linear_hold_deadband
        )

        if heading_hold_request:
            if self.heading_ref is None:
                self.heading_ref = self.current_yaw

            # Normalized heading error to [-pi, pi]
            heading_error = math.atan2(
                math.sin(self.current_yaw - self.heading_ref),
                math.cos(self.current_yaw - self.heading_ref),
            )
            if abs(heading_error) < self.params.yaw_deadband:
                correction = 0.0
            else:
                correction = max(
                    -self.params.max_correction,
                    min(-self.params.kp * heading_error, self.params.max_correction),
                )

            filtered = Twist()
            filtered.linear.x = msg.linear.x
            filtered.linear.y = msg.linear.y
            filtered.linear.z = msg.linear.z
            filtered.angular.x = msg.angular.x
            filtered.angular.y = msg.angular.y
            filtered.angular.z = msg.angular.z + correction
            self.cmd_pub.publish(filtered)

            # Update heading reference for next iteration to prevent drift when turning
            self.heading_ref += msg.angular.z * dt
            return

        self.heading_ref = None
        self.cmd_pub.publish(msg)


if __name__ == "__main__":
    rclpy.init()
    executor = EventsExecutor()
    node = HeadingController()
    executor.add_node(node)

    try:
        executor.spin()
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()
