# Copyright 2026 Toyota Research Institute.  All rights reserved.

"""
YAM robot MCAP-to-AnyData converter.

Converts protobuf-encoded MCAP recordings from the YAM bimanual robot into
AnyData's Unified frame format. Handles h264 video decode (via ffmpeg),
camera intrinsics (foxglove.CameraCalibration), joint-space lowdim
(RobotState/GripperState), and language annotations.

Source data layout (S3):
    s3://xdof-yam-data/<date_camera>/<skill>/<episode_uuid>/output.mcap

MCAP channels:
    - /top-left-camera/image-raw      (foxglove.CompressedVideo, h264)
    - /top-right-camera/image-raw      (foxglove.CompressedVideo, h264)
    - /left-wrist-camera/image-raw     (foxglove.CompressedVideo, h264)
    - /right-wrist-camera/image-raw    (foxglove.CompressedVideo, h264)
    - /top-left-camera/camera-info     (foxglove.CameraCalibration)
    - /top-right-camera/camera-info    (foxglove.CameraCalibration)
    - /left-arm-proprio, /right-arm-proprio   (RobotState: pos[6]+vel[6]+torque[6])
    - /left-eef-proprio, /right-eef-proprio   (GripperState: pos[1]+vel[1]+torque[1])
    - /left-arm-leader, /right-arm-leader     (RobotState: pos[6], teleop commands)
    - /left-eef-leader, /right-eef-leader     (GripperState: pos[1], teleop commands)
    - /instruction                     (task-level language, e.g. "roll the ties")
    - /subtask-annotation              (timestamped subtask labels)

All lowdim is joint-space (6-DOF arms, 1-DOF grippers).
Wrist cameras do not have calibration data; only top cameras have intrinsics.
No depth or extrinsics are available.
Framerate is ~30 Hz for both ZED and RealSense setups (ZED: 30.0, RealSense: ~29.7).

Prerequisites:
    uv pip install mcap mcap-protobuf-support
    ffmpeg and ffprobe must be on PATH

Usage:
    # 1. Download data from S3
    AWS_PROFILE=manip-cluster aws s3 sync \\
        s3://xdof-yam-data/2026_03_30_zed/roll_the_ties/ \\
        /data/cv_downloaded/YAM_MCAP/roll_the_ties/

    # 2. Convert to AnyData unified format
    uv run python anydata/converters/mcapyam.py \\
        /data/cv_downloaded/YAM_MCAP/roll_the_ties --frames

    # Convert all skills at once
    uv run python anydata/converters/mcapyam.py \\
        /data/cv_downloaded/YAM_MCAP --frames --num_procs 4

    # Key flags:
    #   --frames / --videos     Storage mode (required)
    #   --num_procs N           Parallel workers (default 16, 0 for single-threaded)
    #   --first N               Only convert first N episodes
    #   --restart               Re-process even if tmp files exist
    #   --official              Drop _debug suffix from output path
    #   --upload                Sync results to S3
    #   --local_path PATH       Override /data root (or set ANYDATA_LOCAL_ROOT)

    # Output goes to:
    #   /data/cv_unified/frames_debug/<dataset_name>/   (without --official)
    #   /data/cv_unified/frames/<dataset_name>/          (with --official)
"""

import json
import os
import subprocess
import tempfile
import numpy as np

from glob import glob
from collections import defaultdict

from anydata.utils.read import read_image
from anydata.utils.write import write_json, write_lowdim, write_labels
from anydata.converters.utils import run, add_key_to_dict, fill_metadata, parse_dst_seq, frame_name, prepare_lowdim

#######################################################

# Camera topic -> canonical camera name
# Covers both ZED (stereo top) and RealSense (single top) setups.
CAMERA_TOPICS = {
    '/top-left-camera/image-raw': 'top_left_camera',
    '/top-right-camera/image-raw': 'top_right_camera',
    '/top-camera/image-raw': 'top_camera',
    '/left-wrist-camera/image-raw': 'left_wrist_camera',
    '/right-wrist-camera/image-raw': 'right_wrist_camera',
}

# Camera info topic -> image topic it corresponds to
CAMERA_INFO_MAP = {
    '/top-left-camera/image-raw': '/top-left-camera/camera-info',
    '/top-right-camera/image-raw': '/top-right-camera/camera-info',
    '/top-camera/image-raw': '/top-camera/camera-info',
    '/left-wrist-camera/image-raw': '/left-wrist-camera/camera-info',
    '/right-wrist-camera/image-raw': '/right-wrist-camera/camera-info',
}

