"""
Standalone MCAP Reader for extracting robotics data.

A simple, self-contained reader for MCAP ROS 2 recordings that provides
getter methods for images, camera intrinsics/extrinsics, and low-dimensional
signals. Does not require ray or the full preprocessing pipeline.

Usage: 
pip install rosbags
python mcap_reader.py <path_to_mcap_file>

Ex: python3 mcap_reader.py datasets/unitree_g1_dex3/mcap/stack_cubes_ordered/real/teleop/0000_20251215_134924/0000_20251215_134924_0.mcap 
"""

import logging
from collections import defaultdict
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

import cv2
import numpy as np

from scipy.interpolate import interp1d
from scipy.signal import butter, filtfilt
from scipy.spatial.transform import Rotation as R
from rosbags.highlevel import AnyReader

logger = logging.getLogger(__name__)

JPEG_QUALITY = 95


# ---------------------------------------------------------------------------
# Image extraction helpers
# ---------------------------------------------------------------------------


def extract_depth_from_msg(msg: Any) -> Optional[np.ndarray]:
    """
    Extract depth image from a ROS 2 Image message with 16UC1 or 32FC1 encoding.

    Returns:
        uint16 depth array [H, W] (values in mm), or None if not a depth message.
    """
    if not hasattr(msg, "height") or not hasattr(msg, "width") or not hasattr(msg, "encoding"):
        return None

    height, width, encoding = msg.height, msg.width, msg.encoding

    if "16UC1" in encoding:
        depth = np.frombuffer(msg.data, dtype=np.uint16).reshape(height, width)
        return depth
    elif "32FC1" in encoding:
        depth_f = np.frombuffer(msg.data, dtype=np.float32).reshape(height, width)
        # Convert meters to mm uint16 (same as HE convention)
        depth = (depth_f * 1000).clip(0, 65535).astype(np.uint16)
        return depth

    return None


def extract_image_from_msg(msg: Any, return_numpy: bool = False) -> Optional[Union[bytes, np.ndarray]]:
    """
    Extract image data from a ROS 2 image message.

    Args:
        msg: Deserialized ROS 2 message (sensor_msgs/CompressedImage or sensor_msgs/Image)
        return_numpy: If True, return decoded numpy RGB array instead of JPEG bytes

    Returns:
        JPEG bytes or numpy RGB array [H, W, 3], or None if unsupported format
    """
    # CompressedImage: has 'format' and 'data' but no 'height'
    if hasattr(msg, "format") and hasattr(msg, "data") and not hasattr(msg, "height"):
        if "jpeg" in msg.format.lower() or "jpg" in msg.format.lower():
            if return_numpy:
                img_array = cv2.imdecode(np.frombuffer(bytes(msg.data), dtype=np.uint8), cv2.IMREAD_COLOR)
                if img_array is None:
                    raise ValueError(f"Failed to decode JPEG CompressedImage (format={msg.format})")
                return cv2.cvtColor(img_array, cv2.COLOR_BGR2RGB)
            return bytes(msg.data)

        elif "png" in msg.format.lower():
            img_array = cv2.imdecode(np.frombuffer(bytes(msg.data), dtype=np.uint8), cv2.IMREAD_COLOR)
            if img_array is None:
                raise ValueError(f"Failed to decode PNG CompressedImage (format={msg.format})")
            img_rgb = cv2.cvtColor(img_array, cv2.COLOR_BGR2RGB)
            if return_numpy:
                return img_rgb
            success, jpeg_data = cv2.imencode(".jpg", img_rgb, [cv2.IMWRITE_JPEG_QUALITY, JPEG_QUALITY])
            if not success:
                raise ValueError("JPEG encoding failed for PNG CompressedImage")
            return jpeg_data.tobytes()

        return None

    # Raw Image: requires height, width, encoding
    if not hasattr(msg, "height") or not hasattr(msg, "width") or not hasattr(msg, "encoding"):
        return None

    height, width, encoding = msg.height, msg.width, msg.encoding
    data = msg.data

    # Depth images handled separately by extract_depth_from_msg
    if "UC1" in encoding or "FC1" in encoding:
        return None

    if "rgb8" in encoding or "bgr8" in encoding:
        img_array = np.frombuffer(data, dtype=np.uint8).reshape(height, width, 3)
        if "bgr8" in encoding:
            img_array = cv2.cvtColor(img_array, cv2.COLOR_BGR2RGB)
    elif "mono8" in encoding or "gray" in encoding or encoding == "8UC1":
        img_array = np.frombuffer(data, dtype=np.uint8).reshape(height, width)
        img_array = cv2.cvtColor(img_array, cv2.COLOR_GRAY2RGB)
    else:
        logger.debug(f"Unsupported image encoding: {encoding}")
        return None

    if return_numpy:
        return img_array

    # Convert RGB to BGR for JPEG encoding
    img_bgr = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
    success, jpeg_data = cv2.imencode(".jpg", img_bgr, [cv2.IMWRITE_JPEG_QUALITY, JPEG_QUALITY])
    if not success:
        raise ValueError(f"JPEG encoding failed for raw Image (encoding={encoding})")
    return jpeg_data.tobytes()


