#!/usr/bin/env python
"""TidyBot (wheeled holonomic mobile manipulator) few-shot: does morphology
conditioning (C) adapt to a genuinely MOBILE morphology better than learned-ID (B)?
jaco (another arm) gave C≈B because its descriptor ≈ training arms. TidyBot is
is_mobile=1 / base_holonomic=1 — far from the fixed-arm training cluster, so this
is a sharper test. Its whole-body action is already EE-space (pos+quat+gripper),
so it fits the canonical head directly.

TidyBot has 110 episodes → FT on first N in {30, 80}, eval on episodes [95:110].
"""
from __future__ import annotations
import json, subprocess
from pathlib import Path

ROOT = Path("/home/alexw/tinyvla")
PY = str(Path.home() / "tinyvla_venv/bin/python")
DS = "heldout_tidybot"
DS_ROOT = str(Path.home() / "tinyvla_data/heldout/heldout_tidybot")
TEST_START = 95   # eval on eps [95:110]; FT on first N (<= 80), disjoint
SIZES = [30, 80]
VARIANTS = {
    "B": {"dir": "tv2_B_id_canon", "cond": "id", "space": "canonical", "emb": 5, "mkey": None},
    "C": {"dir": "tv2_C_morph_canon", "cond": "morph", "space": "canonical", "emb": 0, "mkey": "tidybot"},
}
OUT = ROOT / "outputs" / "tidybot_fewshot_results.json"


def write_config(v, spec, n):
    cfg = ROOT / "outputs" / f"tidybot_{v}_{n}.yaml"
    extra = f"    morph_key: {spec['mkey']}\n" if spec["mkey"] else ""
    cfg.write_text(f"""output_dir: {ROOT}/outputs/tidybot_{v}_{n}
wandb: null
seed: 42
morphology_descriptors: {ROOT}/configs/morphology/descriptors.yaml
resume_from: {ROOT}/outputs/{spec['dir']}/final
resume_step: 0
datasets:
  - repo_id: {DS}
    root: {DS_ROOT}
    episodes: {n}
    weight: 1.0
    embodiment_id: {spec['emb']}
{extra}policy:
  chunk_size: 50
  n_action_steps: 50
  image_size: 256
  freeze_lm: true
  freeze_vision_encoder: true
  num_embodiments: 16
  conditioning: {spec['cond']}
  action_space: {spec['space']}
batch_size: 32
grad_accum: 1
num_workers: 8
lr: 5.0e-5
warmup_steps: 100
steps: 4000
grad_clip: 10.0
log_freq: 500
save_freq: 4000
mixed_precision: bf16
""")
    return cfg


def run(cmd, log):
    with open(log, "w") as f:
        return subprocess.run(cmd, stdout=f, stderr=subprocess.STDOUT).returncode


def evalcmd(spec, ckpt_dir):
    c = [PY, str(ROOT / "scripts/eval_canonical.py"), "--checkpoint", str(ckpt_dir),
         "--dataset", DS, "--root", DS_ROOT, "--ep-start", str(TEST_START), "--episodes", "15"]
    return c + (["--morph-key", "tidybot"] if spec["mkey"] else ["--oracle-ids", "5"])


def main():
    res = {}
    for v, spec in VARIANTS.items():  # zero-shot first
        ev = ROOT / "outputs" / f"tidybot_{v}_0_eval.log"
        run(evalcmd(spec, ROOT / f"outputs/{spec['dir']}/final"), ev)
        line = [l for l in ev.read_text().splitlines() if "BEST" in l]
        res[f"{v}_0"] = line[-1] if line else "?"
        print(f"{v}_0: {res[f'{v}_0']}", flush=True); OUT.write_text(json.dumps(res, indent=1))
    for v, spec in VARIANTS.items():
        for n in SIZES:
            tag = f"{v}_{n}"; cfg = write_config(v, spec, n)
            print(f"[FT] {tag}", flush=True)
            run([PY, str(ROOT / "scripts/train.py"), "--config", str(cfg)], ROOT / f"outputs/tidybot_{tag}_ft.log")
            ckpt = ROOT / f"outputs/tidybot_{tag}/final"
            if not ckpt.exists():
                res[tag] = "no ckpt"; continue
            ev = ROOT / f"outputs/tidybot_{tag}_eval.log"
            run(evalcmd(spec, ckpt), ev)
            line = [l for l in ev.read_text().splitlines() if "BEST" in l]
            res[tag] = line[-1] if line else "?"
            print(f"{tag}: {res[tag]}", flush=True); OUT.write_text(json.dumps(res, indent=1))
    print("\n=== TIDYBOT FEW-SHOT DONE ===")
    for k, val in res.items():
        print(k, val)


if __name__ == "__main__":
    main()