scripts/prepare_ctrlworld_humanoid_multiview_grid.py

"""
Prepare datasets/humanoid/multiview (AIRBOT_MMK2 LeRobot) into Ctrl-World TRAINING
format for the 4-VIEW 2x2 GRID variant.

Same pipeline as prepare_ctrlworld_humanoid_multiview.py, but keeps FOUR views
(instead of dropping the 2nd external camera) ordered ROW-MAJOR for a 2x2 grid:

    grid:   [ TL  TR ]     TL = main external   (static)  cam_high_rgb  OR cam_head_rgb
            [ BL  BR ]     TR = 2nd external    (static)  cam_third_view OR cam_front_rgb
                           BL = cam_left_wrist_rgb        (dynamic)
                           BR = cam_right_wrist_rgb       (dynamic)

Humanoid camera names differ across tasks (two groups):
    Group A: cam_high_rgb, cam_third_view,  cam_left_wrist_rgb, cam_right_wrist_rgb
    Group B: cam_head_rgb, cam_front_rgb,   cam_left_wrist_rgb, cam_right_wrist_rgb
A few tasks have BOTH; the candidate lists below are ordered so Group A wins,
matching the 3-view prep script's precedence. `resolve_views` picks, per task,
the first available candidate for each of the two static slots.

The dataloader (dataset_bimanual_multiview.py, grid-aware) composes the 4 per-view
latents into a single (F, 4, 48, 80) canvas:
    view 0 -> [0:24,  0:40]  (top-left)      view 1 -> [0:24, 40:80]  (top-right)
    view 2 -> [24:48, 0:40]  (bottom-left)   view 3 -> [24:48,40:80]  (bottom-right)
so the per-view .pt files MUST be saved in this order (0,1,2,3).

Action/state condition = full 36-D observation.state, normalized by a freshly
generated 1%/99% stat.json.

Downsampling: --down-sample D (default 6 -> 30fps->5fps). Video + state downsampled
by the same factor so stored arrays are frame-aligned (dataset uses down_sample=1).

Output layout:
  {output_dir}/{name}/annotation/{split}/{id}.json
  {output_dir}/{name}/videos/{split}/{id}/{0,1,2,3}.mp4          (resized 192x320)
  {output_dir}/{name}/latent_videos/{split}/{id}/{0,1,2,3}.pt    (SVD-VAE latents)

{name} = humanoid_multiview_4view_grid, {id} = "{subset}_{task}__{episode_index:06d}".
"""

import argparse
import json
import os
from pathlib import Path

import numpy as np
import pandas as pd
import torch
import mediapy
from diffusers.models import AutoencoderKLTemporalDecoder


DATASET_BASE = "/pfss/mlde/workspaces/mlde_wsp_IAS_SAMMerge/VLA/doanh/video_world/video_gen_physics/datasets/humanoid/multiview"
OUTPUT_BASE = "/pfss/mlde/workspaces/mlde_wsp_IAS_SAMMerge/VLA/doanh/video_world/video_gen_physics/models/Ctrl-World/dataset_example"
SVD_PATH = "/pfss/mlde/workspaces/mlde_wsp_IAS_SAMMerge/VLA/doanh/video_world/video_gen_physics/checkpoints/stabilityai/stable-video-diffusion-img2vid"

# 4-view 2x2 grid mapping. Each static slot takes the first candidate present.
MAIN_VIEW_CANDIDATES = ["cam_high_rgb", "cam_head_rgb"]        # TL (static)
SECOND_VIEW_CANDIDATES = ["cam_third_view", "cam_front_rgb"]   # TR (static)
WRIST_LEFT = "cam_left_wrist_rgb"                              # BL (dynamic)
WRIST_RIGHT = "cam_right_wrist_rgb"                            # BR (dynamic)

NUM_VIEWS = 4
TARGET_H = 192
TARGET_W = 320
STATE_DIM = 36  # full observation.state


def find_tasks(subset_dir):
    tasks = []
    for d in sorted(Path(subset_dir).iterdir()):
        if d.is_dir() and (d / "data").exists():
            tasks.append(d.name)
    return tasks


def find_episodes(task_dir):
    data_dir = Path(task_dir) / "data"
    episodes = []
    for chunk_dir in sorted(data_dir.glob("chunk-*")):
        for pq in sorted(chunk_dir.glob("episode_*.parquet")):
            ep_id = int(pq.stem.split("_")[1])
            chunk_id = int(chunk_dir.name.split("-")[1])
            episodes.append((chunk_id, ep_id))
    return episodes


