import rclpy
from rclpy.node import Node
from rclpy.executors import MultiThreadedExecutor
from sensor_msgs.msg import PointCloud2
from geometry_msgs.msg import TransformStamped
import sensor_msgs_py.point_cloud2 as pc2
from tf2_ros import StaticTransformBroadcaster
import numpy as np
import matplotlib.pyplot as plt
import struct
from io import BytesIO
import threading
import time
import open3d as o3d
from scipy.spatial.transform import Rotation as R

class PointCloudSaver(Node):
    def __init__(self, node_name: str, pointcloud_topic: str, buffer, timeout_ms: int):
        super().__init__(node_name)
        self.subscription = self.create_subscription(
            PointCloud2,
            pointcloud_topic,
            self.callback,
            10
        )
        self.buffer = buffer
        self.finished = False
        self.points = []
        self.end_time = self.get_clock().now().nanoseconds + (timeout_ms * 1_000_000)
        self.cmap = plt.get_cmap('jet') 

    def callback(self, msg):
        now = self.get_clock().now().nanoseconds
        for p in pc2.read_points(msg, field_names=("x", "y", "z", "intensity"), skip_nans=True):
            self.points.append([p[0], p[1], p[2], p[3]])

        if now > self.end_time:
            if not self.points:
                self.get_logger().warn("No points received!")
                self.destroy_node()
                self.finished = True
                return

            np_points = np.array(self.points, dtype=np.float32)
            intensities = np_points[:, 3]
            norm_int = (intensities - intensities.min()) / (intensities.ptp() + 1e-8)

            # Map normalized intensity to RGB colormap
            colors = self.cmap(norm_int)[:, :3]  # RGB 0-1
            colors = (colors * 255).astype(np.uint8)
            rgb_int = np.left_shift(colors[:,0].astype(np.uint32), 16) | \
                      np.left_shift(colors[:,1].astype(np.uint32), 8)  | \
                      colors[:,2].astype(np.uint32)

            filename = "pointcloud.pcd"
            self.write_pcd_with_intensity_rgb(filename, np_points, rgb_int)
            self.get_logger().info(f"Saved {filename}")
            self.destroy_node()
            self.finished = True

    def write_pcd_with_intensity_rgb(self, filename, points, rgb_int):
        header = f"""# .PCD v0.7 - Point Cloud Data file format
VERSION 0.7
FIELDS x y z intensity rgb
SIZE 4 4 4 4 4
TYPE F F F F U
COUNT 1 1 1 1 1
WIDTH {points.shape[0]}
HEIGHT 1
VIEWPOINT 0 0 0 1 0 0 0
POINTS {points.shape[0]}
DATA binary
"""
        self.buffer.write(header.encode('ascii'))
        for i in range(points.shape[0]):
            self.buffer.write(struct.pack('ffffI', points[i,0], points[i,1], points[i,2], points[i,3], rgb_int[i]))