# Lowdim topics: organizing states (proprio) and actions (leader)
PROPRIO_TOPICS = {
    'left_arm_proprio':  '/left-arm-proprio',
    'right_arm_proprio': '/right-arm-proprio',
    'left_eef_proprio':  '/left-eef-proprio',
    'right_eef_proprio': '/right-eef-proprio',
}

LEADER_TOPICS = {
    'left_arm_leader':  '/left-arm-leader',
    'right_arm_leader': '/right-arm-leader',
    'left_eef_leader':  '/left-eef-leader',
    'right_eef_leader': '/right-eef-leader',
}

#######################################################


def decode_h264_to_jpgs(packets, dst_dir):
    """Decode h264 NAL-unit packets to individual JPEG files on disk.

    Writes to dst_dir/{frame_name(i)}.jpg. Returns (height, width, n_frames).
    """
    bitstream = b''.join(packets)

    with tempfile.NamedTemporaryFile(suffix='.h264', delete=False) as tmp:
        tmp.write(bitstream)
        tmp.flush()
        tmp_path = tmp.name

    try:
        # Probe resolution
        probe_cmd = [
            'ffprobe', '-v', 'quiet', '-print_format', 'json',
            '-show_streams', tmp_path,
        ]
        probe_result = subprocess.run(probe_cmd, capture_output=True, timeout=30)
        info = json.loads(probe_result.stdout)
        stream = next(s for s in info['streams'] if s['codec_type'] == 'video')
        width, height = int(stream['width']), int(stream['height'])

        # Decode to JPEG frames on disk (1-indexed by ffmpeg)
        os.makedirs(dst_dir, exist_ok=True)
        pattern = os.path.join(dst_dir, '%010d.jpg')
        cmd = [
            'ffmpeg', '-y', '-hide_banner', '-loglevel', 'error',
            '-i', tmp_path,
            '-q:v', '2',
            pattern,
        ]
        result = subprocess.run(cmd, capture_output=True, timeout=600)
        if result.returncode != 0:
            raise RuntimeError(f'ffmpeg h264 decode failed: {result.stderr.decode()}')
    finally:
        os.remove(tmp_path)

    # Rename from 1-indexed to 0-indexed frame_name format
    raw_files = sorted(glob(os.path.join(dst_dir, '*.jpg')))
    n_frames = len(raw_files)
    if n_frames == 0:
        raise RuntimeError(f'ffmpeg produced no frames in {dst_dir}')

    for idx, src in enumerate(raw_files):
        target = os.path.join(dst_dir, f'{frame_name(idx)}.jpg')
        if src != target:
            os.rename(src, target)

    return height, width, n_frames