def available_camera_dirs(task_dir):
    """Return set of camera key names (e.g. 'cam_high_rgb') that have a video dir."""
    vid_root = Path(task_dir) / "videos"
    cams = set()
    for chunk_dir in vid_root.glob("chunk-*"):
        for cam_dir in chunk_dir.iterdir():
            if cam_dir.is_dir() and cam_dir.name.startswith("observation.images."):
                cams.add(cam_dir.name.split("observation.images.")[-1])
    return cams


def resolve_views(task_dir):
    """Return ordered list of 4 camera KEYS [main, second, left_wrist, right_wrist]
    (row-major for the 2x2 grid), or None if the task lacks any of them."""
    cams = available_camera_dirs(task_dir)
    main = next((c for c in MAIN_VIEW_CANDIDATES if c in cams), None)
    second = next((c for c in SECOND_VIEW_CANDIDATES if c in cams), None)
    if main is None or second is None or WRIST_LEFT not in cams or WRIST_RIGHT not in cams:
        return None
    return [main, second, WRIST_LEFT, WRIST_RIGHT]


def load_episode_instructions(task_dir):
    """Map episode_index -> instruction from meta/episodes.jsonl."""
    out = {}
    p = Path(task_dir) / "meta" / "episodes.jsonl"
    if p.exists():
        with open(p) as f:
            for line in f:
                obj = json.loads(line)
                tasks = obj.get("tasks") or []
                out[obj["episode_index"]] = tasks[0] if tasks else ""
    return out


def video_path(task_dir, chunk_id, cam_key, episode_id):
    return (
        Path(task_dir) / "videos" / f"chunk-{chunk_id:03d}" /
        f"observation.images.{cam_key}" / f"episode_{episode_id:06d}.mp4"
    )


def encode_view(video_file, vae, device, down_sample):
    video = mediapy.read_video(str(video_file))
    frames = torch.tensor(np.array(video)).permute(0, 3, 1, 2).float() / 255.0 * 2 - 1
    if down_sample > 1:
        frames = frames[::down_sample]
    x = torch.nn.functional.interpolate(
        frames, size=(TARGET_H, TARGET_W), mode="bilinear", align_corners=False
    )
    resized = ((x / 2.0 + 0.5).clamp(0, 1) * 255)
    resized = resized.permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8)

    x = x.to(device)
    with torch.no_grad():
        latents = []
        for i in range(0, len(x), 32):
            batch = x[i:i + 32]
            latent = vae.encode(batch).latent_dist.sample().mul_(vae.config.scaling_factor).cpu()
            latents.append(latent)
        latent = torch.cat(latents, dim=0)
    return resized, latent


def process_episode(task_dir, task_name, view_keys, chunk_id, episode_id, instruction,
                    out_root, split, vae, device, down_sample):
    parquet_path = (
        Path(task_dir) / "data" / f"chunk-{chunk_id:03d}" /
        f"episode_{episode_id:06d}.parquet"
    )
    if not parquet_path.exists():
        return None

    df = pd.read_parquet(parquet_path)
    raw_length = len(df)

    state_full = np.stack(df["observation.state"].values)  # (T, 36)
    state_ds = state_full[::down_sample]  # (n, 36)

    ep_id_str = f"{task_name}__{episode_id:06d}"

    resized_views = []
    latent_views = []
    for cam_key in view_keys:
        vf = video_path(task_dir, chunk_id, cam_key, episode_id)
        if not vf.exists():
            print(f"  Missing video: {vf}")
            return None
        resized, latent = encode_view(vf, vae, device, down_sample)
        resized_views.append(resized)
        latent_views.append(latent)

    n_video = min(v.shape[0] for v in latent_views)
    n_frames = min(n_video, len(state_ds))
    state_ds = state_ds[:n_frames]

    for view_idx in range(NUM_VIEWS):
        vid_dir = Path(out_root) / "videos" / split / ep_id_str
        vid_dir.mkdir(parents=True, exist_ok=True)
        mediapy.write_video(
            str(vid_dir / f"{view_idx}.mp4"),
            resized_views[view_idx][:n_frames],
            fps=max(1, int(round(30 / down_sample))),
        )
        lat_dir = Path(out_root) / "latent_videos" / split / ep_id_str
        lat_dir.mkdir(parents=True, exist_ok=True)
        torch.save(latent_views[view_idx][:n_frames], str(lat_dir / f"{view_idx}.pt"))

    state_list = state_ds.tolist()
    annotation = {
        "texts": [instruction or "robot manipulation task"],
        "episode_id": ep_id_str,
        "task_name": task_name,
        "raw_episode_id": episode_id,
        "view_keys": view_keys,
        "success": True,
        "video_length": n_frames,
        "state_length": n_frames,
        "raw_length": raw_length,
        "down_sample": down_sample,
        "videos": [{"video_path": f"videos/{split}/{ep_id_str}/{i}.mp4"} for i in range(NUM_VIEWS)],
        "latent_videos": [{"latent_video_path": f"latent_videos/{split}/{ep_id_str}/{i}.pt"} for i in range(NUM_VIEWS)],
        "states": state_list,
        "observation.state.qpos": state_list,  # 36-D full state used by the training dataset
    }

    anno_dir = Path(out_root) / "annotation" / split
    anno_dir.mkdir(parents=True, exist_ok=True)
    with open(anno_dir / f"{ep_id_str}.json", "w") as f:
        json.dump(annotation, f)

    return {"id": ep_id_str, "length": n_frames, "split": split, "instruction": instruction}


