"""Small PyRe proof: stdlib CPU by default, explicit PyTorch on GPU lanes."""

from __future__ import annotations

import time

# Match this value to the explicit mode selected in the PyRe desktop.
BACKEND = "cpu"  # "cpu", "cuda", "xpu", or "rocm"


def run_stdlib_cpu() -> None:
    started = time.perf_counter()
    result = sum(value * value for value in range(100_000))
    expected = 333_328_333_350_000
    if result != expected:
        raise SystemExit("CPU_PARITY_FAILED")
    elapsed = time.perf_counter() - started
    print("backend=cpu actual_device=cpu parity=PASS")
    print(f"transfer_inclusive_seconds={elapsed:.6f} result={result}")


def require_torch_device(torch, backend: str):
    if backend == "cuda":
        if torch.version.hip is not None or not torch.cuda.is_available():
            raise SystemExit("NVIDIA_CUDA_UNAVAILABLE")
        return torch.device("cuda"), torch.cuda.synchronize
    if backend == "rocm":
        if torch.version.hip is None or not torch.cuda.is_available():
            raise SystemExit("AMD_ROCM_UNAVAILABLE")
        # ROCm intentionally uses PyTorch's torch.cuda API.
        return torch.device("cuda"), torch.cuda.synchronize
    if backend == "xpu":
        if not hasattr(torch, "xpu") or not torch.xpu.is_available():
            raise SystemExit("INTEL_XPU_UNAVAILABLE")
        return torch.device("xpu"), torch.xpu.synchronize
    raise SystemExit(f"UNSUPPORTED_EXPLICIT_BACKEND:{backend}")


def run_torch_accelerator(backend: str) -> None:
    try:
        import torch
    except ImportError:
        raise SystemExit("PYTORCH_NOT_INSTALLED_IN_SELECTED_INTERPRETER") from None

    device, synchronize = require_torch_device(torch, backend)
    cpu_input = torch.arange(262_144, dtype=torch.float32, device="cpu")
    expected = cpu_input * 2.0 + 1.0

    synchronize()
    started = time.perf_counter()
    device_input = cpu_input.to(device)
    device_output = device_input * 2.0 + 1.0
    actual_device = device_output.device.type
    cpu_output = device_output.to("cpu")
    synchronize()
    elapsed = time.perf_counter() - started

    expected_device = "cuda" if backend in {"cuda", "rocm"} else "xpu"
    if actual_device != expected_device:
        raise SystemExit(
            f"DEVICE_MISMATCH:requested={backend}:actual={actual_device}"
        )
    torch.testing.assert_close(cpu_output, expected, rtol=0, atol=0)
    print(f"backend={backend} actual_device={actual_device} parity=PASS")
    print(
        f"transfer_inclusive_seconds={elapsed:.6f} "
        f"result={float(cpu_output.sum()):.1f}"
    )


def main() -> None:
    backend = BACKEND.strip().lower()
    if backend not in {"cpu", "cuda", "xpu", "rocm"}:
        raise SystemExit(f"INVALID_BACKEND:{backend}")
    if backend == "cpu":
        run_stdlib_cpu()
    else:
        run_torch_accelerator(backend)


if __name__ == "__main__":
    main()