def read_yam_mcap(mcap_path):
    """Read a YAM robot MCAP file and return all channels as structured data.

    Returns dict with keys:
        - camera_packets: {topic: {'timestamps': array, 'h264_data': [bytes, ...]}}
        - camera_info: {topic: calibration_dict}
        - lowdim: {topic: {'timestamps': array, 'values': array}}
        - instruction: str
        - annotations: [(timestamp, text), ...]
    """
    # anydata/converters/mcap.py shadows the 'mcap' package when this file
    # is run as __main__ (sys.path[0] = 'anydata/converters'). Temporarily
    # remove that directory from sys.path so the installed mcap package is found.
    import sys
    converters_dir = os.path.join(os.path.dirname(__file__))
    path_was_modified = False
    if converters_dir in sys.path:
        sys.path.remove(converters_dir)
        path_was_modified = True
    from mcap.reader import make_reader
    from mcap_protobuf.decoder import DecoderFactory
    if path_was_modified:
        sys.path.insert(0, converters_dir)

    camera_timestamps = defaultdict(list)
    camera_h264 = defaultdict(list)
    camera_info = {}
    lowdim = defaultdict(lambda: {'timestamps': [], 'values': []})
    instruction = ''
    annotations = []

    with open(mcap_path, 'rb') as f:
        reader = make_reader(f, decoder_factories=[DecoderFactory()])

        for schema, channel, message, decoded in reader.iter_decoded_messages():
            topic = channel.topic
            ts = message.log_time / 1e9

            if topic == '/instruction':
                instruction = decoded.data

            elif topic == '/subtask-annotation':
                annotations.append((ts, decoded.data))

            elif topic.endswith('/camera-info'):
                K = np.array(list(decoded.K), dtype=np.float64).reshape(3, 3) if len(decoded.K) == 9 else None
                camera_info[topic] = {
                    'width': decoded.width,
                    'height': decoded.height,
                    'K': K,
                    'D': np.array(list(decoded.D), dtype=np.float64) if len(decoded.D) > 0 else None,
                    'distortion_model': decoded.distortion_model,
                }

            elif topic.endswith('/image-raw'):
                camera_timestamps[topic].append(ts)
                camera_h264[topic].append(bytes(decoded.data))

            elif schema.name == 'RobotState':
                values = list(decoded.position)
                if decoded.velocity:
                    values.extend(decoded.velocity)
                if decoded.torque:
                    values.extend(decoded.torque)
                lowdim[topic]['timestamps'].append(ts)
                lowdim[topic]['values'].append(np.array(values, dtype=np.float32))

            elif schema.name == 'GripperState':
                values = list(decoded.position)
                if decoded.velocity:
                    values.extend(decoded.velocity)
                if decoded.torque:
                    values.extend(decoded.torque)
                lowdim[topic]['timestamps'].append(ts)
                lowdim[topic]['values'].append(np.array(values, dtype=np.float32))

    # Convert lowdim lists to arrays
    for topic in lowdim:
        lowdim[topic]['timestamps'] = np.array(lowdim[topic]['timestamps'])
        lowdim[topic]['values'] = np.array(lowdim[topic]['values'])

    # Pack camera data
    camera_packets = {}
    for topic in camera_timestamps:
        camera_packets[topic] = {
            'timestamps': np.array(camera_timestamps[topic]),
            'h264_data': camera_h264[topic],
        }

    return {
        'camera_packets': camera_packets,
        'camera_info': camera_info,
        'lowdim': dict(lowdim),
        'instruction': instruction,
        'annotations': annotations,
    }


#######################################################


def get_lowdim_at_time(all_lowdim, topic, target_time):
    """Get nearest value from a lowdim topic at a given timestamp."""
    if topic not in all_lowdim:
        return None
    data = all_lowdim[topic]
    timestamps, values = data['timestamps'], data['values']
    if len(timestamps) == 0:
        return None
    idx = np.argmin(np.abs(timestamps - target_time))
    return values[idx]


def get_states_and_actions(all_lowdim, target_time):
    """Build states (proprio feedback) and actions (leader commands) at a given time.

    States are the proprioceptive readings (current joint positions, velocities, torques).
    Actions are the leader/teleop commands that the robot should track.
    """
    states = {}
    for name, topic in PROPRIO_TOPICS.items():
        val = get_lowdim_at_time(all_lowdim, topic, target_time)
        if val is not None:
            states[name] = val

    actions = {}
    for name, topic in LEADER_TOPICS.items():
        val = get_lowdim_at_time(all_lowdim, topic, target_time)
        if val is not None:
            actions[name] = val

    return states, actions


#######################################################

def get_annotation_at_time(annotations, target_time):
    """Get the active subtask annotation at a given timestamp.

    Annotations are (timestamp, text) pairs marking the start of each subtask.
    Returns the text of the most recent annotation at or before target_time,
    or None if no annotation precedes the target time.
    """
    active = None
    for ts, text in annotations:
        if ts <= target_time:
            active = text
        else:
            break
    return active

#######################################################

def get_sequences(args):
    seqs = glob(f'{args.src}/**/output.mcap', recursive=True)
    return sorted(seqs)


def parse_sequence(seq, args):
    return parse_dst_seq(seq, args, remove=[-1])

#######################################################