# ---------------------------------------------------------------------------
# Structured message extraction
# ---------------------------------------------------------------------------


def _matrix_to_rot_6d(matrix: np.ndarray) -> np.ndarray:
    """Convert 3x3 rotation matrix to 6D representation (first two columns flattened)."""
    return matrix[:, :2].T.flatten().astype(np.float32)


def extract_structured_msg(msg: Any) -> Optional[Dict[str, np.ndarray]]:
    """
    Extract structured key-value data from standard ROS 2 messages.

    Supported:
        sensor_msgs/JointState: {"__<joint_name>": array([position])}
        geometry_msgs/PoseStamped, Pose: {"__xyz": array, "__rot_6d": array, "__quat": array}

    Returns None for unsupported message types.
    """
    # sensor_msgs/JointState
    if hasattr(msg, "name") and hasattr(msg, "position") and hasattr(msg, "velocity"):
        if not msg.name:
            return None
        result = {}
        for joint_name, pos in zip(msg.name, msg.position, strict=True):
            result[f"__{joint_name}"] = np.asarray([pos], dtype=np.float32)
        return result

    # geometry_msgs/PoseStamped or Pose
    pose = None
    if hasattr(msg, "pose") and hasattr(msg.pose, "position") and hasattr(msg.pose, "orientation"):
        pose = msg.pose
    elif hasattr(msg, "position") and hasattr(msg, "orientation") and not hasattr(msg, "velocity"):
        pose = msg

    if pose is not None:
        xyz = np.array([pose.position.x, pose.position.y, pose.position.z], dtype=np.float32)
        quat = np.array(
            [pose.orientation.x, pose.orientation.y, pose.orientation.z, pose.orientation.w], dtype=np.float32
        )
        rot_matrix = R.from_quat(quat).as_matrix()
        rot_6d = _matrix_to_rot_6d(rot_matrix)
        return {"__xyz": xyz, "__rot_6d": rot_6d, "__quat": quat}

    return None


def extract_array_from_msg(msg: Any) -> Optional[np.ndarray]:
    """
    Extract a flat numeric array from a ROS 2 message.

    Handles Imu, WrenchStamped, Wrench, and generic messages with numeric fields.
    """
    arrays = []

    # sensor_msgs/Imu
    if hasattr(msg, "orientation") and hasattr(msg, "angular_velocity") and hasattr(msg, "linear_acceleration"):
        arrays.append(np.array([msg.orientation.x, msg.orientation.y, msg.orientation.z, msg.orientation.w]))
        arrays.append(np.array([msg.angular_velocity.x, msg.angular_velocity.y, msg.angular_velocity.z]))
        arrays.append(np.array([msg.linear_acceleration.x, msg.linear_acceleration.y, msg.linear_acceleration.z]))

    # geometry_msgs/WrenchStamped
    elif hasattr(msg, "wrench") and hasattr(msg.wrench, "force"):
        wrench = msg.wrench
        arrays.append(np.array([wrench.force.x, wrench.force.y, wrench.force.z]))
        arrays.append(np.array([wrench.torque.x, wrench.torque.y, wrench.torque.z]))

    # geometry_msgs/Wrench
    elif hasattr(msg, "force") and hasattr(msg, "torque") and hasattr(msg.force, "x"):
        arrays.append(np.array([msg.force.x, msg.force.y, msg.force.z]))
        arrays.append(np.array([msg.torque.x, msg.torque.y, msg.torque.z]))

    # Generic: recursively extract numeric fields
    else:
        arrays = _extract_numeric_fields(msg)

    return np.concatenate(arrays).astype(np.float32) if arrays else None


