# Written by yams_any4d agent, Jul 2026.

"""
Yukon-recorded single-arm YAM session -> AnyData Unified converter.

Converts yukon `record.py`-style sessions (per-frame jpg + per-frame lowdim pkl)
— e.g. the pickplace_70 cup pick-and-place set — into AnyData's Unified format,
mirroring the output schema of `yam_xdof.py` / `raiden_yam.py` so Any4D
checkpoints trained on YAMxdof run on this data unchanged.

Source data layout:
    <src>/<ep>/rgb/scene_camera/<frame>.jpg        960x540, ~15 fps
    <src>/<ep>/rgb/right_wrist_camera/<frame>.jpg  1280x720, frame-synced to scene
    <src>/<ep>/lowdim/<frame>.pkl                  joints(7), timestamp, intrinsics,
                                                   extrinsics, save_wh
    <src>/<ep>/calibration_results.json            scene extrinsics + wrist hand-eye

Camera mapping:
    scene_camera        -> top_left_camera   (primary / reference timeline)
    right_wrist_camera  -> right_wrist_camera

FPS: source is ~15 fps; Any4D trains at 10 fps, and 15/10 is not an integer
stride, so this converter RESAMPLES the reference timeline to TARGET_FPS
(nearest-frame, <=33 ms error) instead of relying on webbing stride. Web the
result with --stride 1.

Lowdim mapping (names must match yam_xdof.py so unified_anyact builds the same
action vector). Single right arm only; left-arm channels are ZERO-filled, and
only joint positions exist (no vel/eff -> zeros):
    states  right_arm_proprio = concat(joints[:6], zeros(6), zeros(6))
            right_eef_proprio = concat(joints[6:7], zeros(1), zeros(1))
            left_*_proprio    = zeros (same shapes)
    actions right_arm_leader  = joints[:6]   (teleop follower pos ~= command)
            right_eef_leader  = joints[6:7]
            left_*_leader     = zeros

No intrinsics/extrinsics are written for v1 (YAMxdof was webbed without them;
the s52h ckpt doesn't consume them). The calibration IS available per episode
— wiring it in unlocks camera-conditioned tasks later.

Usage (on the machine holding the raw episodes, e.g. yukon):
    ANYDATA_LOCAL_ROOT=$HOME/any4d_work/anydata_root \
    python anydata/converters/yukon_yam.py \
        $HOME/any4d_work/raw/pickplace_70 --videos --num_procs 8
"""

import json
import os
import numpy as np

from glob import glob
from PIL import 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

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

# Yukon camera name -> canonical camera name. First entry = reference camera.
CAMERA_MAP = [
    ('scene_camera', 'top_left_camera'),
    ('right_wrist_camera', 'right_wrist_camera'),
]

TARGET_FPS = 10.0   # Any4D training frame rate; source ~15 fps gets resampled

TASK = 'pickplace'
INSTRUCTION = 'pick up the cup and place it on the target.'

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


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}/*/lowdim'))
    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):

    ### Per-frame lowdim pkls carry timestamps + joints
    import pickle
    pkl_paths = sorted(glob(f'{seq}/lowdim/*.pkl'))
    lowdims = []
    for p in pkl_paths:
        with open(p, 'rb') as f:
            lowdims.append(pickle.load(f))

    src_ts = np.array([d['timestamp'] for d in lowdims], dtype=np.float64)
    joints = np.stack([np.asarray(d['joints'], dtype=np.float32) for d in lowdims])  # (T, 7)

    ### Frame paths per camera, index-synced to the pkls
    frame_paths = {}
    for yukon_name, canon_name in CAMERA_MAP:
        paths = sorted(glob(f'{seq}/rgb/{yukon_name}/*.jpg'))
        assert len(paths) == len(pkl_paths), \
            f'{seq}: {yukon_name} has {len(paths)} frames vs {len(pkl_paths)} lowdim'
        frame_paths[canon_name] = paths
    cameras = [c for _, c in CAMERA_MAP]

    ### Resample reference timeline to TARGET_FPS (source ~15 fps)
    dst_ts = np.arange(src_ts[0], src_ts[-1], 1.0 / TARGET_FPS)
    src_idx_per_frame = nearest_indices(src_ts, dst_ts)
    n_frames_ref = len(src_idx_per_frame)

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

    language = dict(
        task=TASK,
        description=[INSTRUCTION],
    )

    ### Precompute state/action arrays (single right arm; left zero-filled)
    zeros6 = np.zeros_like(joints[:, :6])
    zeros1 = np.zeros_like(joints[:, 6:7])
    states_full = {
        'left_arm_proprio':  np.concatenate([zeros6, zeros6, zeros6], axis=1),
        'right_arm_proprio': np.concatenate([joints[:, :6], zeros6, zeros6], axis=1),
        'left_eef_proprio':  np.concatenate([zeros1, zeros1, zeros1], axis=1),
        'right_eef_proprio': np.concatenate([joints[:, 6:7], zeros1, zeros1], axis=1),
    }
    actions_full = {
        'left_arm_leader':  zeros6,
        'right_arm_leader': joints[:, :6],
        'left_eef_leader':  zeros1,
        'right_eef_leader': joints[:, 6:7],
    }

    ### Per camera: gather resampled RGB + write labels/lowdim
    for cam_name in cameras:
        dense = {label: dict() for label in dense_labels}

        for ref_idx, src_idx in enumerate(src_idx_per_frame):
            img = np.asarray(Image.open(frame_paths[cam_name][src_idx]).convert('RGB'))
            dense['rgb'][ref_idx] = img

        ### Lowdim per frame
        for ref_idx, src_idx in enumerate(src_idx_per_frame):
            frame = frame_name(ref_idx)
            prepare_lowdim(lowdim, dst, cam_name, frame)
            filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam_name}/{frame}.npz')

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

        ######## 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', 'single_arm', 'teleop']

    filename = f'{dst}/metadata.json'
    seq_metadata = fill_metadata(
        args=args,
        info=dict(
            name='YAMyukon',
            tags=tags,
            raw_id=seq.replace(f'{args.src}/', ''),
        ),
        labels=labels,
        cameras=cameras,
        resolution=resolution,
        num_frames=num_frames,
        framerate=int(TARGET_FPS),
        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='yukon',
            control='teleop',
        ),
    )
    write_json(filename, seq_metadata)

    return dst

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

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

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