#!/usr/bin/env python3
"""Replay the published training curves into your own Weights & Biases project.

The CSVs next to this script are the full W&B history of every run in this repo. This
turns them back into real W&B runs in *your* account, so you can put your own training
next to ours on the same axes instead of eyeballing a PNG.

    pip install wandb
    python replay_to_wandb.py --project my-project

    # just the runs you care about ("91m" means every 91M run)
    python replay_to_wandb.py --project my-project --runs 91m/trunk 350m/anneal-T300000
    python replay_to_wandb.py --project my-project --runs 91m

    # log locally first, upload when you feel like it
    WANDB_MODE=offline python replay_to_wandb.py --project my-project
    wandb sync --sync-all

Each replayed run keeps its original step numbers, so an anneal branch starts at 0.8·T
rather than at 0 — it is a continuation of its trunk, not a fresh run. `tokens` is logged
as a metric too: pick it as the x-axis in the W&B UI to compare runs by data seen rather
than by step.

If you cloned the whole repo (not just `training-curves/`), each run's `config.json` is
found automatically and becomes the run's W&B config, so the config panel and W&B's
run-comparison filters work the way they would on a run you trained yourself.

No dependency but `wandb` — this is deliberately standalone, so it keeps working after
the rest of this repo has gone stale.
"""

from __future__ import annotations

import argparse
import csv
import json
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
STEP_COLUMN = "step"


def discover(curves_dir: Path) -> list[tuple[str, Path]]:
    """Every published history, as (run key, csv path). Key is "<model>/<branch>"."""
    found = []
    for path in sorted(curves_dir.glob("*/*.csv")):
        found.append((f"{path.parent.name}/{path.stem}", path))
    return found


def selected(runs: list[tuple[str, Path]], wanted: list[str] | None) -> list[tuple[str, Path]]:
    if not wanted:
        return runs
    keep = [(key, path) for key, path in runs if key in wanted or key.split("/")[0] in wanted]
    unknown = set(wanted) - {k for k, _ in runs} - {k.split("/")[0] for k, _ in runs}
    if unknown:
        raise SystemExit(f"unknown run(s): {', '.join(sorted(unknown))}")
    return keep


def config_for(repo_root: Path, key: str) -> dict:
    """The run's own config.json, if this is a full checkout of the repo.

    `91m/trunk` lives at `91m/trunk/`, `91m/anneal-T018750` at `91m/anneal/T018750/` —
    the CSV name flattens the checkpoint tree's two levels into one.
    """
    model, branch = key.split("/", 1)
    if branch == "trunk":
        path = repo_root / model / "trunk" / "config.json"
    else:
        path = repo_root / model / "anneal" / branch.removeprefix("anneal-") / "config.json"
    if not path.is_file():
        return {}
    raw = json.loads(path.read_text(encoding="utf-8"))
    # Flatten the config sections into one namespace, as the training loop logged it.
    flat: dict = {}
    for section in ("model_config", "train_config", "hardware_config", "tokenizers_config"):
        flat.update(raw.get(section, {}))
    return flat


def rows_of(path: Path) -> list[dict]:
    """History rows, with unlogged cells dropped rather than sent as empty strings.

    Two cadences share one table (evals every 500 steps, the rest every 100), so most
    rows carry only the log-cadence metrics. Logging the blanks would put nulls in W&B
    where the run simply did not measure anything.
    """
    out = []
    with path.open(newline="", encoding="utf-8") as handle:
        for raw in csv.DictReader(handle):
            row = {key: float(value) for key, value in raw.items() if value not in ("", None)}
            if STEP_COLUMN in row:
                out.append(row)
    return out


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("--project", required=True, help="W&B project to log into.")
    parser.add_argument("--entity", help="W&B entity (default: your default entity).")
    parser.add_argument(
        "--curves-dir", type=Path, default=HERE, help="Where the CSVs are (default: here)."
    )
    parser.add_argument(
        "--runs",
        nargs="+",
        metavar="MODEL[/BRANCH]",
        help='Limit to these runs, e.g. "91m/trunk" or "91m" for all 91M runs.',
    )
    parser.add_argument(
        "--name-prefix",
        default="",
        help='Prepended to each run name, e.g. "ref-" -> "ref-91m-trunk".',
    )
    parser.add_argument(
        "--dry-run", action="store_true", help="List what would be logged, log nothing."
    )
    args = parser.parse_args(argv)

    runs = selected(discover(args.curves_dir), args.runs)
    if not runs:
        raise SystemExit(f"no curve CSVs under {args.curves_dir}")

    if not args.dry_run:
        import wandb

    repo_root = args.curves_dir.parent
    for key, path in runs:
        rows = rows_of(path)
        name = args.name_prefix + key.replace("/", "-")
        print(f"  {key:26} {len(rows):>6,} steps -> {name}", flush=True)
        if args.dry_run:
            continue
        with wandb.init(
            entity=args.entity,
            project=args.project,
            name=name,
            config=config_for(repo_root, key),
            tags=["chess-autocomplete-v1", "replay"],
        ) as run:
            for row in rows:
                step = int(row.pop(STEP_COLUMN))
                run.log(row, step=step)

    print(f"\n{len(runs)} run(s) {'listed' if args.dry_run else 'logged'}")
    return 0


if __name__ == "__main__":
    sys.exit(main())