def _extract_numeric_fields(msg: Any) -> List[np.ndarray]:
    """Recursively extract numeric fields from a message."""
    arrays = []

    try:
        attributes = list(vars(msg).keys())
    except TypeError:
        attributes = [a for a in dir(msg) if not a.startswith("_")]

    for field_name in attributes:
        if field_name.startswith("_"):
            continue

        try:
            field = getattr(msg, field_name)
        except AttributeError:
            continue

        if callable(field):
            continue

        if isinstance(field, (list, tuple)):
            try:
                arr = np.array(field, dtype=np.float32)
                if arr.size > 0 and np.issubdtype(arr.dtype, np.number):
                    arrays.append(arr.flatten())
            except (ValueError, TypeError):
                pass
        elif isinstance(field, (int, float, np.number)):
            arrays.append(np.array([field], dtype=np.float32))
        elif hasattr(field, "__dict__") or hasattr(field, "__slots__"):
            nested_array = extract_array_from_msg(field)
            if nested_array is not None:
                arrays.append(nested_array)

    return arrays


# ---------------------------------------------------------------------------
# Camera info extraction
# ---------------------------------------------------------------------------


def extract_camera_info(msg: Any) -> Optional[Dict[str, Any]]:
    """
    Extract camera calibration info from sensor_msgs/CameraInfo message.

    Returns dict with:
        - K: 3x3 intrinsic matrix
        - D: distortion coefficients
        - R: 3x3 rectification matrix (if available)
        - P: 3x4 projection matrix (if available)
        - width, height: image dimensions
        - distortion_model: distortion model name
    """
    if not (hasattr(msg, "k") and hasattr(msg, "width") and hasattr(msg, "height")):
        return None

    info = {
        "width": msg.width,
        "height": msg.height,
        "K": np.array(msg.k, dtype=np.float64).reshape(3, 3) if len(msg.k) == 9 else None,
        "D": np.array(msg.d, dtype=np.float64) if hasattr(msg, "d") and len(msg.d) > 0 else None,
    }

    if hasattr(msg, "r") and len(msg.r) == 9:
        info["R"] = np.array(msg.r, dtype=np.float64).reshape(3, 3)

    if hasattr(msg, "p") and len(msg.p) == 12:
        info["P"] = np.array(msg.p, dtype=np.float64).reshape(3, 4)

    if hasattr(msg, "distortion_model"):
        info["distortion_model"] = msg.distortion_model

    return info


def extract_transform(msg: Any) -> Optional[Dict[str, np.ndarray]]:
    """
    Extract transform from geometry_msgs/TransformStamped or Transform.

    Returns dict with:
        - translation: [x, y, z]
        - rotation_quat: [x, y, z, w]
        - rotation_matrix: 3x3 rotation matrix
        - transform_matrix: 4x4 homogeneous transform
    """
    transform = None
    frame_id = None
    child_frame_id = None

    if hasattr(msg, "transform"):
        transform = msg.transform
        if hasattr(msg, "header") and hasattr(msg.header, "frame_id"):
            frame_id = msg.header.frame_id
        if hasattr(msg, "child_frame_id"):
            child_frame_id = msg.child_frame_id
    elif hasattr(msg, "translation") and hasattr(msg, "rotation"):
        transform = msg

    if transform is None:
        return None

    translation = np.array(
        [transform.translation.x, transform.translation.y, transform.translation.z], dtype=np.float64
    )

    quat = np.array(
        [transform.rotation.x, transform.rotation.y, transform.rotation.z, transform.rotation.w], dtype=np.float64
    )

    rot_matrix = R.from_quat(quat).as_matrix()

    # Build 4x4 homogeneous transform
    transform_matrix = np.eye(4, dtype=np.float64)
    transform_matrix[:3, :3] = rot_matrix
    transform_matrix[:3, 3] = translation

    result = {
        "translation": translation,
        "rotation_quat": quat,
        "rotation_matrix": rot_matrix,
        "transform_matrix": transform_matrix,
    }

    if frame_id is not None:
        result["frame_id"] = frame_id
    if child_frame_id is not None:
        result["child_frame_id"] = child_frame_id

    return result


# ---------------------------------------------------------------------------
# Resampling
# ---------------------------------------------------------------------------


