"""Open-loop evaluation of a trained TurboVLA checkpoint on the SO-101 left-arm set.

Same protocol as openpi's scripts/eval_so101_open_loop.py, GR00T's open_loop_eval.py and
the SmolVLA script, so the four numbers mean the same thing: walk each trajectory in
strides of the scored horizon, take the predicted chunk at each anchor, compare against
the ground-truth chunk, report MSE/MAE in the dataset's own action units.

The model is loaded in-process with baseframework.from_pretrained, the same call the
websocket policy server makes. predict_action returns NORMALIZED actions, so they are
unnormalized here with the run's own dataset_statistics.json, mirroring
experiments/robotwin/evaluation/model2robotwin_interface.py.

One thing is deliberately NOT copied from that interface: its `continuous_mask` is
hardcoded to RoboTwin's 14-dim layout (12 continuous + 2 binary). SO-101 is 5 continuous
arm joints plus 1 binary gripper, matching the normalization_modes declared in
experiments/so101/data_registry/data_config.py.

Usage:
  python eval_so101_open_loop.py \
      --ckpt results/Checkpoints/turbovla_so101_2gpu \
      --traj-ids 0 50 100 150 200 251 301 351 401 451 501 \
      --eval-horizon 16 --norm-std 28.67
"""

import argparse
import json
import pathlib

import numpy as np
import pandas as pd
import torch
from PIL import Image

from starVLA.model.framework.base_framework import baseframework

# Which action/state dimensions were normalized continuously (min_max) rather than
# binarized. This MUST match normalization_modes in
# experiments/so101/data_registry/data_config.py: the mask is applied both to the state
# fed in and to the actions read back, so a stale value corrupts the model's INPUT as well
# as its output -- which is how a gripper-only mismatch degraded every joint, not just the
# gripper. All six SO-101 dimensions are continuous; the gripper is an aperture, not a
# switch.
CONTINUOUS_MASK = np.array([True] * 6, dtype=bool)
BINARY_THRESHOLD = 0.49


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--ckpt", required=True, help="Run directory or a checkpoints/*.pt file")
    p.add_argument(
        "--data-root",
        default="/mnt/data/sftp/data/hunght23/DATA_ZOO/so101_leftarm_multitask_v21/"
        "hungho77/so101-left-arm-multitask",
    )
    p.add_argument("--traj-ids", type=int, nargs="+",
                   default=[0, 50, 100, 150, 200, 251, 301, 351, 401, 451, 501])
    p.add_argument("--steps", type=int, default=400)
    p.add_argument("--eval-horizon", type=int, default=16)
    p.add_argument("--norm-std", type=float, default=None,
                   help="ABSOLUTE action std; 28.67 for this dataset.")
    p.add_argument("--embodiment", default="new_embodiment")
    return p.parse_args()


def bounds(stats):
    return np.asarray(stats["max"], dtype=np.float64), np.asarray(stats["min"], dtype=np.float64)


def normalize_state(state, stats):
    hi, lo = bounds(stats)
    valid = CONTINUOUS_MASK & (hi != lo)
    out = np.where(valid, (state - lo) / np.where(hi == lo, 1.0, hi - lo) * 2 - 1, state)
    return np.where(~CONTINUOUS_MASK, (out > BINARY_THRESHOLD).astype(out.dtype), out)


def unnormalize_actions(norm_actions, stats):
    hi, lo = bounds(stats)
    norm_actions = np.clip(norm_actions, -1, 1)
    return np.where(
        CONTINUOUS_MASK,
        0.5 * (norm_actions + 1) * (hi - lo) + lo,
        (norm_actions > BINARY_THRESHOLD).astype(norm_actions.dtype),
    )


def main():
    args = parse_args()
    ckpt = pathlib.Path(args.ckpt)
    run_dir = ckpt if ckpt.is_dir() else ckpt.parent.parent

    stats_all = json.loads((run_dir / "dataset_statistics.json").read_text())[args.embodiment]
    a_stats, s_stats = stats_all["action"], stats_all["state"]

    model = baseframework.from_pretrained(str(ckpt)).to("cuda").eval()
    horizon = int(getattr(model, "action_horizon", args.eval_horizon))
    score = min(args.eval_horizon or horizon, horizon)

    root = pathlib.Path(args.data_root)
    episodes = [json.loads(l) for l in (root / "meta" / "episodes.jsonl").open()]
    tasks = {t["task_index"]: t["task"] for t in
             (json.loads(l) for l in (root / "meta" / "tasks.jsonl").open())}
    lengths = {e["episode_index"]: e["length"] for e in episodes}

    per_traj = {}
    for tid in args.traj_ids:
        if tid not in lengths:
            print(f"  traj {tid}: khong ton tai, bo qua")
            continue
        df = pd.read_parquet(root / "data" / "chunk-000" / f"episode_{tid:06d}.parquet")
        task = tasks[int(df["task_index"].iloc[0])]
        actions = np.stack(df["action"].to_numpy())          # (L, 6) absolute
        states = np.stack(df["observation.state"].to_numpy())
        n = min(args.steps, lengths[tid] - score)
        if n <= 0:
            print(f"  traj {tid}: qua ngan, bo qua")
            continue

        # PyAV, not torchcodec: av==12.3.0 is what this repo's own dataloader uses
        # (video_backend: pyav) and torchcodec is not in its environment. The whole
        # episode is decoded once into memory -- a few hundred 360x640 frames -- rather
        # than seeking per anchor, which with inter-frame compression is both slower and
        # easy to get subtly wrong.
        import av

        frames = {}
        for key in ("top", "wrist"):
            path = root / "videos" / "chunk-000" / f"observation.images.{key}" / f"episode_{tid:06d}.mp4"
            with av.open(str(path)) as container:
                frames[key] = [f.to_image() for f in container.decode(video=0)]

        errs = []
        for t in range(0, n, score):
            imgs = [frames[key][min(t, len(frames[key]) - 1)] for key in ("top", "wrist")]
            example = {
                "image": imgs,
                "lang": task,
                "state": normalize_state(states[t].astype(np.float64), s_stats),
            }
            with torch.no_grad():
                out = model.predict_action([example])
            pred = unnormalize_actions(out["normalized_actions"][0], a_stats)[:score]
            gt = actions[t : t + score]
            errs.append(pred[: len(gt)] - gt)

        e = np.concatenate(errs, axis=0)
        mse, mae = float((e**2).mean()), float(np.abs(e).mean())
        per_traj[tid] = (mse, mae)
        per_dim = np.abs(e).mean(axis=0)
        print(f"  traj {tid:3d}: MSE {mse:.6f}  MAE {mae:.6f}   ({len(e)} steps compared)")
        print(f"            MAE theo chieu: {np.round(per_dim, 3)}")

    if not per_traj:
        print("nothing evaluated")
        return
    maes = [v[1] for v in per_traj.values()]
    mses = [v[0] for v in per_traj.values()]
    print(f"\nScored horizon: {score} of {horizon} predicted steps")
    print(f"Average MSE: {np.mean(mses):.6f}")
    print(f"Average MAE: {np.mean(maes):.6f}")
    if args.norm_std:
        print(f"Normalized MAE: {np.mean(maes) / args.norm_std * 100:.2f}% of action scale "
              f"(divided by absolute action std {args.norm_std})")


if __name__ == "__main__":
    main()