"""Convert raw YOR episodes (thor2 recorder format) to patch-policy training format.

Raw (per episode):
    robot_data.csv   wide tick rows, incl. frame_idx (odin) + t1_ns clock
    odin.mp4         head cam, one frame per unique image
    wrist_left/right.mp4 + *_index.csv (own clock domain, same machine)

Output (patch-policy style, like cube/pusht):
    <out-dir>/
        seq_lengths.pkl      list[int], subsampled lengths
        actions.pth          FloatTensor [N, Tmax, 58] zero-padded
        obses/episode_00000.pth ...  uint8 Tensor [T, V=3, H, W, C]
        meta.json            action column order, stride, img size, episode names

Action (58D, all commanded where available):
    [left_ee_target(7), right_ee_target(7),
     left_hand_qpos(20, wuji-commanded), right_hand_qpos(20),
     lift(1), base_pose_target(3)]

Images: center-crop square -> resize 224 (DINO/14 native -> 256 patches).
Views: V=3 [odin, wrist_left, wrist_right], timestamp-matched per tick.
"""
import argparse
import csv
import json
import pickle
from pathlib import Path

import cv2
import numpy as np
import torch

EE_L = [f"left_ee_target_wxyz_xyz_{i}" for i in range(7)]
EE_R = [f"right_ee_target_wxyz_xyz_{i}" for i in range(7)]
HAND_L = [f"left_hand_qpos_{i}" for i in range(20)]
HAND_R = [f"right_hand_qpos_{i}" for i in range(20)]
LIFT = ["lift"]
BASE = ["base_pose_target_0", "base_pose_target_1", "base_pose_target_2"]
ACTION_COLS = EE_L + EE_R + HAND_L + HAND_R + LIFT + BASE
assert len(ACTION_COLS) == 58


def read_csv_rows(path):
    with open(path, newline="") as f:
        r = csv.DictReader(f)
        return r.fieldnames, list(r)


def decode_resize(video_path, size=224):
    """Decode mp4 -> list of uint8 HxWxC resized frames (center-crop square first)."""
    cap = cv2.VideoCapture(str(video_path))
    if not cap.isOpened():
        raise RuntimeError(f"cannot open {video_path}")
    frames = []
    while True:
        ok, img = cap.read()
        if not ok:
            break
        h, w = img.shape[:2]
        side = min(h, w)
        x0 = (w - side) // 2
        y0 = (h - side) // 2
        crop = img[y0:y0 + side, x0:x0 + side]
        small = cv2.resize(crop, (size, size), interpolation=cv2.INTER_AREA)
        small = cv2.cvtColor(small, cv2.COLOR_BGR2RGB)
        frames.append(small)
    cap.release()
    return np.stack(frames)  # [F, H, W, C] uint8


def nearest_idx(sorted_ts, queries):
    """For each query, index of closest entry in sorted_ts (both int64 ns)."""
    pos = np.searchsorted(sorted_ts, queries)
    pos = np.clip(pos, 0, len(sorted_ts) - 1)
    prev = np.clip(pos - 1, 0, len(sorted_ts) - 1)
    pick_next = np.abs(sorted_ts[pos] - queries) < np.abs(sorted_ts[prev] - queries)
    return np.where(pick_next, pos, prev)


def convert_episode(epdir, size=224, stride=3):
    _, ticks = read_csv_rows(epdir / "robot_data.csv")
    _, wleft_idx = read_csv_rows(epdir / "wrist_left_index.csv")
    _, wright_idx = read_csv_rows(epdir / "wrist_right_index.csv")

    print(f"  decoding {epdir.name} ...", flush=True)
    odin = decode_resize(epdir / "odin.mp4", size)
    wl = decode_resize(epdir / "wrist_left.mp4", size)
    wr = decode_resize(epdir / "wrist_right.mp4", size)
    print(f"    odin={odin.shape} wl={wl.shape} wr={wr.shape}", flush=True)

    wleft_ts = np.array([int(r["img_recv_ts_ns"]) for r in wleft_idx], dtype=np.int64)
    wright_ts = np.array([int(r["img_recv_ts_ns"]) for r in wright_idx], dtype=np.int64)

    keep = list(range(0, len(ticks), stride))
    T = len(keep)
    V = 3
    obs = np.empty((T, V, size, size, 3), dtype=np.uint8)
    act = np.zeros((T, len(ACTION_COLS)), dtype=np.float32)

    tick_t1 = np.array([int(ticks[k]["t1_ns"]) for k in keep], dtype=np.int64)
    wl_match = nearest_idx(wleft_ts, tick_t1)
    wr_match = nearest_idx(wright_ts, tick_t1)

    prev = np.zeros(len(ACTION_COLS), dtype=np.float32)
    for j, k in enumerate(keep):
        row = ticks[k]
        fi = int(float(row["frame_idx"]))
        fi = max(0, min(fi, len(odin) - 1))
        obs[j, 0] = odin[fi]
        obs[j, 1] = wl[wl_match[j]]
        obs[j, 2] = wr[wr_match[j]]
        vals = []
        for c in ACTION_COLS:
            v = row.get(c, "")
            try:
                vals.append(float(v) if v != "" else np.nan)
            except ValueError:
                vals.append(np.nan)
        a = np.array(vals, dtype=np.float32)
        nan = np.isnan(a)
        a[nan] = prev[nan]  # forward-fill; row0 NaN -> 0
        prev = a
        act[j] = a
    return obs, act


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--raw-root", type=str,
                    default=str(Path.home() / "data/yor_episodes/5c829ca0"))
    ap.add_argument("--out-dir", type=str,
                    default=str(Path.home() / "data/yor_dataset/yor_dataset"))
    ap.add_argument("--stride", type=int, default=3)
    ap.add_argument("--img-size", type=int, default=224)
    args = ap.parse_args()

    raw_root = Path(args.raw_root)
    out = Path(args.out_dir)
    (out / "obses").mkdir(parents=True, exist_ok=True)

    epdirs = sorted(raw_root.glob("ep_*"))
    assert epdirs, f"no episodes in {raw_root}"
    print(f"found {len(epdirs)} episodes")

    all_act, lengths, names = [], [], []
    for i, epdir in enumerate(epdirs):
        obs, act = convert_episode(epdir, args.img_size, args.stride)
        torch.save(torch.from_numpy(obs),
                   out / "obses" / f"episode_{i:05d}.pth")
        all_act.append(torch.from_numpy(act))
        lengths.append(act.shape[0])
        names.append(epdir.name)
        print(f"  [{i}] {epdir.name}: T={act.shape[0]} act_mean={act.mean():.4f}")

    Tmax = max(lengths)
    A = len(ACTION_COLS)
    actions = torch.zeros(len(all_act), Tmax, A, dtype=torch.float32)
    for i, a in enumerate(all_act):
        actions[i, :a.shape[0]] = a
    torch.save(actions, out / "actions.pth")
    with open(out / "seq_lengths.pkl", "wb") as f:
        pickle.dump(lengths, f)
    with open(out / "meta.json", "w") as f:
        json.dump({
            "action_columns": ACTION_COLS,
            "action_dim": A,
            "views": 3,
            "view_names": ["odin", "wrist_left", "wrist_right"],
            "img_size": args.img_size,
            "stride": args.stride,
            "episodes": names,
            "seq_lengths": lengths,
        }, f, indent=2)
    print(f"wrote {out}: N={len(all_act)} Tmax={Tmax} A={A}")
    print(f"  actions range [{actions.min():.3f}, {actions.max():.3f}]")


if __name__ == "__main__":
    main()