"""Native torch-neuronx behaviours the Neuron port depends on (or works around).

Run on a NeuronCore (no DeepInfer import, plain torch-neuronx):

    NEURON_RT_VISIBLE_CORES=0 python tests/neuron/device/test_device_primitives.py

Each check compares the device result with the host. ``expect="ok"`` checks must pass; ``expect="known"``
checks document an observed torch-neuronx 2.14 limitation the port works around -- if one of those starts
passing, the workaround can go. Exit status is non-zero only when an ``ok`` check fails.
"""

from __future__ import annotations

import os
import sys

import torch
import torch.distributed as dist

DEV = torch.device("neuron", 0)
RESULTS: list[tuple[str, str, str]] = []


def check(name: str, expect: str, fn) -> None:
    try:
        ok = bool(fn())
        status = "ok" if ok else "DIFF"
    except Exception as e:  # noqa: BLE001
        status = f"EXC {type(e).__name__}: {str(e).splitlines()[0][:100]}"
    RESULTS.append((name, expect, status))


def same(dev: torch.Tensor, host: torch.Tensor) -> bool:
    d = dev.cpu()
    if d.dtype in (torch.float8_e4m3fn, torch.float8_e8m0fnu):
        d, host = d.view(torch.uint8), host.view(torch.uint8)
    return torch.equal(d, host)