def process_sequence(i, seq, dst, args):

    ### Read MCAP (h264 packets stay as bytes, not decoded to numpy yet)
    data = read_yam_mcap(seq)

    ### Determine cameras present in this episode
    cameras = []
    cam_topics = []
    for topic, cam_name in CAMERA_TOPICS.items():
        if topic in data['camera_packets']:
            cameras.append(cam_name)
            cam_topics.append(topic)

    if not cameras:
        raise RuntimeError(f'No camera data found in {seq}')

    ### Initialize lists and dicts
    num_frames = {cam: dict() for cam in cameras}
    resolution = {cam: dict() for cam in cameras}
    labels, lowdim = [], {}
    dense_labels = ['rgb']

    ### Language
    skill = os.path.basename(os.path.dirname(os.path.dirname(seq)))
    language = dict(
        task=skill,
        description=[skill.replace('_', ' ')],
    )
    if data['annotations']:
        language['annotations'] = [
            {'timestamp': t, 'text': text} for t, text in data['annotations']
        ]

    ### Decode h264 and write RGB per camera (one at a time to limit memory)
    from anydata.converters.utils import extract_mp4
    videos = {}
    timestamps = {}
    for topic, cam_name in zip(cam_topics, cameras):
        cam_data = data['camera_packets'][topic]
        timestamps[cam_name] = cam_data['timestamps']
        bitstream = b''.join(cam_data['h264_data'])
        with tempfile.NamedTemporaryFile(suffix='.h264', delete=False) as tmp:
            tmp.write(bitstream)
            tmp.flush()
            tmp_path = tmp.name
            videos[cam_name] = extract_mp4(tmp.name)

    ref_cam = cameras[0]
    n_frames = min(len(timestamps[ref_cam]), videos[ref_cam].shape[0])

    videos[ref_cam] = videos[ref_cam][:n_frames]
    timestamps[ref_cam] = timestamps[ref_cam][:n_frames]

    ### Use first camera as reference timeline
    ref_topic = cam_topics[0]
    ref_timestamps = timestamps[ref_cam] 

    ### Decode h264 and write RGB per camera (one at a time to limit memory)
    for topic, cam_name in zip(cam_topics, cameras):
        dense = {label: dict() for label in dense_labels}

        cam_data = data['camera_packets'][topic]
        cam_timestamps = timestamps[cam_name]
        video = videos[cam_name]

        # For non-reference cameras, sync to reference timeline
        if topic != ref_topic:
            # Build sync index: for each ref timestamp, find nearest cam timestamp
            indices = np.array([np.argmin(np.abs(cam_timestamps[:n_frames] - t)) for t in ref_timestamps])
            indices = indices[:n_frames]
            for ref_idx, cam_idx in enumerate(indices):
                dense['rgb'][ref_idx] = video[cam_idx]
        else:
            for ref_idx, cam_idx in enumerate(range(n_frames)):
                dense['rgb'][ref_idx] = video[cam_idx]

        ### Get intrinsics for this camera (if available)
        info_topic = CAMERA_INFO_MAP.get(topic)
        K = None
        if info_topic and info_topic in data['camera_info']:
            K = data['camera_info'][info_topic]['K']

        ### Write lowdim per frame
        n_ref = len(ref_timestamps)
        for frame_idx in range(n_ref):
            frame = frame_name(frame_idx)
            t = ref_timestamps[frame_idx]

            prepare_lowdim(lowdim, dst, cam_name, frame)
            filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam_name}/{frame}.npz')

            if K is not None:
                lowdim[filename_lowdim]['intrinsics'] = K

            states, actions = get_states_and_actions(data['lowdim'], t)
            lowdim[filename_lowdim]['action'] = {'actions': actions, 'states': states}

            ######## LANGUAGE (per-frame subtask annotation)
            if data['annotations']:
                annotation = get_annotation_at_time(data['annotations'], t)
                if annotation is None: annotation = ''
                lowdim[filename_lowdim]['language'] = dict(prompt=[annotation])

        # Free h264 data for this camera
        del data['camera_packets'][topic]

        ######## WRITE LABELS
        write_labels(dst, cam_name, args.storage, dense, labels, resolution, num_frames)

    ######## WRITE LOWDIM
    write_lowdim(args, dst, labels, num_frames, lowdim)

    ######## METADATA
    tags = ['dynamic', 'real', 'robotics', 'bimanual', 'teleop']

    filename = f'{dst}/metadata.json'
    seq_metadata = fill_metadata(
        args=args,
        info=dict(
            name='YAMxdof',
            tags=tags,
            raw_id=seq.replace(f'{args.src}/', ''),
        ),
        labels=labels,
        cameras=cameras,
        resolution=resolution,
        num_frames=num_frames,
        framerate=30,  # Both ZED (30.0 Hz) and RealSense (~29.7 Hz) are approx 30
        rgb=dict(extension='jpg'),
        intrinsics=dict(model='plumb_bob'),
        extrinsics=None,
        depth=None,
        semantic=None,
        action=dict(format='joint_raw'),
        language=language,
        specific=dict(
            robot='yam',
            domain='real',
            instruction=data['instruction'],
        ),
    )
    write_json(filename, seq_metadata)

    return dst

#######################################################

if __name__ == '__main__':
    converter = os.path.basename(__file__)
    run(converter, get_sequences, parse_sequence, process_sequence)

#######################################################