#!/usr/bin/env python3
"""Odometry fusion node for SaltyBot.

Fuses wheel odometry, visual odometry, and IMU data using complementary filtering.
Publishes fused /odom and broadcasts odom→base_link TF at 50Hz.

Published topics:
  /odom (nav_msgs/Odometry) - Fused odometry with covariance

TF broadcasts:
  odom→base_link - Position and orientation transform

Subscribed topics:
  /saltybot/wheel_odom (nav_msgs/Odometry) - Wheel encoder odometry
  /rtab_map/odom (nav_msgs/Odometry) - Visual odometry (SLAM)
  /imu/data (sensor_msgs/Imu) - Inertial measurement unit
"""

import math
import json
from typing import Optional, Tuple
from dataclasses import dataclass, field

import numpy as np
import rclpy
from rclpy.node import Node
from rclpy.timer import Timer

from geometry_msgs.msg import Pose, Twist, TransformStamped, Quaternion
from nav_msgs.msg import Odometry
from sensor_msgs.msg import Imu
from std_msgs.msg import String
from tf2_ros import TransformBroadcaster
from tf_transformations import quaternion_from_euler, euler_from_quaternion


@dataclass
class OdomState:
    """Fused odometry state."""
    x: float = 0.0  # Position in odom frame
    y: float = 0.0
    theta: float = 0.0  # Orientation (yaw)

    vx: float = 0.0  # Velocity in body frame
    vy: float = 0.0
    omega: float = 0.0  # Angular velocity

    cov_pos: np.ndarray = field(default_factory=lambda: np.eye(2) * 0.01)  # Position covariance
    cov_theta: float = 0.01  # Orientation covariance
    cov_vel: np.ndarray = field(default_factory=lambda: np.eye(2) * 0.01)  # Velocity covariance


