"""Train a Go1 velocity policy on one terrain curriculum.

    python train.py hard --iters 10000

Checkpoints and the per-iteration curve stream to HuggingFace as training runs;
Kaggle kills sessions at 12 h and truncates logs, so nothing is written only at
the end.
"""

import argparse, csv, glob, os, re, shutil, subprocess, threading, time

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

REPO = "mitanshugoel/go1-terrain"
RUNDIR, OUT = "/tmp/run", "/kaggle/working/out"

p = argparse.ArgumentParser()
p.add_argument("terrain", choices=["flat", "rough", "hard"])
# 4000, not the shipped 10000, for two measured reasons:
#   hard terrain runs at 18k steps/s vs flat's 33k, so 10k iters is 15 h on hard
#   and Kaggle kills at 12 h -- one finished arm vs two truncated ones is not a
#   comparison. 4000 fits every arm inside the cap.
#   It also stops before the command curriculum widens at 5000, which would
#   otherwise change the task mid-run and confound the terrain comparison.
p.add_argument("--iters", type=int, default=4000)
p.add_argument("--envs", type=int, default=4096)
p.add_argument("--sync-every", type=int, default=600)
a = p.parse_args()

os.makedirs(RUNDIR, exist_ok=True)
os.makedirs(OUT, exist_ok=True)

from huggingface_hub import HfApi  # noqa: E402

api = HfApi(token=os.environ["HF_TOKEN"])
api.create_repo(REPO, repo_type="model", private=False, exist_ok=True)

CSV = f"{OUT}/curve.csv"
COLS = ["iteration", "reward", "ep_len", "steps_per_s", "fell_over",
        "value_loss", "action_std", "track_lin_vel", "err_vel_xy"]
PATS = {c: p for c, p in zip(COLS[1:], [
  r"Mean reward:\s*(-?[\d.]+)", r"Mean episode length:\s*([\d.]+)",
  r"Steps per second:\s*(\d+)", r"Episode_Termination/fell_over:\s*([\d.]+)",
  r"Mean value loss:\s*(-?[\d.]+)", r"Mean action std:\s*([\d.]+)",
  r"Episode_Reward/track_linear_velocity:\s*(-?[\d.]+)",
  r"Metrics/twist/error_vel_xy:\s*([\d.]+)"])}


def sync():
  """Mirror the whole run dir; upload_folder dedups so repeats are cheap."""
  try:
    api.upload_folder(folder_path=RUNDIR, path_in_repo=f"{a.terrain}/run",
                      repo_id=REPO, repo_type="model")
    api.upload_file(path_or_fileobj=CSV, path_in_repo=f"{a.terrain}/curve.csv",
                    repo_id=REPO, repo_type="model")
  except Exception as e:
    print("[hf]", e, flush=True)


def syncer():
  while True:
    time.sleep(a.sync_every)
    sync()


threading.Thread(target=syncer, daemon=True).start()

# Terrain is injected by patching the task config before the trainer imports it.
shutil.copy(os.path.join(os.path.dirname(__file__), "terrain.py"), "/tmp/terrain.py")
# All arms use the ROUGH task and differ only in terrain_generator. The flat
# task is not interchangeable: it deletes height_scan from the observations
# (config/go1/env_cfgs.py:288), so a flat-task policy has a smaller input
# vector and cannot be evaluated on rough. Swapping only the generator keeps
# observations, rewards and terminations identical across the matrix.
patch = f"""
import sys; sys.path.insert(0, "/tmp")
from terrain import TERRAINS
from mjlab.tasks.registry import load_env_cfg as _orig
def load_env_cfg(task_id, **kw):
    cfg = _orig(task_id, **kw)
    gen = TERRAINS["{a.terrain}"]
    gen.curriculum = True          # parity with mjlab's rough default
    cfg.scene.terrain.terrain_generator = gen
    return cfg
import mjlab.tasks.registry as R; R.load_env_cfg = load_env_cfg
import mjlab.scripts.train as T; T.load_env_cfg = load_env_cfg
T.main()
"""
open("/tmp/entry.py", "w").write(patch)

cf = open(CSV, "w", newline="")
w = csv.writer(cf)
w.writerow(COLS)
TASK = "Mjlab-Velocity-Rough-Unitree-Go1"   # same task for every arm; see above
cmd = (f"cd {RUNDIR} && python /tmp/entry.py {TASK} "
       f"--env.scene.num-envs {a.envs} --agent.max-iterations {a.iters}")
print(">>>", cmd, flush=True)

proc = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE,
                        stderr=subprocess.STDOUT, text=True, bufsize=1)
cur = {}
for line in proc.stdout:
  m = re.search(r"Learning iteration (\d+)/", line)
  if m:
    if cur:
      w.writerow([cur.get(c) for c in COLS])
      cf.flush()
    cur = {"iteration": int(m.group(1))}
  for k, pat in PATS.items():
    g = re.search(pat, line)
    if g:
      cur[k] = float(g.group(1))
  if re.search(r"Learning iteration|Mean reward|Error|Traceback", line):
    print(line.rstrip(), flush=True)

proc.wait()
if cur:
  w.writerow([cur.get(c) for c in COLS])
cf.close()
sync()
print(f"done rc={proc.returncode}  https://huggingface.co/{REPO}/tree/main/{a.terrain}",
      flush=True)