"""Prove that a real training step runs on this machine's accelerator, and say what it cost.

Purpose: the per-track smoke test for a training environment. Picks the device the way
         every later script does, runs a short optimisation of a small stack of linear
         layers in the chosen precision, prints the loss falling, and reports peak
         device memory where the backend can report it.
Platform: all (CUDA on Tracks S and N, ROCm on Track X, MPS on Track M, or the CPU)
Minimum memory: 8 GB
Assumes: torch installed in the active environment. Nothing is downloaded and nothing
         is written to disk.

Usage: python verify-training-step.py [--steps 50] [--precision auto|bf16|fp32] [--size 2048]
"""
from __future__ import annotations

import argparse
import contextlib
import time

import torch
from torch import nn


def pick_device() -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda")
    mps = getattr(torch.backends, "mps", None)
    if mps is not None and mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def describe(device: torch.device) -> str:
    if device.type == "cuda":
        # The ROCm build of PyTorch also reports AMD GPUs through the cuda device name.
        return f"cuda / {torch.cuda.get_device_name(0)} (torch {torch.__version__})"
    if device.type == "mps":
        return f"mps / Apple silicon GPU (torch {torch.__version__})"
    return f"cpu (torch {torch.__version__})"


def bf16_supported(device: torch.device) -> bool:
    if device.type == "cuda":
        return torch.cuda.is_bf16_supported()
    return False


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--steps", type=int, default=50)
    parser.add_argument("--size", type=int, default=2048, help="width of each linear layer")
    parser.add_argument("--layers", type=int, default=8)
    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument("--precision", choices=["auto", "bf16", "fp32"], default="auto")
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args()

    torch.manual_seed(args.seed)
    device = pick_device()
    print(f"device: {describe(device)}")

    want_bf16 = args.precision == "bf16" or (args.precision == "auto" and bf16_supported(device))
    if want_bf16 and not bf16_supported(device) and device.type != "cpu":
        print("note: bfloat16 requested but this device does not report support for it; continuing anyway")
    print(f"precision: {'bfloat16 autocast' if want_bf16 else 'float32'}")

    layers: list[nn.Module] = []
    for _ in range(args.layers):
        layers += [nn.Linear(args.size, args.size), nn.GELU()]
    model = nn.Sequential(*layers, nn.Linear(args.size, 1)).to(device)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"parameters: {n_params:,}")
    print(f"weights at float32: {n_params * 4 / 1e9:.3f} GB   "
          f"gradients: {n_params * 4 / 1e9:.3f} GB   Adam states: {n_params * 8 / 1e9:.3f} GB")

    optimiser = torch.optim.AdamW(model.parameters(), lr=1e-3)
    inputs = torch.randn(args.batch_size, args.size, device=device)
    targets = torch.randn(args.batch_size, 1, device=device)
    loss_fn = nn.MSELoss()

    if device.type == "cuda":
        torch.cuda.reset_peak_memory_stats()

    # torch.amp documents autocast for device types including cuda and cpu; rather than assume
    # anything about the others, the script only enters autocast when it actually wants bfloat16.
    def precision_context():
        if not want_bf16:
            return contextlib.nullcontext()
        return torch.autocast(device_type=device.type, dtype=torch.bfloat16)

    first = last = None
    started = time.time()
    for step in range(1, args.steps + 1):
        with precision_context():
            loss = loss_fn(model(inputs), targets)   # forward
        optimiser.zero_grad(set_to_none=True)
        loss.backward()                              # backward
        optimiser.step()                             # step
        value = loss.item()
        first = value if first is None else first
        last = value
        if step == 1 or step % 10 == 0:
            print(f"step {step:3d}  loss {value:.6f}")
    if device.type == "cuda":
        torch.cuda.synchronize()
    elapsed = time.time() - started

    print(f"loss went from {first:.6f} to {last:.6f} in {args.steps} steps ({elapsed:.1f} s)")
    if device.type == "cuda":
        print(f"peak device memory: {torch.cuda.max_memory_allocated() / 1e9:.3f} GB allocated, "
              f"{torch.cuda.max_memory_reserved() / 1e9:.3f} GB reserved")
    else:
        print("peak device memory: not reported by this backend; watch the system memory monitor instead")

    if last is not None and first is not None and last >= first:
        raise SystemExit("the loss did not fall: something is wrong with this environment, not with the model")
    print("training step verified on this machine")


if __name__ == "__main__":
    main()