#!/usr/bin/env python3
"""Run one nanoGPT training run and log everything needed to reproduce it.

Usage (from 06-experiments/):
    .venv/bin/python run.py NAME CONFIG.py [key=value ...] [--note "why this run"]

Example:
    .venv/bin/python run.py spike-qk-on configs/spike.py qk_norm=True learning_rate=3e-3

Creates runs/<YYYYmmdd-HHMMSS>_<NAME>/ containing:
    run.json       what was run and how: git commit (+ dirty flag), nanoGPT upstream commit,
                   full resolved config, dataset checksums, hardware/software, outcome, summary
    config.json    every hyperparameter nanoGPT actually used (written by train.py)
    metrics.jsonl  one line per logged iteration / eval
    diff.patch     uncommitted changes at launch time (only if the repo was dirty)
    stdout.log     raw output (git-ignored)
    ckpt.pt        checkpoint (git-ignored)
and appends one row to runs/index.csv — the at-a-glance table of every run.

A run is stopped and marked "diverged" if the training loss becomes NaN/inf or
exceeds --diverge-loss (default 20; the initial loss is ~10.8 for a 50k vocab).
"""
import csv
import datetime as dt
import json
import math
import os
import platform
import re
import shlex
import socket
import subprocess
import sys
import time

HERE = os.path.dirname(os.path.abspath(__file__))
RUNS = os.path.join(HERE, "runs")
UPSTREAM_NANOGPT = "3adf61e154c3fe3fca428ad6bc3818b27a3b8291"  # see nanogpt/README-thesis.md

ITER_RE = re.compile(r"^iter (\d+): loss ([-\w.]+), time ([\d.]+)ms")
EVAL_RE = re.compile(r"^step (\d+): train loss ([-\w.]+), val loss ([-\w.]+)")
PARAMS_RE = re.compile(r"number of parameters: ([\d.]+)M")


def sh(*cmd):
    try:
        return subprocess.run(cmd, cwd=HERE, capture_output=True, text=True, check=True).stdout.strip()
    except Exception:
        return ""


def git_info(run_dir):
    sha = sh("git", "rev-parse", "HEAD")
    dirty = sh("git", "status", "--porcelain", "--untracked-files=no")
    info = {"commit": sha, "branch": sh("git", "rev-parse", "--abbrev-ref", "HEAD"), "dirty": bool(dirty)}
    if dirty:
        open(os.path.join(run_dir, "diff.patch"), "w").write(sh("git", "diff", "HEAD"))
        info["diff_file"] = "diff.patch"
    return info


def env_info():
    import torch  # imported here so `--help` works without torch

    return {
        "host": socket.gethostname(),
        "platform": platform.platform(),
        "machine": platform.machine(),
        "cpu": sh("sysctl", "-n", "machdep.cpu.brand_string") or platform.processor(),
        "python": platform.python_version(),
        "torch": torch.__version__,
        "mps_available": torch.backends.mps.is_available(),
        "cuda_available": torch.cuda.is_available(),
    }


def data_info(dataset):
    path = os.path.join(HERE, "data", dataset, "prep.json")
    info = json.load(open(path)) if os.path.exists(path) else {"warning": "no prep.json found"}
    for split in ("train", "val"):
        p = os.path.join(HERE, "data", dataset, f"{split}.bin")
        if os.path.exists(p):
            info[f"{split}_bin_bytes"] = os.path.getsize(p)
    return info


def parse_value(v):
    for cast in (int, float):
        try:
            return cast(v)
        except ValueError:
            pass
    if v in ("True", "False"):
        return v == "True"
    return v


def num(x):
    try:
        return float(x)
    except ValueError:
        return float("nan")


