Train it

"""Pretrain from random weights, selecting a checkpoint by validation only."""

import argparse
from dataclasses import asdict
from datetime import datetime
import json
import math
import os
from pathlib import Path
import time
from zoneinfo import ZoneInfo

import torch

from .checkpoint import load_checkpoint, save_checkpoint
from .evaluation import better_validation, evaluate
from .model import next_token_loss
from .training_data import digest, load_training_data


ROOT = Path(__file__).resolve().parents[1]


def write_json(path, content):
    path.write_text(json.dumps(content, indent=2, sort_keys=True, allow_nan=False) + "\n", encoding="utf-8")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--run-dir", type=Path, default=ROOT / "runs" / "pretrain-v2")
    parser.add_argument("--initial", type=Path, default=ROOT / "checkpoints" / "initial.pt")
    parser.add_argument("--device", choices=("auto", "cpu", "mps"), default="auto")
    parser.add_argument("--epochs", type=int, default=200)
    parser.add_argument("--batch-size", type=int, default=256)
    parser.add_argument("--learning-rate", type=float, default=0.001)
    parser.add_argument("--eval-every", type=int, default=5)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--target-accuracy", type=float, default=1.0)
    parser.add_argument("--stable-checks", type=int, default=3)
    args = parser.parse_args()
    if (min(args.epochs, args.batch_size, args.eval_every, args.stable_checks) <= 0
            or not math.isfinite(args.learning_rate) or args.learning_rate <= 0
            or not 0 <= args.target_accuracy <= 1):
        parser.error("Positive epoch/batch/evaluation counts and learning rate are required; accuracy must be 0-1")
    if args.run_dir.exists():
        parser.error("Run directory exists. Choose a new --run-dir to preserve its checkpoints and log")
    if not args.initial.is_file():
        parser.error("Initial checkpoint missing. Create it with: "
                     ".venv/bin/python -m tiny_gpt.inspect_model --save-initial checkpoints/initial.pt")
    device = ("mps" if torch.backends.mps.is_available() else "cpu") if args.device == "auto" else args.device
    if device == "mps" and not torch.backends.mps.is_available():
        parser.error("MPS unavailable; use --device cpu")
    if device == "mps" and os.environ.get("PYTORCH_ENABLE_MPS_FALLBACK") == "1":
        parser.error("Disable automatic CPU fallback for a recorded MPS training run")
    torch.set_num_threads(4)
    torch.manual_seed(args.seed)
    model, initial_metadata = load_checkpoint(args.initial, device=device)
    if initial_metadata.get("arithmetic_training_steps") != 0:
        parser.error("Step 6 pretraining requires the untrained initial checkpoint")
    data = load_training_data()
    train = data["train"]
    validation = data["validation"]
    inputs = train["batch"]["input_ids"].to(device)
    targets = train["batch"]["target_ids"].to(device)
    generator = torch.Generator().manual_seed(args.seed)
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.learning_rate,
                                 betas=(0.9, 0.95), weight_decay=0.01, foreach=False)
    batches_per_epoch = math.ceil(len(train["rows"]) / args.batch_size)
    maximum_steps = batches_per_epoch * args.epochs
    warmup_steps = min(100, max(1, maximum_steps // 10))
    args.run_dir.mkdir(parents=True)
    config = {"model": asdict(model.config), "seed": args.seed, "device": device,
              "epochs_limit": args.epochs, "batch_size": args.batch_size,
              "learning_rate": args.learning_rate, "minimum_lr_fraction": 0.1,
              "warmup_steps": warmup_steps, "optimizer": "AdamW", "betas": [0.9, 0.95],
              "weight_decay": 0.01, "gradient_clip_norm": 1.0,
              "eval_every": args.eval_every, "target_accuracy": args.target_accuracy,
              "stable_checks": args.stable_checks, "loss": "all non-padding next tokens",
              "checkpoint_selection": "highest validation exact accuracy, then lowest answer NLL",
              "initial_checkpoint": str(args.initial), "initial_sha256": digest(args.initial),
              "initial_metadata": initial_metadata,
              "data": {name: {"examples": len(part["rows"]), "source_sha256": part["source_sha256"],
                              "token_sha256": part["token_sha256"]} for name, part in data.items()},
              "excluded_from_training_and_selection": ["calibration", "test", "ood"]}
    write_json(args.run_dir / "config.json", config)
    report_path = args.run_dir / "report.json"
    log_path = args.run_dir / "metrics.jsonl"
    report = {"status": "running", "started_at": datetime.now(ZoneInfo("Europe/Amsterdam")).isoformat(),
              "config": config, "final_test_evaluated": False}
    write_json(report_path, report)
    started = time.perf_counter()
    baseline = evaluate(model, validation, batch_size=args.batch_size)
    report["initial_validation"] = baseline
    print(f"Device {device}; {len(train['rows'])} training examples; {len(validation['rows'])} validation examples", flush=True)
    print(f"Before training: validation exact {baseline['accuracy']:.2%}; answer NLL {baseline['answer_nll']:.4f}", flush=True)
    best = None
    best_epoch = 0
    step = 0
    stable = 0
    try:
        for epoch in range(1, args.epochs + 1):
            model.train()
            order = torch.randperm(len(train["rows"]), generator=generator).to(device)
            accumulated_loss = 0.0
            real_tokens = 0
            for start in range(0, len(order), args.batch_size):
                if step < warmup_steps:
                    fraction = (step + 1) / warmup_steps
                else:
                    progress = (step - warmup_steps) / max(1, maximum_steps - warmup_steps - 1)
                    fraction = 0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * progress))
                lr = args.learning_rate * fraction
                for group in optimizer.param_groups:
                    group["lr"] = lr
                chosen = order[start:start + args.batch_size]
                x, y = inputs[chosen], targets[chosen]
                optimizer.zero_grad(set_to_none=True)
                loss = next_token_loss(model(x), y)
                if not bool(torch.isfinite(loss)):
                    raise RuntimeError("Training loss became non-finite")
                loss.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0, error_if_nonfinite=True)
                optimizer.step()
                count = int((y != 0).sum().cpu())
                accumulated_loss += float(loss.detach().cpu()) * count
                real_tokens += count
                step += 1
            entry = {"epoch": epoch, "step": step, "train_next_token_loss": accumulated_loss / real_tokens,
                     "learning_rate": lr, "elapsed_seconds": time.perf_counter() - started}
            if epoch == 1 or epoch % args.eval_every == 0 or epoch == args.epochs:
                metrics = evaluate(model, validation, batch_size=args.batch_size)
                entry["validation"] = metrics
                metadata = {"stage": "pretrained-addition", "epoch": epoch,
                            "arithmetic_training_steps": step, "seed": args.seed,
                            "validation": metrics, "training_config": config}
                save_checkpoint(model, args.run_dir / "last.pt", metadata=metadata)
                if better_validation(metrics, best):
                    best = metrics
                    best_epoch = epoch
                    save_checkpoint(model, args.run_dir / "best.pt", metadata=metadata)
                stable = stable + 1 if metrics["accuracy"] >= args.target_accuracy else 0
                print(f"Epoch {epoch:3d}: train loss {entry['train_next_token_loss']:.4f}; "
                      f"validation exact {metrics['correct']}/{metrics['examples']} ({metrics['accuracy']:.2%}); "
                      f"answer NLL {metrics['answer_nll']:.4f}; elapsed {entry['elapsed_seconds']:.1f}s", flush=True)
                report.update({"last_epoch": epoch, "step": step, "best_epoch": best_epoch,
                               "best_validation": best, "elapsed_seconds": time.perf_counter() - started})
                write_json(report_path, report)
            with log_path.open("a", encoding="utf-8") as file:
                file.write(json.dumps(entry, sort_keys=True, allow_nan=False) + "\n")
            if stable >= args.stable_checks:
                print(f"Stopping: validation target met at {args.stable_checks} consecutive scheduled checks.", flush=True)
                break
        best_model, metadata = load_checkpoint(args.run_dir / "best.pt", device=device)
        selected_validation = evaluate(best_model, validation, batch_size=args.batch_size)
        selected_training = evaluate(best_model, train, batch_size=args.batch_size)
        if selected_validation != best:
            # Numerical agreement is checked by counts and tolerances, not exact float dict equality.
            if selected_validation["correct"] != best["correct"] or abs(selected_validation["answer_nll"] - best["answer_nll"]) > 1e-4:
                raise RuntimeError("Reloaded checkpoint validation differs from selection results")
        report.update({"status": "complete", "finished_at": datetime.now(ZoneInfo("Europe/Amsterdam")).isoformat(),
                       "best_epoch": best_epoch, "step": step, "elapsed_seconds": time.perf_counter() - started,
                       "selected_checkpoint": str(args.run_dir / "best.pt"),
                       "selected_checkpoint_sha256": digest(args.run_dir / "best.pt"),
                       "selected_validation": selected_validation, "selected_training": selected_training,
                       "stopping_reason": "stable validation target" if stable >= args.stable_checks else "epoch limit",
                       "final_test_evaluated": False, "calibration_performed": False})
        write_json(report_path, report)
        print(f"Saved best checkpoint (epoch {best_epoch}): {args.run_dir / 'best.pt'}", flush=True)
        print(f"Selected validation accuracy: {selected_validation['accuracy']:.2%}; "
              f"seen training accuracy: {selected_training['accuracy']:.2%}", flush=True)
        print("Final test and calibration have not been evaluated.", flush=True)
    except BaseException as error:
        report.update({"status": "interrupted" if isinstance(error, KeyboardInterrupt) else "failed",
                       "error": str(error), "step": step, "elapsed_seconds": time.perf_counter() - started})
        write_json(report_path, report)
        raise


if __name__ == "__main__":
    main()