import contextlib
import io
import sys
import time
import torch

BATCH = 64
IN_DIM = 128
HIDDEN = 256
OUT_DIM = 64
LR = 0.001
STEPS = 100

# Optional second argument: fp32 (default), fp16 or bf16. The 16-bit runs
# cast the same fp32 data and weights with .to(dtype) and train entirely in
# that dtype (cuBLAS still accumulates matmuls in fp32), like train_cuda.sc's
# fp16/bf16 modes.
DTYPES = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}


def run(device_name, dtype=torch.float32):
    torch.manual_seed(42)
    device = torch.device(device_name)

    X = (torch.rand(BATCH, IN_DIM, device=device) - 0.5).to(dtype)
    target = (torch.rand(BATCH, OUT_DIM, device=device) - 0.5).to(dtype)

    W1 = ((torch.rand(IN_DIM, HIDDEN, device=device) - 0.5) * 0.1).to(dtype).requires_grad_(True)
    W2 = ((torch.rand(HIDDEN, OUT_DIM, device=device) - 0.5) * 0.1).to(dtype).requires_grad_(True)

    if device_name == "mps":
        torch.mps.synchronize()
    t0 = time.perf_counter()
    last_loss = 0.0
    for step in range(STEPS):
        if W1.grad is not None:
            W1.grad = None
        if W2.grad is not None:
            W2.grad = None
        H = torch.relu(X @ W1)
        Y = H @ W2
        diff = Y - target
        loss = (diff * diff).sum()
        loss.backward()
        with torch.no_grad():
            W1 -= LR * W1.grad
            W2 -= LR * W2.grad
        last_loss = loss.item()
    if device_name == "mps":
        torch.mps.synchronize()
    t1 = time.perf_counter()
    train_ms = (t1 - t0) * 1000.0
    print(f"[{device_name}] train_ms={train_ms:.3f} last_loss={last_loss:.6f}")
    print(f"[{device_name}] throughput_samples_per_sec={(STEPS * BATCH) / (train_ms / 1000.0):.2f}")

    # Inference-only
    with torch.no_grad():
        if device_name == "mps":
            torch.mps.synchronize()
        t2 = time.perf_counter()
        checksum = 0.0
        for _ in range(1000):
            H2 = torch.relu(X @ W1)
            Y2 = H2 @ W2
            checksum += Y2[0, 0].item()
        if device_name == "mps":
            torch.mps.synchronize()
        t3 = time.perf_counter()
    inf_ms = (t3 - t2) * 1000.0
    print(f"[{device_name}] inference_ms_per_1000={inf_ms:.3f} checksum={checksum:.6f}")

if __name__ == "__main__":
    device_name = sys.argv[1] if len(sys.argv) > 1 else "cpu"
    dtype = DTYPES[sys.argv[2] if len(sys.argv) > 2 else "fp32"]
    # Warm-up pass, not reported: the framework's first backward / first
    # kernel launch pays one-time init (on the ROCm PyTorch build the first
    # CPU backward alone takes ~10 s). run() reseeds, so the timed pass
    # starts from the same weights.
    with contextlib.redirect_stdout(io.StringIO()):
        run(device_name, dtype)
    run(device_name, dtype)