#!/usr/bin/env python
"""UAV (quadrotor) few-shot: does morphology conditioning (C) beat learned-ID (B)
when the held-out robot is a genuinely different MORPHOLOGY (flying free-body, not
another arm)? The jaco result (C≈B) is expected — jaco's descriptor ≈ training
arms. The uav descriptor is far outside the training cluster, so this is the real
test of C>B.

For each variant in {B, C} and N in {100, 500}: FT the pretrained final on uav[0:N],
eval endpoint error on uav[test:test+20]. A is skipped (native already shown to
diverge; and uav native == canonical 6-DOF so it'd be redundant).
"""
from __future__ import annotations
import json, subprocess
from pathlib import Path

ROOT = Path("/home/alexw/tinyvla")
PY = str(Path.home() / "tinyvla_venv/bin/python")
UAV = "heldout_uav"
UAV_ROOT = str(Path.home() / "tinyvla_data/heldout/heldout_uav")
TEST_START = 1250   # 1275 eps; test on last 20 (leave a gap), FT on first N
SIZES = [100, 500]
VARIANTS = {
    "B": {"dir": "tv2_B_id_canon", "cond": "id", "space": "canonical", "emb": 4, "mkey": None},
    "C": {"dir": "tv2_C_morph_canon", "cond": "morph", "space": "canonical", "emb": 0, "mkey": "uav"},
}
OUT = ROOT / "outputs" / "uav_fewshot_results.json"


def write_config(v, spec, n):
    cfg = ROOT / "outputs" / f"uav_{v}_{n}.yaml"
    extra = f"    morph_key: {spec['mkey']}\n" if spec["mkey"] else ""
    cfg.write_text(f"""output_dir: {ROOT}/outputs/uav_{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: {UAV}
    root: {UAV_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: 5000
grad_clip: 10.0
log_freq: 500
save_freq: 5000
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 main():
    res = {}
    # zero-shot (N=0) baselines first
    for v, spec in VARIANTS.items():
        ev = ROOT / "outputs" / f"uav_{v}_0_eval.log"
        cmd = [PY, str(ROOT/"scripts/eval_canonical.py"), "--checkpoint",
               str(ROOT/f"outputs/{spec['dir']}/final"), "--dataset", UAV, "--root", UAV_ROOT,
               "--ep-start", str(TEST_START), "--episodes", "20"]
        cmd += ["--morph-key", "uav"] if spec["mkey"] else ["--oracle-ids", "4"]
        print(f"[ZS] {v} N=0", flush=True); run(cmd, ev)
        line=[l for l in ev.read_text().splitlines() if "BEST" in l]
        res[f"{v}_0"]={"eval":line[-1] if line else "?"}; print("  ",res[f"{v}_0"]); OUT.write_text(json.dumps(res,indent=1))
    # few-shot
    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/uav_{tag}_ft.log")
            ckpt=ROOT/f"outputs/uav_{tag}/final"
            if not ckpt.exists(): res[tag]={"error":"no ckpt"}; continue
            ev=ROOT/f"outputs/uav_{tag}_eval.log"
            cmd=[PY,str(ROOT/"scripts/eval_canonical.py"),"--checkpoint",str(ckpt),"--dataset",UAV,"--root",UAV_ROOT,"--ep-start",str(TEST_START),"--episodes","20"]
            cmd += ["--morph-key","uav"] if spec["mkey"] else ["--oracle-ids","4"]
            print(f"[EVAL] {tag}", flush=True); run(cmd, ev)
            line=[l for l in ev.read_text().splitlines() if "BEST" in l]
            res[tag]={"eval":line[-1] if line else "?"}; print("  ",res[tag]); OUT.write_text(json.dumps(res,indent=1))
    print("\n=== UAV FEW-SHOT DONE ==="); [print(k,v) for k,v in res.items()]


if __name__ == "__main__":
    main()