"""
clankerDiffusion — single 80GB A100 training launcher.

Target: 0.8B model (~08b config), 8k train context with NTK RoPE scaling
(rope_scale=16) so it extends to 128k at inference. Trains up to --hours
(12h/session, interruptible) and RESUMES from the latest checkpoint, so you
can chain several 12h sessions to reach the 20h budget.

Run on the A100 box (after `pip install torch ... numpy tokenizers datasets
safetensors huggingface_hub` and `export HF_TOKEN=...`):
    python a100_train.py --hours 12 --size 08b --scale 4

Data is regenerated in-container by prep.py (fast HF egress); the scale is
chosen so the 0.8B model sees enough tokens within the budget.
"""
import os
import subprocess
import argparse

HERE = os.path.dirname(os.path.abspath(__file__))
DATA_DIR = os.path.join(HERE, "data")
CKPT_DIR = os.path.join(HERE, "checkpoints")

# 0.8B config (also in modal_train.py SIZES["08b"]).
SIZES = {
    "08b": dict(d_model=1536, n_layers=20, n_heads=16, d_ff=4096, batch=12),
    "base": dict(d_model=768, n_layers=12, n_heads=12, d_ff=2048, batch=24),
    "large": dict(d_model=2048, n_layers=24, n_heads=16, d_ff=5504, batch=8),
}


def main(hours=12.0, size="08b", scale=4.0, batch=None, ckpt_every=250,
         seq_len=8192, rope_scale=16.0):
    os.makedirs(DATA_DIR, exist_ok=True)
    os.makedirs(CKPT_DIR, exist_ok=True)

    # pull latest code so cloud + local stay in sync (optional; needs HF_TOKEN)
    tok_path = os.path.join(DATA_DIR, "tokenizer.json")
    if not os.path.exists(tok_path) and os.environ.get("HF_TOKEN"):
        try:
            from huggingface_hub import hf_hub_download
            hf_hub_download(repo_id="coderofpears/clankerDiffusion-base",
                            filename="data/tokenizer.json", repo_type="model",
                            local_dir=DATA_DIR, token=os.environ.get("HF_TOKEN"))
        except Exception as e:
            print(f"[a100] tokenizer fetch skipped: {e}")

    # (re)generate data if incomplete
    meta_path = os.path.join(DATA_DIR, "meta.json")
    if not os.path.exists(meta_path):
        stale = os.path.join(DATA_DIR, "train.bin")
        if os.path.exists(stale):
            os.remove(stale)
        print(f"[a100] generating data (scale={scale}) ...")
        subprocess.run(["python", "prep.py", "--out-dir", DATA_DIR,
                         "--scale", str(scale)], check=True)

    spec = dict(SIZES.get(size, SIZES["08b"]))
    if batch:
        spec["batch"] = batch

    cmd = [
        "python", "train.py",
        "--hours", str(hours),
        "--batch", str(spec["batch"]),
        "--ckpt-every", str(ckpt_every),
        "--data-dir", DATA_DIR,
        "--ckpt-dir", CKPT_DIR,
        "--hf-repo", "coderofpears/clankerDiffusion-checkpoints",
        "--d-model", str(spec["d_model"]),
        "--n-layers", str(spec["n_layers"]),
        "--n-heads", str(spec["n_heads"]),
        "--d-ff", str(spec["d_ff"]),
        "--seq-len", str(seq_len),
        "--rope-scale", str(rope_scale),
    ]
    print("[a100] launching:", " ".join(cmd), flush=True)
    subprocess.run(cmd, check=True)
    print("[a100] session finished (resume next session with same command)",
          flush=True)


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--hours", type=float, default=12.0)
    ap.add_argument("--size", default="08b")
    ap.add_argument("--scale", type=float, default=4.0)
    ap.add_argument("--batch", type=int, default=None)
    ap.add_argument("--ckpt-every", type=int, default=250)
    ap.add_argument("--seq-len", type=int, default=8192)
    ap.add_argument("--rope-scale", type=float, default=16.0)
    a = ap.parse_args()
    main(hours=a.hours, size=a.size, scale=a.scale, batch=a.batch,
         ckpt_every=a.ckpt_every, seq_len=a.seq_len, rope_scale=a.rope_scale)