def choose_split(global_idx):
    return "val" if global_idx % 20 == 19 else "train"


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--subset", choices=["makovian", "non_makovian", "both"], default="both")
    parser.add_argument("--output-dir", type=str, default=OUTPUT_BASE)
    parser.add_argument("--dataset-base", type=str, default=DATASET_BASE)
    parser.add_argument("--svd-path", type=str, default=SVD_PATH)
    parser.add_argument("--down-sample", type=int, default=6,
                        help="Take every D-th frame. 6 = 30fps->5fps (default), 1 = keep 30fps.")
    parser.add_argument("--name", type=str, default="humanoid_multiview_4view_grid")
    parser.add_argument("--limit-episodes", type=int, default=None)
    args = parser.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Loading SVD VAE from {args.svd_path} on {device} ...")
    vae = AutoencoderKLTemporalDecoder.from_pretrained(args.svd_path, subfolder="vae").to(device)
    vae.requires_grad_(False)

    subsets = ["makovian", "non_makovian"] if args.subset == "both" else [args.subset]
    out_root = os.path.join(args.output_dir, args.name)

    results = []
    global_idx = 0
    for subset in subsets:
        subset_dir = os.path.join(args.dataset_base, subset)
        tasks = find_tasks(subset_dir)
        print(f"\n[{subset}] {len(tasks)} tasks")
        for task_name in tasks:
            task_dir = os.path.join(subset_dir, task_name)
            view_keys = resolve_views(task_dir)
            if view_keys is None:
                print(f"  SKIP task {task_name}: missing 4-view set (need one of "
                      f"{MAIN_VIEW_CANDIDATES} + one of {SECOND_VIEW_CANDIDATES} + both wrists)")
                continue
            instr_map = load_episode_instructions(task_dir)
            episodes = find_episodes(task_dir)
            for (chunk_id, ep_id) in episodes:
                if args.limit_episodes is not None and global_idx >= args.limit_episodes:
                    break
                split = choose_split(global_idx)
                res = process_episode(
                    task_dir, f"{subset}_{task_name}", view_keys, chunk_id, ep_id,
                    instr_map.get(ep_id, ""), out_root, split, vae, device, args.down_sample,
                )
                if res:
                    results.append(res)
                    print(f"  [{global_idx}] {res['id']} ({split}) "
                          f"views={view_keys[0]}+{view_keys[1]} -> {res['length']} frames")
                else:
                    print(f"  [{global_idx}] {subset}/{task_name} ep {ep_id:06d} -> SKIPPED")
                global_idx += 1
            if args.limit_episodes is not None and global_idx >= args.limit_episodes:
                break
        if args.limit_episodes is not None and global_idx >= args.limit_episodes:
            break

    summary = {
        "name": args.name,
        "down_sample": args.down_sample,
        "grid": "2x2 row-major [TL=main_external, TR=second_external, BL=left_wrist, BR=right_wrist]",
        "view_mapping": {
            "view0_TL": MAIN_VIEW_CANDIDATES,
            "view1_TR": SECOND_VIEW_CANDIDATES,
            "view2_BL": WRIST_LEFT,
            "view3_BR": WRIST_RIGHT,
        },
        "total_episodes": len(results),
        "n_train": sum(1 for r in results if r["split"] == "train"),
        "n_val": sum(1 for r in results if r["split"] == "val"),
        "episodes": results,
    }
    os.makedirs(out_root, exist_ok=True)
    with open(os.path.join(out_root, "preparation_summary.json"), "w") as f:
        json.dump(summary, f, indent=2)

    print(f"\nDone: {len(results)} episodes "
          f"(train={summary['n_train']}, val={summary['n_val']}) -> {out_root}")


if __name__ == "__main__":
    main()