class LidarTransformPublisher(Node):
    def __init__(self, lidar1_buffer, lidar1_frame, lidar2_buffer, lidar2_frame):
        super().__init__('static_transform_lidar_offsets')
        self.br = StaticTransformBroadcaster(self)
        self.lidar1_buffer = lidar1_buffer
        self.lidar2_buffer = lidar2_buffer
        self.lidar1_frame = lidar1_frame
        self.lidar2_frame = lidar2_frame

    def publish(self):
        self.pcd_1 = self.pcd_buffer_to_o3d(self.lidar1_buffer)
        self.pcd_2 = self.pcd_buffer_to_o3d(self.lidar2_buffer)

        self.T = self.compute_transform()
        self.get_logger().info(f"Computed initial transform:\n{self.T}")

        T_copy = np.array(self.T, copy=True)
        trans = T_copy[:3, 3]
        rot_quat = R.from_matrix(T_copy[:3, :3]).as_quat()

        t = TransformStamped()
        t.header.stamp.sec = 0
        t.header.stamp.nanosec = 0
        t.header.frame_id = self.lidar1_frame
        t.child_frame_id = self.lidar2_frame
        t.transform.translation.x = trans[0]
        t.transform.translation.y = trans[1]
        t.transform.translation.z = trans[2]
        t.transform.rotation.x = rot_quat[0]
        t.transform.rotation.y = rot_quat[1]
        t.transform.rotation.z = rot_quat[2]
        t.transform.rotation.w = rot_quat[3]

        self.br.sendTransform(t)
        self.get_logger().info("Published static transform.")

    def pcd_buffer_to_o3d(self, buffer: BytesIO):
        buffer.seek(0)
        header_lines = []
        while True:
            line = buffer.readline().decode('ascii').strip()
            if not line:
                continue
            header_lines.append(line)
            if line.startswith("DATA"):
                data_line = line
                break

        if not data_line.lower().startswith("data binary"):
            raise NotImplementedError("Only binary PCD supported")

        num_points = 0
        for line in header_lines:
            if line.startswith("POINTS"):
                num_points = int(line.split()[1])
        if num_points == 0:
            raise ValueError("PCD header missing point count")

        dtype = np.dtype([
            ('x', 'f4'), ('y', 'f4'), ('z', 'f4'),
            ('intensity', 'f4'), ('rgb', 'u4')
        ])

        data = buffer.read(num_points * dtype.itemsize)
        points_array = np.frombuffer(data, dtype=dtype, count=num_points)

        pcd = o3d.geometry.PointCloud()
        xyz = np.stack([points_array['x'], points_array['y'], points_array['z']], axis=-1)
        pcd.points = o3d.utility.Vector3dVector(xyz)
        return pcd

    def compute_transform(self):
        voxel_size = 0.2  # coarse-to-fine pyramid base

        # --- Feature extraction ---
        def preprocess(pcd, voxel):
            pcd_down = pcd.voxel_down_sample(voxel)
            pcd_down.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=voxel*2, max_nn=30))
            fpfh = o3d.pipelines.registration.compute_fpfh_feature(
                pcd_down, o3d.geometry.KDTreeSearchParamHybrid(radius=voxel*5, max_nn=100))
            return pcd_down, fpfh

        src_down, src_fpfh = preprocess(self.pcd_2, voxel_size)
        tgt_down, tgt_fpfh = preprocess(self.pcd_1, voxel_size)

        # --- Global alignment with RANSAC ---
        result_ransac = o3d.pipelines.registration.registration_ransac_based_on_feature_matching(
            src_down, tgt_down, src_fpfh, tgt_fpfh, True,
            1.5,  # distance threshold
            o3d.pipelines.registration.TransformationEstimationPointToPoint(False),
            4,
            [
                o3d.pipelines.registration.CorrespondenceCheckerBasedOnEdgeLength(0.9),
                o3d.pipelines.registration.CorrespondenceCheckerBasedOnDistance(1.5)
            ],
            o3d.pipelines.registration.RANSACConvergenceCriteria(4000000, 500)
        )

        trans_init = result_ransac.transformation

        # --- Multi-scale ICP refinement ---
        voxel_radii = [0.2, 0.1, 0.05]
        max_iters = [50, 30, 14]
        transformation = trans_init
        for radius, iters in zip(voxel_radii, max_iters):
            src_down = self.pcd_2.voxel_down_sample(radius)
            tgt_down = self.pcd_1.voxel_down_sample(radius)
            src_down.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=radius*2, max_nn=30))
            tgt_down.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=radius*2, max_nn=30))

            result_icp = o3d.pipelines.registration.registration_icp(
                src_down, tgt_down, radius*2, transformation,
                o3d.pipelines.registration.TransformationEstimationPointToPlane()
            )
            transformation = result_icp.transformation

        return transformation


def monitor_nodes(nodes, publisher):
    while rclpy.ok():
        if all(node.finished for node in nodes):
            print("All pointclouds captured")
            publisher.publish()
            return
        time.sleep(0.1)

def main():
    rclpy.init()
    buffer_velodyne = BytesIO()
    buffer_livox = BytesIO()

    executor = MultiThreadedExecutor()
    nodes = [
        PointCloudSaver('velodyne_pcd_saver', '/velodyne_points', buffer_velodyne, 350),
        PointCloudSaver('livox_pcd_saver', '/livox/lidar', buffer_livox, 1000),
    ]
    publisher = LidarTransformPublisher(buffer_velodyne, 'velodyne', buffer_livox, 'frame_default')

    monitor_thread = threading.Thread(target=monitor_nodes, args=(nodes,publisher), daemon=True)
    monitor_thread.start()

    for node in nodes:
        executor.add_node(node)
    try:
        executor.spin()
    finally:
        monitor_thread.join()
        rclpy.shutdown()


if __name__ == "__main__":
    main()
