#!/usr/bin/env python3
"""Export matched /image and /odom_combined samples from a ROS2 bag."""

import argparse
import bisect
import math
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Sequence, Tuple


DEFAULT_IMAGE_TOPIC = "/image"
DEFAULT_ODOM_TOPIC = "/odom_combined"
DEFAULT_OUTPUT_ROOT = Path("/home/sunrise/yiliao_ws/bag_outputs")
DEFAULT_MAX_DELTA_S = 0.2


@dataclass(frozen=True)
class OdomSample:
    stamp: float
    x: float
    y: float
    cos_yaw: float
    sin_yaw: float


@dataclass(frozen=True)
class ImageRecord:
    stamp: float
    msg: object
    msg_type: str


def stamp_to_seconds(stamp) -> float:
    return float(stamp.sec) + float(stamp.nanosec) * 1e-9


def message_time_seconds(msg, fallback_nanoseconds: int) -> float:
    header = getattr(msg, "header", None)
    stamp = getattr(header, "stamp", None)
    if stamp is not None and (getattr(stamp, "sec", 0) != 0 or getattr(stamp, "nanosec", 0) != 0):
        return stamp_to_seconds(stamp)
    return fallback_nanoseconds * 1e-9


def yaw_sin_cos_from_quaternion(q) -> Tuple[float, float]:
    siny_cosp = 2.0 * (q.w * q.z + q.x * q.y)
    cosy_cosp = 1.0 - 2.0 * (q.y * q.y + q.z * q.z)
    yaw = math.atan2(siny_cosp, cosy_cosp)
    return math.cos(yaw), math.sin(yaw)


def odom_sample_from_msg(stamp: float, msg) -> OdomSample:
    pose = msg.pose.pose
    cos_yaw, sin_yaw = yaw_sin_cos_from_quaternion(pose.orientation)
    return OdomSample(
        stamp=stamp,
        x=float(pose.position.x),
        y=float(pose.position.y),
        cos_yaw=cos_yaw,
        sin_yaw=sin_yaw,
    )


def find_nearest_odom(
    odom_samples: Sequence[OdomSample],
    image_stamp: float,
    max_delta_s: float,
) -> Optional[OdomSample]:
    if not odom_samples:
        return None

    stamps = [sample.stamp for sample in odom_samples]
    index = bisect.bisect_left(stamps, image_stamp)
    candidates = []
    if index < len(odom_samples):
        candidates.append(odom_samples[index])
    if index > 0:
        candidates.append(odom_samples[index - 1])
    nearest = min(candidates, key=lambda sample: abs(sample.stamp - image_stamp))
    if abs(nearest.stamp - image_stamp) > max_delta_s:
        return None
    return nearest


def read_bag_records(
    bag_path: Path,
    image_topic: str,
    odom_topic: str,
) -> Tuple[List[ImageRecord], List[OdomSample]]:
    from rclpy.serialization import deserialize_message
    from rosbag2_py import ConverterOptions, SequentialReader, StorageOptions
    from rosidl_runtime_py.utilities import get_message

    reader = SequentialReader()
    reader.open(
        StorageOptions(uri=str(bag_path), storage_id="sqlite3"),
        ConverterOptions(input_serialization_format="cdr", output_serialization_format="cdr"),
    )

    topic_types: Dict[str, str] = {
        topic_metadata.name: topic_metadata.type for topic_metadata in reader.get_all_topics_and_types()
    }
    required = [topic for topic in (image_topic, odom_topic) if topic not in topic_types]
    if required:
        available = ", ".join(sorted(topic_types))
        raise ValueError(f"topic(s) missing from bag: {', '.join(required)}. Available topics: {available}")

    msg_classes = {
        image_topic: get_message(topic_types[image_topic]),
        odom_topic: get_message(topic_types[odom_topic]),
    }

    image_records: List[ImageRecord] = []
    odom_samples: List[OdomSample] = []
    while reader.has_next():
        topic, data, bag_time_ns = reader.read_next()
        if topic not in msg_classes:
            continue
        msg = deserialize_message(data, msg_classes[topic])
        stamp = message_time_seconds(msg, bag_time_ns)
        if topic == image_topic:
            image_records.append(ImageRecord(stamp=stamp, msg=msg, msg_type=topic_types[image_topic]))
        elif topic == odom_topic:
            odom_samples.append(odom_sample_from_msg(stamp, msg))

    odom_samples.sort(key=lambda sample: sample.stamp)
    image_records.sort(key=lambda record: record.stamp)
    return image_records, odom_samples


