#!/usr/bin/env python3
"""
average_tf_publisher.py
Compute the mean of all 4x4 transform .npy files in a folder
and broadcast it as a static TF at 50 Hz.

Usage:
    ros2 run <your_package> average_tf_publisher.py [directory] [parent_frame] [child_frame]

Default parameters:
    directory     = "calibration"
    parent_frame  = "world"
    child_frame   = "average_lidar"
"""

import os
import sys
import glob
import numpy as np
from scipy.spatial.transform import Rotation as R

import rclpy
from rclpy.node import Node
from tf2_ros.static_transform_broadcaster import StaticTransformBroadcaster
from geometry_msgs.msg import TransformStamped


class AverageTFPublisher(Node):
    def __init__(self):
        super().__init__("average_tf_publisher")
        self.declare_parameter('calibration', 'calibration')
        self.declare_parameter('parent_frame', 'velodyne')
        self.declare_parameter('child_frame', 'livox_base')
        self.declare_parameter('hz', 50.0)

        # Read the values
        directory    = self.get_parameter('calibration').value
        self.parent_frame = self.get_parameter('parent_frame').value
        self.child_frame  = self.get_parameter('child_frame').value

        self.get_logger().info(f"Loading transforms from: {directory}")
        transforms = self._load_transforms(directory)
        if len(transforms) == 0:
            raise RuntimeError(f"No valid .npy transforms found in {directory}")

        self.get_logger().info(f"Loaded {len(transforms)} transforms.")
        self.mean_t, self.mean_q = self._compute_mean_transform(transforms)
        self.get_logger().info(
            f"Average translation: {self.mean_t}\n"
            f"Average quaternion:   {self.mean_q}"
        )

        self.br = StaticTransformBroadcaster(self)

        # Publish at 50 Hz (static TF can be published once, but continuous
        # publishing helps late-joining nodes)
        period = 1.0 / self.get_parameter('hz').value
        self.timer = self.create_timer(period, self._timer_cb)

    def _load_transforms(self, directory):
        """Load all 4x4 numpy transforms from directory."""
        files = sorted(glob.glob(os.path.join(directory, "*.npy")))
        mats = []
        for f in files:
            try:
                T = np.load(f)
                if T.shape == (4, 4):
                    mats.append(T)
                else:
                    self.get_logger().warn(f"Skipping {f}: wrong shape {T.shape}")
            except Exception as e:
                self.get_logger().warn(f"Failed loading {f}: {e}")
        return mats

    def _compute_mean_transform(self, mats):
        """Return (mean_translation, mean_quaternion)"""
        translations = []
        quats = []
        for T in mats:
            translations.append(T[:3, 3])
            q = R.from_matrix(T[:3, :3]).as_quat()  # [x,y,z,w]
            quats.append(q)

        # mean translation
        t_mean = np.mean(np.vstack(translations), axis=0)

        # quaternion average using eigenvector method
        Q = np.vstack(quats)
        A = Q.T @ Q
        eigvals, eigvecs = np.linalg.eigh(A)
        q_mean = eigvecs[:, np.argmax(eigvals)]
        if q_mean[3] < 0:  # ensure w positive for consistency
            q_mean = -q_mean
        return t_mean, q_mean  # q_mean = [x,y,z,w]

    def _timer_cb(self):
        t = TransformStamped()
        t.header.stamp = self.get_clock().now().to_msg()
        t.header.frame_id = self.parent_frame
        t.child_frame_id = self.child_frame

        t.transform.translation.x = float(self.mean_t[0])
        t.transform.translation.y = float(self.mean_t[1])
        t.transform.translation.z = float(self.mean_t[2])
        t.transform.rotation.x = float(self.mean_q[0])
        t.transform.rotation.y = float(self.mean_q[1])
        t.transform.rotation.z = float(self.mean_q[2])
        t.transform.rotation.w = float(self.mean_q[3])

        self.br.sendTransform(t)


def main(argv=None):
    rclpy.init(args=argv)

    try:
        node = AverageTFPublisher()
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        rclpy.shutdown()


if __name__ == "__main__":
    main()