def main() -> int:
    # torch-neuronx requires the neuron process group to exist before the runtime initializes (first device op).
    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
    os.environ.setdefault("MASTER_PORT", "29761")
    dist.init_process_group("gloo", rank=0, world_size=1)
    g = dist.new_group([0], backend="neuron")

    torch.manual_seed(0)
    x = torch.randn(11, 512, dtype=torch.bfloat16)
    xd = x.to(DEV)

    # dtypes used by the kernels
    f8 = torch.randn(64, 64).clamp(-400, 400).to(torch.float8_e4m3fn)
    check("fp8 e4m3 H2D/D2H bit-exact", "ok", lambda: same(f8.to(DEV), f8))
    check("fp8 -> bf16 upcast on device", "ok", lambda: same(f8.to(DEV).to(torch.bfloat16), f8.to(torch.bfloat16)))
    e8 = torch.randint(100, 140, (8, 8), dtype=torch.uint8).view(torch.float8_e8m0fnu)
    check("e8m0 H2D bit-exact", "ok", lambda: same(e8.to(DEV), e8))
    f8b = torch.randn(4, 8).clamp(-400, 400).to(torch.float8_e4m3fn)
    e8b = torch.randint(110, 130, (4, 8), dtype=torch.uint8).view(torch.float8_e8m0fnu)
    check("torch.cat of fp8 e4m3 on device", "known", lambda: same(torch.cat([f8.to(DEV)[:4, :8], f8b.to(DEV)]), torch.cat([f8[:4, :8], f8b])))
    check("torch.cat of e8m0 on device", "known", lambda: same(torch.cat([e8.to(DEV)[:4], e8b.to(DEV)]), torch.cat([e8[:4], e8b])))
    check("torch.cat of fp8 via uint8 views", "ok", lambda: same(
        torch.cat([e8.to(DEV)[:4].view(torch.uint8), e8b.to(DEV).view(torch.uint8)]).view(torch.float8_e8m0fnu), torch.cat([e8[:4], e8b])))
    b = torch.randint(-128, 127, (4, 16), dtype=torch.int8)
    check("int8 nibble unpack", "ok", lambda: same(b.to(DEV) & 0xF, b & 0xF) and same((b.to(DEV) >> 4) & 0xF, (b >> 4) & 0xF))
    s = torch.randint(118, 130, (4, 16), dtype=torch.uint8)
    check("ue8m0 bits -> fp32 (<<23, view)", "ok",
          lambda: same((s.to(DEV).to(torch.int32) << 23).view(torch.float32), (s.to(torch.int32) << 23).view(torch.float32)))
    u = torch.tensor([1, 2**40 + 7, 2**63 + 5], dtype=torch.uint64)
    check("uint64 round trip", "known", lambda: same(u.to(DEV), u))
    i64 = torch.tensor([[2, 0, 1, 7]])
    m64 = torch.tensor([[True, True, True, False]])
    big = torch.iinfo(torch.int64).max
    check("int64.max sentinel through torch.where (implicit int32 downcast)", "known", lambda: same(
        torch.where(m64.to(DEV), i64.to(DEV), torch.full_like(i64.to(DEV), big)), torch.where(m64, i64, torch.full_like(i64, big))))
    check("uint64 add", "known", lambda: same(u.to(DEV) + 3, u + 3))

    # views / in-place semantics
    check("column-slice view -> host", "ok", lambda: same(xd[:, 256:], x[:, 256:]))
    check("transpose view -> host", "ok", lambda: same(xd.t(), x.t()))

    def inplace_slice():
        o, oc = torch.zeros(11, 1024, dtype=torch.bfloat16, device=DEV), torch.zeros(11, 1024, dtype=torch.bfloat16)
        o[:, 512:].copy_(xd)
        oc[:, 512:].copy_(x)
        o[:, :512] += xd
        oc[:, :512] += x
        return same(o, oc)

    check("copy_ / += into column slice", "ok", inplace_slice)
    buf = torch.randint(0, 256, (2 * 902880,), dtype=torch.uint8)
    bufd = buf.to(DEV)
    blk, row = torch.tensor([0, 0, 1, 1, 0]), torch.tensor([0, 5, 7, 128, 3])

    def paged(t: torch.Tensor) -> torch.Tensor:  # KV-cache style strided record view
        return t.as_strided((2, 129, 512), (902880, 528, 1))

    # Same layout without the e4m3 NaN encodings (0x7F/0xFF), so torch.equal compares real values.
    nbuf = torch.where((buf & 0x7F) == 0x7F, torch.zeros_like(buf), buf)
    nbufd = nbuf.to(DEV)
    check("fp8 strided view -> bf16 cast (KV record layout)", "ok",
          lambda: same(paged(nbufd).view(torch.float8_e4m3fn).to(torch.bfloat16), paged(nbuf).view(torch.float8_e4m3fn).to(torch.bfloat16)))
    check("fp8 strided view -> fp32 cast", "ok",
          lambda: same(paged(nbufd)[:, :7].view(torch.float8_e4m3fn).float(), paged(nbuf)[:, :7].view(torch.float8_e4m3fn).float()))
    check("uint8 strided view advanced indexing", "ok",
          lambda: same(paged(bufd)[blk.to(DEV), row.to(DEV)], paged(buf)[blk, row]))
    check("fp8 strided view advanced indexing (incl. NaN bytes)", "known",
          lambda: same(paged(bufd).view(torch.float8_e4m3fn)[blk.to(DEV), row.to(DEV)], paged(buf).view(torch.float8_e4m3fn)[blk, row]))
    check("fp8 strided view advanced indexing (no NaN bytes)", "known",
          lambda: same(paged(nbufd).view(torch.float8_e4m3fn)[blk.to(DEV), row.to(DEV)], paged(nbuf).view(torch.float8_e4m3fn)[blk, row]))
    idx = torch.tensor([1, 1, 4])
    check("index_add_ (duplicates)", "ok", lambda: same(torch.zeros(11, 512, dtype=torch.bfloat16, device=DEV).index_add_(0, idx.to(DEV), xd[:3]),
                                                         torch.zeros(11, 512, dtype=torch.bfloat16).index_add_(0, idx, x[:3])))

    def inference_version_counter():
        with torch.inference_mode():
            t = torch.topk(xd.float(), 4).indices
            return t.is_inference() or t._version >= 0

    check("inference-mode output has is_inference() or a version counter", "known", inference_version_counter)

    # numerics
    a, w = torch.randn(64, 256), torch.randn(256, 128)
    ref = a.double() @ w.double()
    rel = lambda r: ((r.cpu().double() - ref).norm() / ref.norm()).item()  # noqa: E731
    check("fp32 matmul rel err < 1e-5 (vs fp64)", "ok", lambda: rel(a.to(DEV) @ w.to(DEV)) < 1e-5)
    check("bf16 topk indices == host", "ok", lambda: torch.equal(
        torch.topk(xd[0], 64).indices.cpu().sort().values, torch.topk(x[0], 64).indices.sort().values))


    # single-rank process group (EP=1 bring-up)
    check("neuron PG size 1: all_reduce", "ok", lambda: (lambda t: (dist.all_reduce(t, group=g), same(t, x))[1])(xd.clone()))
    check("neuron PG size 1: all_gather_into_tensor", "ok",
          lambda: (lambda o: (dist.all_gather_into_tensor(o, xd, group=g), same(o, x))[1])(torch.empty_like(xd)))
    check("neuron PG size 1: all_to_all_single", "known",
          lambda: (lambda o: (dist.all_to_all_single(o, xd, group=g), same(o, x))[1])(torch.empty_like(xd)))
    dist.destroy_process_group()

    failed = 0
    for name, expect, status in RESULTS:
        if expect == "ok":
            tag = "PASS" if status == "ok" else "FAIL"
            failed += tag == "FAIL"
        else:
            tag = "KNOWN-ISSUE" if status != "ok" else "FIXED-UPSTREAM?"
        print(f"{tag:16s} {name:62s} {status}")
    print(f"\n{failed} unexpected failure(s)")
    return 1 if failed else 0


if __name__ == "__main__":
    sys.exit(main())