def bag_output_name(bag_path: Path) -> str:
    return bag_path.resolve().name


def ensure_clean_output_dirs(image_dir: Path, odom_dir: Path, overwrite: bool) -> None:
    image_dir.mkdir(parents=True, exist_ok=True)
    odom_dir.mkdir(parents=True, exist_ok=True)
    existing = list(image_dir.glob("*.png")) + list(odom_dir.glob("*.txt"))
    if existing and not overwrite:
        raise FileExistsError(
            f"output already contains exported files: {image_dir} or {odom_dir}. "
            "Use --overwrite to replace numbered outputs."
        )
    if overwrite:
        for path in existing:
            path.unlink()


def write_image_png(msg, msg_type: str, path: Path) -> None:
    import cv2
    import numpy as np

    if msg_type == "sensor_msgs/msg/CompressedImage":
        encoded = np.frombuffer(msg.data, dtype=np.uint8)
        image = cv2.imdecode(encoded, cv2.IMREAD_UNCHANGED)
        if image is None:
            raise ValueError("failed to decode compressed image")
    elif msg_type == "sensor_msgs/msg/Image":
        from cv_bridge import CvBridge

        image = CvBridge().imgmsg_to_cv2(msg, desired_encoding="passthrough")
    else:
        raise ValueError(f"unsupported image message type: {msg_type}")

    if not cv2.imwrite(str(path), image):
        raise IOError(f"failed to write image: {path}")


def write_odom_txt(sample: OdomSample, path: Path) -> None:
    path.write_text(
        f"{sample.x:.9f} {sample.y:.9f} {sample.cos_yaw:.9f} {sample.sin_yaw:.9f}\n",
        encoding="utf-8",
    )


def export_matched_records(
    image_records: Iterable[ImageRecord],
    odom_samples: Sequence[OdomSample],
    image_dir: Path,
    odom_dir: Path,
    max_delta_s: float,
) -> Tuple[int, int]:
    exported = 0
    dropped = 0
    for record in image_records:
        odom = find_nearest_odom(odom_samples, record.stamp, max_delta_s)
        if odom is None:
            dropped += 1
            continue
        exported += 1
        stem = f"{exported:04d}"
        write_image_png(record.msg, record.msg_type, image_dir / f"{stem}.png")
        write_odom_txt(odom, odom_dir / f"{stem}.txt")
    return exported, dropped


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Export /image frames and nearest /odom_combined poses from a ROS2 bag."
    )
    parser.add_argument("bag_path", type=Path, help="Path to the ROS2 bag directory.")
    parser.add_argument("--image-topic", default=DEFAULT_IMAGE_TOPIC, help=f"Image topic, default: {DEFAULT_IMAGE_TOPIC}")
    parser.add_argument("--odom-topic", default=DEFAULT_ODOM_TOPIC, help=f"Odom topic, default: {DEFAULT_ODOM_TOPIC}")
    parser.add_argument(
        "--output-root",
        type=Path,
        default=DEFAULT_OUTPUT_ROOT,
        help=f"Output root, default: {DEFAULT_OUTPUT_ROOT}",
    )
    parser.add_argument(
        "--max-delta",
        type=float,
        default=DEFAULT_MAX_DELTA_S,
        help=f"Maximum image/odom timestamp difference in seconds, default: {DEFAULT_MAX_DELTA_S}",
    )
    parser.add_argument("--overwrite", action="store_true", help="Delete existing numbered png/txt files first.")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    bag_path = args.bag_path.expanduser().resolve()
    if not bag_path.exists():
        raise FileNotFoundError(f"bag path does not exist: {bag_path}")

    output_name = bag_output_name(bag_path)
    image_dir = args.output_root / "images" / output_name
    odom_dir = args.output_root / "odom" / output_name
    ensure_clean_output_dirs(image_dir, odom_dir, args.overwrite)

    image_records, odom_samples = read_bag_records(bag_path, args.image_topic, args.odom_topic)
    exported, dropped = export_matched_records(
        image_records=image_records,
        odom_samples=odom_samples,
        image_dir=image_dir,
        odom_dir=odom_dir,
        max_delta_s=args.max_delta,
    )
    print(f"bag: {bag_path}")
    print(f"images read: {len(image_records)}")
    print(f"odom read: {len(odom_samples)}")
    print(f"exported pairs: {exported}")
    print(f"dropped images: {dropped}")
    print(f"image output: {image_dir}")
    print(f"odom output: {odom_dir}")


if __name__ == "__main__":
    main()