class TemporalResampler:
    """Resamples multi-rate signals to a uniform target frequency."""

    def __init__(self, target_hz: float):
        if target_hz <= 0:
            raise ValueError(f"target_hz must be positive, got {target_hz}")
        self.target_hz = target_hz

    def create_target_timeline(self, start_time: float, end_time: float) -> np.ndarray:
        """Create uniform timeline at target frequency."""
        step = 1.0 / self.target_hz
        target_times = np.arange(start_time, end_time, step)
        if len(target_times) == 0:
            target_times = np.array([start_time])
        return target_times

    def resample_continuous(
        self,
        source_times: np.ndarray,
        source_values: np.ndarray,
        target_times: np.ndarray,
        method: str = "linear",
    ) -> np.ndarray:
        """Resample continuous signals with optional anti-aliasing filter."""
        values = source_values

        # Anti-aliasing filter for downsampled signals
        if len(source_times) > 8:
            dt_median = np.median(np.diff(source_times))
            if dt_median > 0:
                source_hz = 1.0 / dt_median
                if source_hz > self.target_hz * 2.0:
                    nyquist = source_hz / 2.0
                    cutoff = self.target_hz / 2.0
                    normalized_cutoff = cutoff / nyquist
                    if 0 < normalized_cutoff < 1:
                        b, a = butter(4, normalized_cutoff, btype="low")
                        values = values.copy()
                        if len(values.shape) == 1:
                            values = filtfilt(b, a, values)
                        else:
                            for i in range(values.shape[1]):
                                values[:, i] = filtfilt(b, a, values[:, i])

        if len(values.shape) == 1:
            interpolator = interp1d(source_times, values, kind=method, bounds_error=False, fill_value="extrapolate")
            return interpolator(target_times)
        else:
            resampled = np.zeros((len(target_times), values.shape[1]))
            for i in range(values.shape[1]):
                interpolator = interp1d(
                    source_times, values[:, i], kind=method, bounds_error=False, fill_value="extrapolate"
                )
                resampled[:, i] = interpolator(target_times)
            return resampled

    def resample_nearest(
        self,
        source_times: np.ndarray,
        source_values: Union[np.ndarray, List],
        target_times: np.ndarray,
    ) -> Union[np.ndarray, List]:
        """Resample using nearest neighbor (for images or discrete signals)."""
        indices = np.array([np.argmin(np.abs(source_times - t)) for t in target_times])

        if isinstance(source_values, list):
            return [source_values[i] for i in indices]
        else:
            return source_values[indices]


# ---------------------------------------------------------------------------
# MCAPReader
# ---------------------------------------------------------------------------


