"""Symbioza example job: gpu-check.

Is there a CUDA GPU, is PyTorch using it, and does it compute correctly? Multiplies two seeded
2048x2048 matrices on the GPU in float32 (TF32 off) and in bfloat16 (float16 where the GPU has no
native bfloat16), and compares both with a float64 reference computed on the CPU.
Writes gpu-check.json. On any failure it writes error.json; SHA256SUMS is written on every exit.
Reports generic facts only: no device name, id, memory size, speed or host identity.
CPU_OK=1 lets it run without CUDA, for local testing only; the published spec never sets it.
"""
import hashlib
import json
import os
import re
import signal
import sys
import traceback
import warnings

warnings.filterwarnings("ignore", message=".*cuda capability")  # torch names the card in this warning
import torch  # noqa: E402

OUT = os.environ.get("ARTIFACTS_DIR", "/workspace/artifacts")
N = 2048
SEED = 1234
STAGE = "start"


class Fail(Exception):
    """A failure the job names itself: an exit code and the reason."""

    def __init__(self, code, msg):
        super().__init__(msg)
        self.code = code


def stage(name):
    global STAGE
    STAGE = name
    print(f"== {name}", flush=True)


def sha256(path):
    h = hashlib.sha256()
    with open(path, "rb") as fh:
        for block in iter(lambda: fh.read(1 << 20), b""):
            h.update(block)
    return h.hexdigest()


def write_sha256sums():
    with open(os.path.join(OUT, "SHA256SUMS"), "w") as fh:
        for name in sorted(os.listdir(OUT)):
            path = os.path.join(OUT, name)
            if name != "SHA256SUMS" and os.path.isfile(path):
                fh.write(f"{sha256(path)}  {name}\n")


def redact(text):
    return re.sub(r"\d+(\.\d+)? ?[KMGT]i?B\b", "<size>", text)  # an out-of-memory message states the card's memory


def run(main, job):
    """Run main(). On ANY failure leave error.json; write SHA256SUMS on every exit path."""
    signal.signal(signal.SIGTERM, lambda *_: sys.exit(143))
    code, err = 0, None
    try:
        main()
    except Fail as e:
        code, err = e.code, str(e)
    except SystemExit as e:
        code = e.code if isinstance(e.code, int) else 1
        err = "stopped by SIGTERM" if code == 143 else f"exit {e.code}"
    except BaseException as e:
        code, err = 1, f"{type(e).__name__}: {e}"
        print(redact(traceback.format_exc()), flush=True)
    os.makedirs(OUT, exist_ok=True)
    for name in os.listdir(OUT):
        if name.endswith(".part"):  # a save the failure interrupted
            os.remove(os.path.join(OUT, name))
    if code:
        err = redact(err)[:2000]
        print(f"{job} FAILED during '{STAGE}' (exit {code}): {err}", flush=True)
        with open(os.path.join(OUT, "error.json"), "w") as fh:
            json.dump({"ok": False, "job": job, "stage": STAGE, "error": err, "exit_code": code}, fh, indent=2)
    write_sha256sums()
    sys.exit(code)


def main():
    stage("check CUDA")
    cuda = torch.cuda.is_available()
    if not cuda and os.environ.get("CPU_OK") != "1":
        raise Fail(3, "torch.cuda.is_available() is False; this job requires a GPU")
    dev = torch.device("cuda" if cuda else "cpu")
    torch.backends.cuda.matmul.allow_tf32 = False  # float32 means float32
    torch.backends.cudnn.allow_tf32 = False
    try:
        bf16 = not cuda or torch.cuda.is_bf16_supported(including_emulation=False)
    except TypeError:  # torch older than 2.4 has no including_emulation
        bf16 = torch.cuda.is_bf16_supported()
    half = torch.bfloat16 if bf16 else torch.float16
    half_name = str(half).replace("torch.", "")

    stage("matmul against a float64 reference")
    g = torch.Generator().manual_seed(SEED)
    a, b = torch.randn(N, N, generator=g), torch.randn(N, N, generator=g)
    ref = a.double() @ b.double()
    scale = ref.abs().max().item()
    ad, bd = a.to(dev), b.to(dev)
    got = {"float32": (ad @ bd).cpu(), half_name: (ad.to(half) @ bd.to(half)).cpu()}
    limits = {"float32": 1e-4, half_name: 5e-2}  # max error relative to max |reference|
    accuracy = {}
    for k, v in got.items():
        err = v.double().sub(ref).abs().max().item()
        accuracy[k] = {"max_abs_err": round(err, 6), "rel_err": float(f"{err / scale:.3e}"),
                       "rel_limit": limits[k], "ok": err / scale < limits[k]}
    ok = all(v["ok"] for v in accuracy.values())
    accuracy["float32_vs_cpu_float32_max_abs_err"] = round(got["float32"].sub(a @ b).abs().max().item(), 6)

    stage("write report")
    report = {
        "ok": ok,
        "cuda_available": cuda,
        "device_count": torch.cuda.device_count() if cuda else 0,
        "cpu_fallback": not cuda,
        "torch": torch.__version__,
        "cuda_build": torch.version.cuda,
        "seed": SEED,
        "tf32": False,
        "accuracy": {"reference": "float64 matmul on the CPU", "n": N, **accuracy},
    }
    os.makedirs(OUT, exist_ok=True)
    with open(os.path.join(OUT, "gpu-check.json"), "w") as fh:
        json.dump(report, fh, indent=2)
    print(json.dumps(report, indent=2), flush=True)
    if not ok:
        raise Fail(4, "a matmul result is outside its accuracy limit (see gpu-check.json)")


if __name__ == "__main__":
    run(main, "gpu-check")