class OdomFusionNode(Node):
    """ROS2 node for odometry fusion."""

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

        # Parameters
        self.declare_parameter("frequency", 50)  # Hz
        self.declare_parameter("wheel_weight", 0.6)  # Weight for wheel odometry
        self.declare_parameter("visual_weight", 0.3)  # Weight for visual odometry
        self.declare_parameter("imu_weight", 0.1)  # Weight for IMU

        frequency = self.get_parameter("frequency").value
        self.wheel_weight = self.get_parameter("wheel_weight").value
        self.visual_weight = self.get_parameter("visual_weight").value
        self.imu_weight = self.get_parameter("imu_weight").value

        # Normalize weights
        total_weight = self.wheel_weight + self.visual_weight + self.imu_weight
        self.wheel_weight /= total_weight
        self.visual_weight /= total_weight
        self.imu_weight /= total_weight

        # State
        self.state = OdomState()
        self.last_time = None

        # Last received messages
        self.last_wheel_odom: Optional[Odometry] = None
        self.last_visual_odom: Optional[Odometry] = None
        self.last_imu: Optional[Imu] = None

        # Subscriptions
        self.create_subscription(
            Odometry, "/saltybot/wheel_odom", self._on_wheel_odom, 10
        )
        self.create_subscription(
            Odometry, "/rtab_map/odom", self._on_visual_odom, 10
        )
        self.create_subscription(Imu, "/imu/data", self._on_imu, 10)

        # Publications
        self.pub_odom = self.create_publisher(Odometry, "/odom", 10)
        self.pub_info = self.create_publisher(String, "/saltybot/odom_fusion_info", 10)

        # TF broadcaster
        self.tf_broadcaster = TransformBroadcaster(self)

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

        self.get_logger().info(
            f"Odometry fusion initialized at {frequency}Hz. "
            f"Weights - wheel: {self.wheel_weight:.2f}, "
            f"visual: {self.visual_weight:.2f}, "
            f"imu: {self.imu_weight:.2f}"
        )

    def _on_wheel_odom(self, msg: Odometry) -> None:
        """Update wheel odometry."""
        self.last_wheel_odom = msg

    def _on_visual_odom(self, msg: Odometry) -> None:
        """Update visual odometry."""
        self.last_visual_odom = msg

    def _on_imu(self, msg: Imu) -> None:
        """Update IMU data."""
        self.last_imu = msg

    def _timer_callback(self) -> None:
        """Fuse odometry sources and publish at 50Hz."""
        now = self.get_clock().now()

        if self.last_time is None:
            self.last_time = now
            return

        dt = (now - self.last_time).nanoseconds / 1e9
        self.last_time = now

        # Fuse available sources
        self._fuse_odometry(dt)

        # Publish fused odometry
        self._publish_odometry(now)

        # Broadcast TF
        self._broadcast_tf(now)

        # Publish debug info
        self._publish_info()

    def _fuse_odometry(self, dt: float) -> None:
        """Fuse odometry sources using complementary filtering."""
        if dt <= 0 or dt > 1.0:  # Safety check on dt
            return

        # Collect available measurements
        measurements = {}

        if self.last_wheel_odom is not None:
            measurements["wheel"] = self.last_wheel_odom

        if self.last_visual_odom is not None:
            measurements["visual"] = self.last_visual_odom

        if self.last_imu is not None:
            measurements["imu"] = self.last_imu

        if not measurements:
            return  # No data to fuse

        # Complementary filter approach
        # High-frequency IMU for attitude, low-frequency vision for position correction

        # Extract velocities from wheel odometry (most reliable for linear motion)
        if "wheel" in measurements:
            wheel_odom = measurements["wheel"]
            self.state.vx = wheel_odom.twist.twist.linear.x * self.wheel_weight
            self.state.vy = wheel_odom.twist.twist.linear.y * self.wheel_weight

        # Integrate IMU angular velocity (high-frequency)
        if "imu" in measurements:
            imu = measurements["imu"]
            # Use IMU for angular velocity (most accurate for rotation)
            self.state.omega = imu.angular_velocity.z * self.imu_weight

            # Update orientation from IMU
            roll, pitch, yaw = euler_from_quaternion([
                imu.orientation.x, imu.orientation.y,
                imu.orientation.z, imu.orientation.w
            ])
            # Blend with existing estimate
            self.state.theta = (0.8 * self.state.theta + 0.2 * yaw)

        # Apply velocity-based position update
        self.state.x += self.state.vx * dt * math.cos(self.state.theta)
        self.state.x -= self.state.vy * dt * math.sin(self.state.theta)

        self.state.y += self.state.vx * dt * math.sin(self.state.theta)
        self.state.y += self.state.vy * dt * math.cos(self.state.theta)

        # Visual odometry correction (low-frequency drift correction)
        if "visual" in measurements and self.visual_weight > 0:
            visual_odom = measurements["visual"]
            visual_x = visual_odom.pose.pose.position.x
            visual_y = visual_odom.pose.pose.position.y

            # Gradually correct position drift using visual odometry
            correction_factor = self.visual_weight * 0.1  # Small correction rate
            self.state.x += (visual_x - self.state.x) * correction_factor
            self.state.y += (visual_y - self.state.y) * correction_factor

            # Correct orientation from visual odometry
            visual_roll, visual_pitch, visual_yaw = euler_from_quaternion([
                visual_odom.pose.pose.orientation.x,
                visual_odom.pose.pose.orientation.y,
                visual_odom.pose.pose.orientation.z,
                visual_odom.pose.pose.orientation.w
            ])
            self.state.theta += (visual_yaw - self.state.theta) * correction_factor

        # Update covariances based on available measurements
        self._update_covariances(measurements)

    def _update_covariances(self, measurements: dict) -> None:
        """Update state covariance based on measurement sources."""
        # Start with baseline
        base_pos_cov = 0.01
        base_theta_cov = 0.01

        # Reduce covariance if we have good measurements
        if "wheel" in measurements:
            base_pos_cov *= 0.5
        if "visual" in measurements:
            base_pos_cov *= 0.7
        if "imu" in measurements:
            base_theta_cov *= 0.3

        # Update covariance matrices
        self.state.cov_pos = np.eye(2) * base_pos_cov
        self.state.cov_theta = base_theta_cov

    def _publish_odometry(self, now) -> None:
        """Publish fused odometry message."""
        odom_msg = Odometry()
        odom_msg.header.stamp = now.to_msg()
        odom_msg.header.frame_id = "odom"
        odom_msg.child_frame_id = "base_link"

        # Position
        odom_msg.pose.pose.position.x = self.state.x
        odom_msg.pose.pose.position.y = self.state.y
        odom_msg.pose.pose.position.z = 0.0

        # Orientation (yaw only, roll/pitch from IMU)
        q = quaternion_from_euler(0, 0, self.state.theta)
        odom_msg.pose.pose.orientation = Quaternion(
            x=q[0], y=q[1], z=q[2], w=q[3]
        )

        # Position covariance (6x6)
        pose_cov = np.zeros((6, 6))
        pose_cov[0:2, 0:2] = self.state.cov_pos
        pose_cov[5, 5] = self.state.cov_theta  # Yaw variance
        odom_msg.pose.covariance = pose_cov.flatten().tolist()

        # Velocity
        odom_msg.twist.twist.linear.x = self.state.vx
        odom_msg.twist.twist.linear.y = self.state.vy
        odom_msg.twist.twist.linear.z = 0.0
        odom_msg.twist.twist.angular.z = self.state.omega

        # Velocity covariance (6x6)
        twist_cov = np.zeros((6, 6))
        twist_cov[0:2, 0:2] = self.state.cov_vel
        twist_cov[5, 5] = 0.01  # Angular velocity variance
        odom_msg.twist.covariance = twist_cov.flatten().tolist()

        self.pub_odom.publish(odom_msg)

    def _broadcast_tf(self, now) -> None:
        """Broadcast odom→base_link TF."""
        t = TransformStamped()
        t.header.stamp = now.to_msg()
        t.header.frame_id = "odom"
        t.child_frame_id = "base_link"

        # Translation
        t.transform.translation.x = self.state.x
        t.transform.translation.y = self.state.y
        t.transform.translation.z = 0.0

        # Rotation (yaw only)
        q = quaternion_from_euler(0, 0, self.state.theta)
        t.transform.rotation = Quaternion(
            x=q[0], y=q[1], z=q[2], w=q[3]
        )

        self.tf_broadcaster.sendTransform(t)

    def _publish_info(self) -> None:
        """Publish fusion debug information."""
        info = {
            "timestamp": self.get_clock().now().nanoseconds / 1e9,
            "position": [self.state.x, self.state.y],
            "theta_deg": math.degrees(self.state.theta),
            "velocity": [self.state.vx, self.state.vy],
            "omega": self.state.omega,
            "sources": {
                "wheel": self.last_wheel_odom is not None,
                "visual": self.last_visual_odom is not None,
                "imu": self.last_imu is not None,
            },
            "weights": {
                "wheel": f"{self.wheel_weight:.2f}",
                "visual": f"{self.visual_weight:.2f}",
                "imu": f"{self.imu_weight:.2f}",
            },
        }
        msg = String(data=json.dumps(info))
        self.pub_info.publish(msg)


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


if __name__ == "__main__":
    main()
