"""Cross-evaluate every policy on every terrain -> transfer matrix.

    python eval.py --policies flat rough hard --terrains flat rough hard

Each cell is one policy on one terrain: how often it falls, how long it
survives, how well it tracks the commanded velocity. Writes results/matrix.csv.
"""

import argparse, csv, json, os, re
from dataclasses import asdict

os.environ.setdefault("MUJOCO_GL", "egl")
os.environ["WANDB_MODE"] = "disabled"

import torch  # noqa: E402
from huggingface_hub import HfApi, hf_hub_download  # noqa: E402

from mjlab.envs import ManagerBasedRlEnv  # noqa: E402
from mjlab.rl import MjlabOnPolicyRunner, RslRlVecEnvWrapper  # noqa: E402
from mjlab.tasks.registry import load_env_cfg, load_rl_cfg, load_runner_cls  # noqa: E402
from mjlab.utils.torch import configure_torch_backends  # noqa: E402

from terrain import TERRAINS  # noqa: E402

REPO = "mitanshugoel/go1-terrain"
TASK = "Mjlab-Velocity-Rough-Unitree-Go1"

p = argparse.ArgumentParser()
p.add_argument("--policies", nargs="+", default=["flat", "rough", "hard"])
p.add_argument("--terrains", nargs="+", default=["flat", "rough", "hard"])
p.add_argument("--envs", type=int, default=512)
p.add_argument("--steps", type=int, default=2000)   # 2 episodes at 1000 steps
p.add_argument("--out", default="results/matrix.csv")
a = p.parse_args()

api = HfApi(token=os.environ["HF_TOKEN"])
files = api.list_repo_files(REPO, repo_type="model")


def latest_ckpt(policy):
  cand = sorted((int(m.group(1)), f) for f in files
                if (m := re.search(rf"^{policy}/run/.*model_(\d+)\.pt$", f)))
  if not cand:
    raise FileNotFoundError(f"no checkpoint for policy '{policy}'")
  return hf_hub_download(REPO, cand[-1][1], repo_type="model",
                         token=os.environ["HF_TOKEN"], local_dir="/tmp/ck")


configure_torch_backends()
dev = "cuda:0"
rows = []

for terrain in a.terrains:
  cfg = load_env_cfg(TASK, play=True)
  agent = load_rl_cfg(TASK)
  cfg.scene.num_envs = a.envs
  cfg.scene.terrain.terrain_generator = TERRAINS[terrain]
  cfg.curriculum = {}          # fixed command range across all cells
  env = ManagerBasedRlEnv(cfg=cfg, device=dev, render_mode=None)
  env = RslRlVecEnvWrapper(env, clip_actions=agent.clip_actions)
  runner = (load_runner_cls(TASK) or MjlabOnPolicyRunner)(env, asdict(agent), device=dev)

  for policy in a.policies:
    runner.load(latest_ckpt(policy), load_cfg={"actor": True}, strict=True,
                map_location=dev)
    pol = runner.get_inference_policy(device=dev)
    obs = env.get_observations()
    rew, dones, alive = [], 0, torch.zeros(a.envs, device=dev)
    lengths = []
    with torch.inference_mode():
      for _ in range(a.steps):
        obs, r, d, _ = env.step(pol(obs))
        rew.append(r.mean().item())
        alive += 1
        if d.any():
          lengths += alive[d.bool()].tolist()
          alive[d.bool()] = 0
          dones += int(d.sum())
    row = dict(policy=policy, terrain=terrain,
               mean_reward=round(sum(rew) / len(rew), 4),
               terminations=dones,
               mean_episode_len=round(sum(lengths) / len(lengths), 1) if lengths else a.steps,
               envs=a.envs, steps=a.steps)
    print(json.dumps(row), flush=True)
    rows.append(row)

  env.close()
  del env, runner
  torch.cuda.empty_cache()

os.makedirs(os.path.dirname(a.out) or ".", exist_ok=True)
with open(a.out, "w", newline="") as f:
  wr = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
  wr.writeheader()
  wr.writerows(rows)

print("\ntransfer matrix (mean episode length, higher = fewer falls)")
print(f"{'policy':>8} " + " ".join(f"{t:>10}" for t in a.terrains))
for pol in a.policies:
  cells = {r["terrain"]: r["mean_episode_len"] for r in rows if r["policy"] == pol}
  print(f"{pol:>8} " + " ".join(f"{cells.get(t, '-'):>10}" for t in a.terrains))