class MCAPReader:
    """
    Standalone MCAP reader with getter methods for robotics data.

    Loads an MCAP file and provides methods to access:
    - Images from camera topics
    - Camera intrinsics and extrinsics
    - Low-dimensional signals (joint states, poses, IMU, etc.)
    - Timestamps and metadata

    Args:
        mcap_path: Path to the MCAP file
        target_hz: Target frequency for resampling (None to keep original rates)
        return_numpy_images: If True, return images as numpy RGB arrays instead of JPEG bytes
    """

    def __init__(
        self,
        mcap_path: str,
        target_hz: Optional[float] = None,
        return_numpy_images: bool = True,
    ):
        self.mcap_path = Path(mcap_path)
        if not self.mcap_path.exists():
            raise FileNotFoundError(f"MCAP file not found: {mcap_path}")

        self.target_hz = target_hz
        self.return_numpy_images = return_numpy_images
        self.resampler = TemporalResampler(target_hz) if target_hz else None

        # Data storage
        self._topic_data: Dict[str, Dict[str, Any]] = defaultdict(lambda: {"timestamps": [], "values": []})
        self._image_topics: set = set()
        self._depth_topics: set = set()
        self._camera_info_topics: set = set()
        self._tf_data: List[Dict] = []
        self._topic_types: Dict[str, str] = {}

        # Load data
        self._load_mcap()

    def _load_mcap(self):
        """Load and parse all data from the MCAP file."""
        logger.info(f"Loading MCAP: {self.mcap_path}")

        with AnyReader([self.mcap_path]) as reader:
            # Store topic types
            for connection in reader.connections:
                self._topic_types[connection.topic] = connection.msgtype

            for connection, timestamp, rawdata in reader.messages():
                topic = connection.topic
                timestamp_sec = timestamp / 1e9

                try:
                    msg = reader.deserialize(rawdata, connection.msgtype)
                    self._process_message(topic, timestamp_sec, msg, connection.msgtype)
                except Exception as e:
                    logger.warning(f"Failed to deserialize message on {topic}: {e}")

        # Convert to numpy arrays
        self._finalize_data()

        logger.info(
            f"Loaded {len(self._topic_data)} topics: "
            f"{len(self._image_topics)} image, "
            f"{len(self._camera_info_topics)} camera_info, "
            f"{len(self._topic_data) - len(self._image_topics) - len(self._camera_info_topics)} lowdim"
        )

    def _process_message(self, topic: str, timestamp_sec: float, msg: Any, msgtype: str):
        """Process a single message and store its data."""
        # Camera info
        if "camera_info" in topic.lower() or "CameraInfo" in msgtype:
            info = extract_camera_info(msg)
            if info is not None:
                self._camera_info_topics.add(topic)
                self._topic_data[topic]["timestamps"].append(timestamp_sec)
                self._topic_data[topic]["values"].append(info)
                return

        # TF transforms
        if topic in ["/tf", "/tf_static"]:
            if hasattr(msg, "transforms"):
                for transform_msg in msg.transforms:
                    tf_data = extract_transform(transform_msg)
                    if tf_data is not None:
                        tf_data["timestamp"] = timestamp_sec
                        tf_data["static"] = topic == "/tf_static"
                        self._tf_data.append(tf_data)
            return

        # Depth images (raw 16UC1/32FC1)
        if "depth" in topic.lower() and ("image" in topic.lower() or "Image" in msgtype):
            depth_data = extract_depth_from_msg(msg)
            if depth_data is not None:
                self._depth_topics.add(topic)
                self._topic_data[topic]["timestamps"].append(timestamp_sec)
                self._topic_data[topic]["values"].append(depth_data)
                return

        # Images
        if "image" in topic.lower() or "Image" in msgtype:
            image_data = extract_image_from_msg(msg, return_numpy=self.return_numpy_images)
            if image_data is not None:
                self._image_topics.add(topic)
                self._topic_data[topic]["timestamps"].append(timestamp_sec)
                self._topic_data[topic]["values"].append(image_data)
                return

        # Structured messages (JointState, PoseStamped)
        structured = extract_structured_msg(msg)
        if structured is not None:
            for suffix, arr in structured.items():
                key = f"{topic}{suffix}"
                self._topic_data[key]["timestamps"].append(timestamp_sec)
                self._topic_data[key]["values"].append(arr)
            return

        # Generic numeric extraction
        array = extract_array_from_msg(msg)
        if array is not None and array.size > 0:
            self._topic_data[topic]["timestamps"].append(timestamp_sec)
            self._topic_data[topic]["values"].append(array)

    def _finalize_data(self):
        """Convert lists to numpy arrays."""
        for key in self._topic_data:
            data = self._topic_data[key]
            data["timestamps"] = np.array(data["timestamps"])

            # Skip image/depth topics and camera info - keep as lists
            if key not in self._image_topics and key not in self._depth_topics and key not in self._camera_info_topics:
                if data["values"]:
                    data["values"] = np.array(data["values"])

    # ---------------------------------------------------------------------------
    # Topic discovery
    # ---------------------------------------------------------------------------

    def get_topics(self) -> List[str]:
        """Get all available topics."""
        return sorted(self._topic_data.keys())

    def get_topic_types(self) -> Dict[str, str]:
        """Get mapping of topic names to message types."""
        return dict(self._topic_types)

    def get_image_topics(self) -> List[str]:
        """Get topics containing image data."""
        return sorted(self._image_topics)

    def get_camera_info_topics(self) -> List[str]:
        """Get topics containing camera calibration info."""
        return sorted(self._camera_info_topics)

    def get_lowdim_topics(self) -> List[str]:
        """Get topics containing low-dimensional data (non-image, non-camera-info)."""
        lowdim = set(self._topic_data.keys()) - self._image_topics - self._camera_info_topics
        return sorted(lowdim)

    def get_duration(self) -> float:
        """Get total duration of recording in seconds."""
        all_times = []
        for data in self._topic_data.values():
            if len(data["timestamps"]) > 0:
                all_times.extend([data["timestamps"][0], data["timestamps"][-1]])
        if not all_times:
            return 0.0
        return max(all_times) - min(all_times)

    def get_time_range(self) -> Tuple[float, float]:
        """Get (start_time, end_time) of the recording."""
        all_times = []
        for data in self._topic_data.values():
            if len(data["timestamps"]) > 0:
                all_times.extend([data["timestamps"][0], data["timestamps"][-1]])
        if not all_times:
            return (0.0, 0.0)
        return (min(all_times), max(all_times))

    # ---------------------------------------------------------------------------
    # Image data
    # ---------------------------------------------------------------------------

    def get_images(
        self, topic: str, resample: bool = False
    ) -> Tuple[np.ndarray, Union[List[np.ndarray], List[bytes]]]:
        """
        Get images from a specific topic.

        Args:
            topic: Image topic name
            resample: If True and target_hz was set, resample to uniform frequency

        Returns:
            Tuple of (timestamps, images)
            - timestamps: np.ndarray of shape [N]
            - images: List of numpy RGB arrays [H, W, 3] or JPEG bytes
        """
        if topic not in self._image_topics:
            available = self.get_image_topics()
            raise ValueError(f"Topic '{topic}' not found in image topics. Available: {available}")

        data = self._topic_data[topic]
        timestamps = data["timestamps"]
        images = data["values"]

        if resample and self.resampler is not None:
            start, end = self.get_time_range()
            target_times = self.resampler.create_target_timeline(start, end)
            images = self.resampler.resample_nearest(timestamps, images, target_times)
            timestamps = target_times

        return timestamps, images

    def get_all_images(self, resample: bool = False) -> Dict[str, Tuple[np.ndarray, List]]:
        """
        Get images from all camera topics.

        Returns:
            Dict mapping topic names to (timestamps, images) tuples
        """
        result = {}
        for topic in self._image_topics:
            result[topic] = self.get_images(topic, resample=resample)
        return result

    def get_depth_topics(self) -> List[str]:
        """Get topics containing depth image data."""
        return sorted(self._depth_topics)

    def get_depth_images(
        self, topic: str
    ) -> Tuple[np.ndarray, List[np.ndarray]]:
        """
        Get depth images from a specific topic.

        Returns:
            Tuple of (timestamps, depth_images)
            - timestamps: np.ndarray of shape [N]
            - depth_images: List of uint16 numpy arrays [H, W] (values in mm)
        """
        if topic not in self._depth_topics:
            available = self.get_depth_topics()
            raise ValueError(f"Topic '{topic}' not found in depth topics. Available: {available}")

        data = self._topic_data[topic]
        return data["timestamps"], data["values"]

    # ---------------------------------------------------------------------------
    # Camera intrinsics and extrinsics
    # ---------------------------------------------------------------------------

    def get_intrinsics(self, topic: Optional[str] = None) -> Optional[Dict[str, Any]]:
        """
        Get camera intrinsics from a camera_info topic.

        Args:
            topic: Camera info topic. If None, returns first available.

        Returns:
            Dict with K (3x3 intrinsic matrix), D (distortion), width, height, etc.
            Returns None if not available.
        """
        if not self._camera_info_topics:
            return None

        if topic is None:
            topic = sorted(self._camera_info_topics)[0]

        if topic not in self._camera_info_topics:
            return None

        # Return the latest camera info (they're typically static)
        values = self._topic_data[topic]["values"]
        if values:
            return values[-1]
        return None

    def get_all_intrinsics(self) -> Dict[str, Dict[str, Any]]:
        """
        Get intrinsics from all camera_info topics.

        Returns:
            Dict mapping topic names to intrinsics dicts
        """
        result = {}
        for topic in self._camera_info_topics:
            intrinsics = self.get_intrinsics(topic)
            if intrinsics is not None:
                result[topic] = intrinsics
        return result

    def get_extrinsics(self, child_frame: Optional[str] = None) -> Optional[Dict[str, np.ndarray]]:
        """
        Get camera extrinsics (transform) from TF data.

        Args:
            child_frame: Child frame ID to look for. If None, returns all transforms.

        Returns:
            Dict with translation, rotation_quat, rotation_matrix, transform_matrix (4x4)
            Returns None if not available.
        """
        if not self._tf_data:
            return None

        if child_frame is None:
            # Return all unique transforms
            return self.get_all_extrinsics()

        for tf in self._tf_data:
            if tf.get("child_frame_id") == child_frame:
                return {
                    "translation": tf["translation"],
                    "rotation_quat": tf["rotation_quat"],
                    "rotation_matrix": tf["rotation_matrix"],
                    "transform_matrix": tf["transform_matrix"],
                    "frame_id": tf.get("frame_id"),
                    "child_frame_id": tf.get("child_frame_id"),
                }

        return None

    def get_all_extrinsics(self) -> Dict[str, Dict[str, np.ndarray]]:
        """
        Get all unique transforms from TF data.

        Returns:
            Dict mapping "parent_frame->child_frame" to transform dicts
        """
        result = {}
        seen = set()

        for tf in self._tf_data:
            parent = tf.get("frame_id", "unknown")
            child = tf.get("child_frame_id", "unknown")
            key = f"{parent}->{child}"

            if key not in seen:
                seen.add(key)
                result[key] = {
                    "translation": tf["translation"],
                    "rotation_quat": tf["rotation_quat"],
                    "rotation_matrix": tf["rotation_matrix"],
                    "transform_matrix": tf["transform_matrix"],
                    "frame_id": parent,
                    "child_frame_id": child,
                    "static": tf.get("static", False),
                }

        return result

    # ---------------------------------------------------------------------------
    # Low-dimensional data
    # ---------------------------------------------------------------------------

    def get_lowdim(self, topic: str, resample: bool = False) -> Tuple[np.ndarray, np.ndarray]:
        """
        Get low-dimensional data from a specific topic.

        Args:
            topic: Topic name (may include structured suffixes like "/joint_states__left_arm")
            resample: If True and target_hz was set, resample to uniform frequency

        Returns:
            Tuple of (timestamps, values)
            - timestamps: np.ndarray of shape [N]
            - values: np.ndarray of shape [N, D]
        """
        if topic not in self._topic_data:
            available = self.get_lowdim_topics()
            raise ValueError(f"Topic '{topic}' not found. Available lowdim topics: {available}")

        if topic in self._image_topics or topic in self._camera_info_topics:
            raise ValueError(f"Topic '{topic}' is not a lowdim topic. Use get_images() or get_intrinsics().")

        data = self._topic_data[topic]
        timestamps = data["timestamps"]
        values = data["values"]

        if resample and self.resampler is not None:
            start, end = self.get_time_range()
            target_times = self.resampler.create_target_timeline(start, end)
            values = self.resampler.resample_continuous(timestamps, values, target_times)
            timestamps = target_times

        return timestamps, values

    def get_all_lowdim(self, resample: bool = False) -> Dict[str, Tuple[np.ndarray, np.ndarray]]:
        """
        Get all low-dimensional data.

        Returns:
            Dict mapping topic names to (timestamps, values) tuples
        """
        result = {}
        for topic in self.get_lowdim_topics():
            result[topic] = self.get_lowdim(topic, resample=resample)
        return result

    def get_lowdim_by_prefix(self, prefix: str, resample: bool = False) -> Dict[str, Tuple[np.ndarray, np.ndarray]]:
        """
        Get all lowdim topics matching a prefix.

        Useful for getting all joint states, all poses, etc.

        Args:
            prefix: Topic prefix to match (e.g., "/joint_states", "/ee_pose")
            resample: If True and target_hz was set, resample to uniform frequency

        Returns:
            Dict mapping matching topic names to (timestamps, values) tuples
        """
        result = {}
        for topic in self.get_lowdim_topics():
            if topic.startswith(prefix):
                result[topic] = self.get_lowdim(topic, resample=resample)
        return result

    # ---------------------------------------------------------------------------
    # Convenience methods
    # ---------------------------------------------------------------------------

    def get_joint_states(self, resample: bool = False) -> Dict[str, Tuple[np.ndarray, np.ndarray]]:
        """Get all joint state data (topics containing 'joint')."""
        return self.get_lowdim_by_prefix("/joint", resample=resample)

    def get_poses(self, resample: bool = False) -> Dict[str, Tuple[np.ndarray, np.ndarray]]:
        """Get all pose data (topics containing '__xyz' or '__rot')."""
        result = {}
        for topic in self.get_lowdim_topics():
            if "__xyz" in topic or "__rot" in topic or "__quat" in topic:
                result[topic] = self.get_lowdim(topic, resample=resample)
        return result

    def get_synchronized_data(
        self,
        image_topics: Optional[List[str]] = None,
        lowdim_topics: Optional[List[str]] = None,
    ) -> Dict[str, Any]:
        """
        Get all data resampled to the same timestamps.

        Args:
            image_topics: List of image topics to include (default: all)
            lowdim_topics: List of lowdim topics to include (default: all)

        Returns:
            Dict with:
                - timestamps: np.ndarray of uniform timestamps
                - images: Dict[topic, List[images]]
                - lowdim: Dict[topic, np.ndarray]
        """
        if self.resampler is None:
            raise ValueError("Cannot synchronize without target_hz. Initialize MCAPReader with target_hz.")

        start, end = self.get_time_range()
        target_times = self.resampler.create_target_timeline(start, end)

        result = {"timestamps": target_times, "images": {}, "lowdim": {}}

        # Images
        topics = image_topics if image_topics else list(self._image_topics)
        for topic in topics:
            if topic in self._image_topics:
                _, images = self.get_images(topic, resample=True)
                result["images"][topic] = images

        # Lowdim
        topics = lowdim_topics if lowdim_topics else self.get_lowdim_topics()
        for topic in topics:
            if topic in self._topic_data and topic not in self._image_topics:
                _, values = self.get_lowdim(topic, resample=True)
                result["lowdim"][topic] = values

        return result

    def __repr__(self) -> str:
        duration = self.get_duration()
        return (
            f"MCAPReader('{self.mcap_path.name}', "
            f"duration={duration:.2f}s, "
            f"images={len(self._image_topics)}, "
            f"lowdim={len(self.get_lowdim_topics())}, "
            f"target_hz={self.target_hz})"
        )
    