def main(argv):
    if len(argv) < 2 or argv[0] in ("-h", "--help"):
        sys.exit(__doc__)
    note = ""
    if "--note" in argv:
        i = argv.index("--note")
        note = argv[i + 1]
        del argv[i:i + 2]
    diverge_loss = 20.0
    if "--diverge-loss" in argv:
        i = argv.index("--diverge-loss")
        diverge_loss = float(argv[i + 1])
        del argv[i:i + 2]
    name, config_file, overrides = argv[0], argv[1], argv[2:]

    # Resolve the config the same way nanoGPT will: config file, then overrides.
    cfg = {}
    exec(open(config_file).read(), {}, cfg)
    cfg = {k: v for k, v in cfg.items() if not k.startswith("_")}
    for o in overrides:
        k, v = o.split("=", 1)
        cfg[k] = parse_value(v)

    run_id = dt.datetime.now().strftime("%Y%m%d-%H%M%S") + "_" + re.sub(r"[^\w.-]", "-", name)
    run_dir = os.path.join(RUNS, run_id)
    os.makedirs(run_dir)
    cmd = [sys.executable, os.path.join(HERE, "nanogpt", "train.py"), config_file,
           f"--out_dir={run_dir}", f"--data_root={os.path.join(HERE, 'data')}"]
    cmd += [f"--{o}" for o in overrides]

    record = {
        "run_id": run_id,
        "name": name,
        "note": note,
        "status": "running",
        "started": dt.datetime.now().astimezone().isoformat(timespec="seconds"),
        "command": " ".join(shlex.quote(c) for c in cmd),
        "config_file": os.path.relpath(config_file, HERE),
        "overrides": overrides,
        "config_requested": cfg,
        "git": git_info(run_dir),
        "nanogpt_upstream_commit": UPSTREAM_NANOGPT,
        "data": data_info(cfg.get("dataset", "?")),
        "environment": env_info(),
    }
    rec_path = os.path.join(run_dir, "run.json")

    def save():
        json.dump(record, open(rec_path, "w"), indent=2)

    save()
    print(f"run {run_id}\n  {record['command']}")

    metrics = open(os.path.join(run_dir, "metrics.jsonl"), "w")
    log = open(os.path.join(run_dir, "stdout.log"), "w")
    t0 = time.time()
    status, diverged_at, params_m = "completed", None, None
    first_loss = last_loss = None
    best_val, last_eval = math.inf, None
    proc = subprocess.Popen(cmd, cwd=HERE, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1,
                            env={**os.environ, "PYTHONUNBUFFERED": "1"})  # stream lines so divergence stops the run early
    try:
        for line in proc.stdout:
            log.write(line)
            line = line.strip()
            if m := PARAMS_RE.search(line):
                params_m = float(m.group(1))
            elif m := ITER_RE.match(line):
                it, loss = int(m.group(1)), num(m.group(2))
                metrics.write(json.dumps({"type": "train", "iter": it, "loss": loss, "ms": float(m.group(3)),
                                          "t": round(time.time() - t0, 2)}) + "\n")
                first_loss = loss if first_loss is None else first_loss
                last_loss = loss
                if not math.isfinite(loss) or loss > diverge_loss:
                    status, diverged_at = "diverged", it
                    print(f"  diverged at iter {it} (loss {loss}); stopping")
                    proc.terminate()
                    break
            elif m := EVAL_RE.match(line):
                it, tr, va = int(m.group(1)), num(m.group(2)), num(m.group(3))
                metrics.write(json.dumps({"type": "eval", "iter": it, "train_loss": tr, "val_loss": va,
                                          "t": round(time.time() - t0, 2)}) + "\n")
                last_eval = {"iter": it, "train_loss": tr, "val_loss": va}
                if math.isfinite(va):
                    best_val = min(best_val, va)
                print(f"  step {it}: train {tr:.4f}  val {va:.4f}")
            metrics.flush()
        proc.wait()
        if status == "completed" and proc.returncode != 0:
            status = "failed"
    except KeyboardInterrupt:
        proc.terminate()
        status = "interrupted"
    finally:
        metrics.close()
        log.close()

    resolved = os.path.join(run_dir, "config.json")
    cfg_used = json.load(open(resolved)) if os.path.exists(resolved) else cfg
    record.update({
        "status": status,
        "finished": dt.datetime.now().astimezone().isoformat(timespec="seconds"),
        "wall_seconds": round(time.time() - t0, 1),
        "summary": {
            "params_millions_nanogpt_count": params_m,
            "first_train_loss": first_loss,
            "last_train_loss": last_loss,
            "last_eval": last_eval,
            "best_val_loss": best_val if math.isfinite(best_val) else None,
            "diverged_at_iter": diverged_at,
            "returncode": proc.returncode,
        },
    })
    save()

    index = os.path.join(RUNS, "index.csv")
    cols = ["run_id", "status", "qk_norm", "learning_rate", "n_layer", "n_embd", "seed", "max_iters",
            "params_M", "best_val_loss", "diverged_at", "wall_s", "git_commit", "dirty", "note"]
    new = not os.path.exists(index)
    with open(index, "a", newline="") as f:
        w = csv.writer(f)
        if new:
            w.writerow(cols)
        w.writerow([run_id, status, cfg_used.get("qk_norm"), cfg_used.get("learning_rate"), cfg_used.get("n_layer"),
                    cfg_used.get("n_embd"), cfg_used.get("seed"), cfg_used.get("max_iters"), params_m,
                    record["summary"]["best_val_loss"], diverged_at, record["wall_seconds"],
                    record["git"]["commit"][:10], record["git"]["dirty"], note])
    print(f"  {status}; best val {record['summary']['best_val_loss']}; {record['wall_seconds']}s -> {os.path.relpath(run_dir, HERE)}")


if __name__ == "__main__":
    main(sys.argv[1:])