# Written by yams_any4d agent, Jul 2026.

"""
Raiden-recorded YAM session -> AnyData Unified converter.

Converts raiden `record.py` sessions (robot_data.npz + ZED SVO2 cameras) from
the russet bimanual YAM station into AnyData's Unified format, mirroring the
output schema of `yam_xdof.py` (same lowdim naming, same metadata layout) so
Any4D checkpoints trained on YAMxdof can run on this data unchanged.

Source data layout (downloaded from S3 raiden_datasets/raw/<task>/):
    <src>/<session>/robot_data.npz          91.5 Hz proprio + commands
    <src>/<session>/metadata.json           task name/instruction, camera list
    <src>/<session>/cameras/<cam>.svo2      ZED SVO2 per camera, ~30 fps

Camera mapping (2026-07-09 decision: use EGO, not scene, as the top view):
    ego_camera         -> top_left_camera   (primary / reference timeline)
    left_wrist_camera  -> left_wrist_camera
    right_wrist_camera -> right_wrist_camera
    scene_camera       -> (excluded)

Lowdim mapping (names must match yam_xdof.py so unified_anyact builds the
same action vector):
    states  left/right_arm_proprio = concat(joint_pos[6], joint_vel[6], joint_eff[6])
            left/right_eef_proprio = concat(gripper_pos[1], vel[1], eff[1])
    actions left/right_arm_leader  = joint_cmd[:6]
            left/right_eef_leader  = joint_cmd[6:7]

No intrinsics/extrinsics are written (YAMxdof was webbed without them too).

Prerequisites: pyzed (ZED SDK python API) importable in the environment.

Usage (on the machine holding the raw sessions, e.g. russet):
    ANYDATA_LOCAL_ROOT=$HOME/any4d_work/anydata_root \
    python anydata/converters/raiden_yam.py \
        $HOME/any4d_work/raw/flip_pink_cup --videos --num_procs 0
"""

import json
import os
import numpy as np

from glob import glob

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

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

# Raiden camera name -> canonical camera name. First entry = reference camera.
CAMERA_MAP = [
    ('ego_camera', 'top_left_camera'),
    ('left_wrist_camera', 'left_wrist_camera'),
    ('right_wrist_camera', 'right_wrist_camera'),
    # ('scene_camera', 'top_right_camera'),  # excluded: ego chosen as the top view
]

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


def normalize_ts(ts):
    """Convert int/float timestamps of unknown unit (s/ms/us/ns) to float seconds."""
    ts = np.asarray(ts, dtype=np.float64)
    m = np.median(ts[ts > 0]) if np.any(ts > 0) else 0.0
    if m > 1e17:    # nanoseconds
        return ts / 1e9
    if m > 1e14:    # microseconds
        return ts / 1e6
    if m > 1e11:    # milliseconds
        return ts / 1e3
    return ts       # seconds


def decode_svo2(svo_path):
    """Decode a ZED SVO2 file to (frames uint8 (T,H,W,3) RGB, timestamps float seconds)."""
    import pyzed.sl as sl

    init = sl.InitParameters()
    init.set_from_svo_file(svo_path)
    init.svo_real_time_mode = False
    init.depth_mode = sl.DEPTH_MODE.NONE

    cam = sl.Camera()
    status = cam.open(init)
    if status != sl.ERROR_CODE.SUCCESS:
        raise RuntimeError(f'ZED open failed on {svo_path}: {status}')

    mat = sl.Mat()
    frames, stamps = [], []
    while True:
        err = cam.grab()
        if err == sl.ERROR_CODE.END_OF_SVOFILE_REACHED:
            break
        if err != sl.ERROR_CODE.SUCCESS:
            break
        cam.retrieve_image(mat, sl.VIEW.LEFT)
        bgra = mat.get_data()
        frames.append(bgra[:, :, 2::-1].copy())    # BGRA -> RGB
        stamps.append(cam.get_timestamp(sl.TIME_REFERENCE.IMAGE).get_nanoseconds())
    cam.close()

    if not frames:
        raise RuntimeError(f'No frames decoded from {svo_path}')
    return np.stack(frames), normalize_ts(np.array(stamps, dtype=np.int64))


def nearest_indices(src_ts, dst_ts):
    """For each dst timestamp, index of nearest src timestamp (both sorted, seconds)."""
    idx = np.searchsorted(src_ts, dst_ts)
    idx = np.clip(idx, 1, len(src_ts) - 1)
    left = src_ts[idx - 1]
    right = src_ts[idx]
    idx -= (dst_ts - left) < (right - dst_ts)
    return np.clip(idx, 0, len(src_ts) - 1)

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