if __name__ == "__main__":
    import sys

    # if len(sys.argv) < 2:
    #     print("Usage: python mcap_reader.py <path_to_mcap_file>")
    #     sys.exit(1)

    # mcap_path = sys.argv[1]
    mcap_path = '/data/cv_downloaded/MCAP/0000_20251215_134924/0000_20251215_134924_0.mcap'

    print("=" * 60)
    print(f"Reading: {mcap_path}")
    print("=" * 60)

    reader = MCAPReader(mcap_path)
    print(f"\n{reader}\n")

    # Topics summary
    print("-" * 60)
    print("ALL TOPICS:")
    print("-" * 60)
    for topic in reader.get_topics():
        print(f"  {topic}")

    print(f"\n{'-' * 60}")
    print("IMAGE TOPICS:")
    print("-" * 60)
    image_topics = reader.get_image_topics()
    if image_topics:
        for topic in image_topics:
            timestamps, images = reader.get_images(topic)
            print(f"  {topic}")
            print(f"    samples: {len(images)}, duration: {timestamps[-1] - timestamps[0]:.2f}s")
            if images and isinstance(images[0], np.ndarray):
                print(f"    shape: {images[0].shape}, dtype: {images[0].dtype}")
    else:
        print("  (none)")

    print(f"\n{'-' * 60}")
    print("LOWDIM TOPICS:")
    print("-" * 60)
    lowdim_topics = reader.get_lowdim_topics()
    if lowdim_topics:
        for topic in lowdim_topics:
            timestamps, values = reader.get_lowdim(topic)
            print(f"  {topic}")
            print(f"    samples: {len(values)}, shape: {values.shape}, duration: {timestamps[-1] - timestamps[0]:.2f}s")
    else:
        print("  (none)")

    # Camera intrinsics
    print(f"\n{'-' * 60}")
    print("CAMERA INTRINSICS:")
    print("-" * 60)
    intrinsics = reader.get_all_intrinsics()
    if intrinsics:
        for topic, info in intrinsics.items():
            print(f"  {topic}:")
            print(f"    resolution: {info.get('width')}x{info.get('height')}")
            if info.get("K") is not None:
                K = info["K"]
                print(f"    K matrix:")
                print(f"      fx={K[0, 0]:.2f}, fy={K[1, 1]:.2f}")
                print(f"      cx={K[0, 2]:.2f}, cy={K[1, 2]:.2f}")
            if info.get("D") is not None:
                print(f"    distortion: {info['D'][:5]}...")
    else:
        print("  (none)")

    # Camera extrinsics
    print(f"\n{'-' * 60}")
    print("CAMERA EXTRINSICS (TF):")
    print("-" * 60)
    extrinsics = reader.get_all_extrinsics()
    if extrinsics:
        for key, tf in extrinsics.items():
            print(f"  {key}:")
            print(f"    translation: [{tf['translation'][0]:.4f}, {tf['translation'][1]:.4f}, {tf['translation'][2]:.4f}]")
            print(f"    rotation (quat): [{tf['rotation_quat'][0]:.4f}, {tf['rotation_quat'][1]:.4f}, {tf['rotation_quat'][2]:.4f}, {tf['rotation_quat'][3]:.4f}]")
            print(f"    static: {tf.get('static', False)}")
    else:
        print("  (none)")

    print(f"\n{'=' * 60}")