def get_sequences(args):
    seqs = sorted(glob(f'{args.src}/*/robot_data.npz'))
    return [os.path.dirname(s) for s in seqs]


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

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


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

    ### Load robot lowdim + session metadata
    npz = np.load(f'{seq}/robot_data.npz', allow_pickle=True)
    with open(f'{seq}/metadata.json') as f:
        session_meta = json.load(f)

    robot_ts = normalize_ts(npz['timestamps'])

    ### Decode cameras (skip any whose SVO2 is missing/broken)
    cam_frames, cam_ts, cameras = {}, {}, []
    for raiden_name, canon_name in CAMERA_MAP:
        svo = f'{seq}/cameras/{raiden_name}.svo2'
        if not os.path.isfile(svo) or os.path.getsize(svo) < 1e6:
            print(f'  [{i}] skipping {raiden_name}: missing/empty SVO2')
            continue
        print(f'  [{i}] decoding {raiden_name} -> {canon_name}...')
        frames, stamps = decode_svo2(svo)
        cam_frames[canon_name] = frames
        cam_ts[canon_name] = stamps
        cameras.append(canon_name)

    if not cameras:
        raise RuntimeError(f'No usable cameras in {seq}')

    ### Reference timeline = first camera (ego / top_left)
    ref_cam = cameras[0]
    ref_ts = cam_ts[ref_cam]
    n_frames_ref = len(ref_ts)

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

    ### Language
    task = session_meta.get('task_name', os.path.basename(os.path.dirname(seq)))
    instruction = session_meta.get('task_instruction', task.replace('_', ' '))
    language = dict(
        task=task,
        description=[instruction],
    )

    ### Precompute robot state/action arrays aligned to canonical names
    def cat(*keys):
        return np.concatenate([np.asarray(npz[k], dtype=np.float32) for k in keys], axis=1)

    states_full = {
        'left_arm_proprio':  cat('follower_l_joint_pos', 'follower_l_joint_vel', 'follower_l_joint_eff'),
        'right_arm_proprio': cat('follower_r_joint_pos', 'follower_r_joint_vel', 'follower_r_joint_eff'),
        'left_eef_proprio':  cat('follower_l_gripper_pos', 'follower_l_gripper_vel', 'follower_l_gripper_eff'),
        'right_eef_proprio': cat('follower_r_gripper_pos', 'follower_r_gripper_vel', 'follower_r_gripper_eff'),
    }
    actions_full = {
        'left_arm_leader':  np.asarray(npz['follower_l_joint_cmd'][:, :6], dtype=np.float32),
        'right_arm_leader': np.asarray(npz['follower_r_joint_cmd'][:, :6], dtype=np.float32),
        'left_eef_leader':  np.asarray(npz['follower_l_joint_cmd'][:, 6:7], dtype=np.float32),
        'right_eef_leader': np.asarray(npz['follower_r_joint_cmd'][:, 6:7], dtype=np.float32),
    }

    robot_idx_per_frame = nearest_indices(robot_ts, ref_ts)

    ### Per camera: sync RGB to reference timeline + write labels/lowdim
    for cam_name in cameras:
        dense = {label: dict() for label in dense_labels}

        if cam_name == ref_cam:
            for ref_idx in range(n_frames_ref):
                dense['rgb'][ref_idx] = cam_frames[cam_name][ref_idx]
        else:
            indices = nearest_indices(cam_ts[cam_name], ref_ts)
            for ref_idx, cam_idx in enumerate(indices):
                dense['rgb'][ref_idx] = cam_frames[cam_name][cam_idx]

        ### Lowdim per frame
        for frame_idx in range(n_frames_ref):
            frame = frame_name(frame_idx)
            prepare_lowdim(lowdim, dst, cam_name, frame)
            filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam_name}/{frame}.npz')

            ridx = robot_idx_per_frame[frame_idx]
            states = {name: arr[ridx] for name, arr in states_full.items()}
            actions = {name: arr[ridx] for name, arr in actions_full.items()}
            lowdim[filename_lowdim]['action'] = {'actions': actions, 'states': states}
            lowdim[filename_lowdim]['language'] = dict(prompt=[instruction])

        # Free frames for this camera before writing (write_labels consumes dense)
        del cam_frames[cam_name]

        ######## 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='YAMraiden',
            tags=tags,
            raw_id=seq.replace(f'{args.src}/', ''),
        ),
        labels=labels,
        cameras=cameras,
        resolution=resolution,
        num_frames=num_frames,
        framerate=int(session_meta.get('camera_fps', 30)),
        rgb=dict(extension='jpg'),
        intrinsics=None,
        extrinsics=None,
        depth=None,
        semantic=None,
        action=dict(format='joint_raw'),
        language=language,
        specific=dict(
            robot='yam',
            domain='real',
            instruction=instruction,
            station='russet',
            control=session_meta.get('control', ''),
        ),
    )
    write_json(filename, seq_metadata)

    return dst

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

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

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