import math

import pytest
import torch

import kernels

mk = kernels.get_kernel("phanerozoic/metakernel", version=1,
                        trust_remote_code=True)

requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(),
                                   reason="CUDA required")

_dossier = None


def dossier():
    global _dossier
    if _dossier is None:
        _dossier = mk.probe_device(quick=True)
    return _dossier


@requires_cuda
@pytest.mark.kernels_ci
def test_clock_measurement_sane():
    ghz = mk.measure_clock()
    assert 0.3 < ghz < 4.5, ghz


@requires_cuda
@pytest.mark.kernels_ci
def test_dossier_bandwidth_hierarchy():
    """Measured L2 bandwidth must exceed the best HBM number; streaming
    patterns must land in plausible ranges; random gather must cost far
    more than streaming."""
    d = dossier()
    hbm = d["hbm_bw"]["value"]
    best_stream = max(hbm["triad_4B"], hbm["triad_8B"], hbm["triad_16B"],
                      hbm["read_only"], hbm["write_only"])
    assert 50 < best_stream < 20000, hbm
    assert hbm["triad_16B"] >= hbm["triad_4B"] * 0.8, hbm
    assert hbm["read_only"] > 50 and hbm["write_only"] > 50, hbm
    assert hbm["random_gather_16B"] < best_stream, hbm
    assert d["l2_bw"]["value"] > best_stream, (d["l2_bw"]["value"],
                                               best_stream)


@requires_cuda
@pytest.mark.kernels_ci
def test_dossier_latency_hierarchy():
    """Dependent-load latency must order smem < L2 < HBM with the L1
    point between cache levels."""
    lat = dossier()["latency"]["value"]
    assert lat["smem_cycles"] < lat["l2_cycles"] < lat["hbm_cycles"], lat
    assert lat["l1_cycles"] < lat["l2_cycles"], lat
    assert 10 < lat["l1_cycles"] < 300, lat
    assert 5 < lat["smem_cycles"] < 150, lat
    assert 50 < lat["hbm_cycles"] < 3000, lat


@requires_cuda
@pytest.mark.kernels_ci
def test_dossier_mma_rates():
    """Tensor cores beat the FMA pipe; int8 is not slower than fp16; fp8
    support matches the architecture."""
    d = dossier()
    mma = d["mma"]["value"]
    fma = d["fma_f32"]["value"]
    assert isinstance(mma["fp16"], float) and mma["fp16"] > fma, (mma, fma)
    assert isinstance(mma["bf16"], float) and mma["bf16"] > fma
    assert isinstance(mma["tf32"], float)
    assert isinstance(mma["int8"], float) and mma["int8"] > 0.9 * mma["fp16"]
    major = int(d["device"]["sm_arch"][2])
    minor = int(d["device"]["sm_arch"][3:])
    if major * 10 + minor >= 89:
        assert isinstance(mma["fp8_e4m3"], float), mma
    else:
        assert mma["fp8_e4m3"] == "unsupported", mma


@requires_cuda
@pytest.mark.kernels_ci
def test_dossier_atomics_contention_curve():
    """Throughput must not fall as the number of distinct addresses grows,
    and the endpoints must differ by a large factor."""
    at = dossier()["atomics_f32"]["value"]
    curve = at["curve_slots"]
    slots = sorted(int(k) for k in curve)
    vals = [curve[str(s)] for s in slots]
    for lo, hi in zip(vals, vals[1:]):
        assert hi >= lo * 0.7, curve  # monotone up to measurement noise
    assert vals[-1] > vals[0] * 4, curve
    assert at["shared_same_addr"] > 0, at


@requires_cuda
@pytest.mark.kernels_ci
def test_dossier_fp64_pipe():
    d = dossier()
    assert 0.02 < d["fma_f64"]["value"] < d["fma_f32"]["value"] * 0.6, \
        (d["fma_f64"]["value"], d["fma_f32"]["value"])


@requires_cuda
@pytest.mark.kernels_ci
def test_dossier_occupancy_knees_timed():
    """The timed counterpart to the API table: a dependent-chain workload
    must lose throughput when dynamic smem forces occupancy down."""
    knees = dossier()["occupancy_knees_fma_tflops"]["value"]
    healthy = knees["t256"]["0KB"]
    starved = knees["t128"]["48KB"]  # 2 blocks x 128 threads per SM
    assert starved < healthy * 0.8, knees
    assert knees["t1024"]["48KB"] >= starved, knees


@requires_cuda
@pytest.mark.kernels_ci
def test_dossier_launch_and_barrier():
    d = dossier()
    lo = d["launch_overhead"]["value"]
    assert 0.2 < lo["queued_us"] < 500, lo
    assert lo["round_trip_us"] > lo["queued_us"] * 0.5, lo
    gb = d["grid_barrier"]["value"]
    assert gb == "unsupported" or 0.05 < gb < 500, gb


@requires_cuda
@pytest.mark.kernels_ci
def test_dossier_occupancy_monotone():
    occ = dossier()["occupancy_blocks_per_sm"]["value"]
    assert occ["t128"]["0KB"] >= occ["t1024"]["0KB"], occ
    assert occ["t256"]["48KB"] <= occ["t256"]["0KB"], occ
    assert occ["t256"]["0KB"] >= 1, occ


@requires_cuda
@pytest.mark.kernels_ci
def test_bench_reports_and_roofline():
    d = dossier()
    n = 4096
    x = torch.randn(n, n, device="cuda", dtype=torch.float16)
    y = torch.randn(n, n, device="cuda", dtype=torch.float16)
    flops = 2 * n ** 3
    byts = 3 * n * n * 2
    rep = mk.bench(lambda: x @ y, iters=48, warmup=8, bytes=byts,
                   flops=flops, dossier=d, dtype="fp16")
    assert rep["median_ms"] > 0
    assert rep["samples"] == 48
    assert rep["rejected_throttled"] <= rep["samples"]
    assert 1 < rep["achieved_tflops"] < 5000, rep
    assert rep["roofline"]["bound"] == "compute", rep["roofline"]

    big = torch.randn(64 << 20, device="cuda")
    dst = torch.empty_like(big)
    rep2 = mk.bench(lambda: dst.copy_(big), iters=32, warmup=4,
                    bytes=2 * big.numel() * 4, flops=big.numel(),
                    dossier=d)
    assert rep2["roofline"]["bound"] == "memory", rep2["roofline"]
    assert rep2["achieved_gbs"] > 20, rep2


@requires_cuda
@pytest.mark.kernels_ci
def test_bench_graph_mode_strips_dispatch():
    """Graph replay removes Python and dispatch overhead: for a train of
    tiny kernels the captured timing must not exceed the eager timing."""
    x = torch.randn(1024, device="cuda")

    def tiny():
        for _ in range(16):
            x.mul_(1.0000001)

    eager = mk.bench(tiny, iters=32, warmup=8)
    graphed = mk.bench(tiny, iters=32, warmup=8, graph=True)
    assert graphed["graph_captured"] and not eager["graph_captured"]
    assert graphed["median_ms"] <= eager["median_ms"] * 1.10, \
        (graphed["median_ms"], eager["median_ms"])


@requires_cuda
@pytest.mark.kernels_ci
def test_compare_multi_output_and_tolerance_route():
    def gen(seed):
        g = torch.Generator(device="cuda").manual_seed(seed)
        return (torch.randn(64, 64, device="cuda", generator=g),)

    ref = lambda x: (x * 2.0, x.sum(dim=1))  # noqa: E731
    same = mk.compare(ref, ref, gen, trials=2)
    assert same["pass"] and same["worst_trial"]["max_ulp"] == 0

    shifted = lambda x: (x * 2.0 + 1e-6, x.sum(dim=1))  # noqa: E731
    ulp_fail = mk.compare(ref, shifted, gen, trials=2)
    assert not ulp_fail["pass"]
    tol_pass = mk.compare(ref, shifted, gen, trials=2, atol=1e-4)
    assert tol_pass["pass"], tol_pass["criterion"]
    assert tol_pass["max_abs"] <= 2e-6, tol_pass["max_abs"]


@requires_cuda
@pytest.mark.kernels_ci
def test_compare_builtin_matmul_k_shuffle():
    """The built-in K-permutation helper: a matmul judged against its own
    accumulation-order band."""
    def gen(seed):
        g = torch.Generator(device="cuda").manual_seed(seed)
        a = torch.randn(96, 4096, device="cuda", generator=g,
                        dtype=torch.float16)
        b = torch.randn(4096, 64, device="cuda", generator=g,
                        dtype=torch.float16)
        return (a, b)

    ref = lambda a, b: a @ b  # noqa: E731
    rep = mk.compare(ref, ref, gen, trials=2,
                     order_shuffle=mk.order_shuffles.matmul_k(),
                     order_trials=4)
    assert rep["pass"], rep["criterion"]

    bad = lambda a, b: (a @ b) * 1.01  # noqa: E731
    rep2 = mk.compare(ref, bad, gen, trials=2,
                      order_shuffle=mk.order_shuffles.matmul_k(),
                      order_trials=4)
    assert not rep2["pass"], rep2["criterion"]


@requires_cuda
@pytest.mark.kernels_ci
def test_ulp_diff_unit_cases():
    a = torch.tensor([1.0, -0.0, float("nan"), 2.0], device="cuda")
    b = torch.tensor([torch.nextafter(torch.tensor(1.0),
                                      torch.tensor(2.0)).item(),
                      0.0, float("nan"), 2.0], device="cuda")
    d = mk.ulp_diff(a, b)
    assert d[0].item() == 1
    assert d[1].item() == 0          # -0.0 and +0.0 are the same value
    assert d[2].item() == 0          # both-NaN agree
    assert d[3].item() == 0
    c = torch.tensor([float("nan")], device="cuda")
    e = torch.tensor([1.0], device="cuda")
    assert mk.ulp_diff(c, e)[0].item() == 1 << 62  # one-sided NaN
    h = torch.tensor([1.0], dtype=torch.float16, device="cuda")
    h2 = torch.nextafter(h, torch.tensor([2.0], dtype=torch.float16,
                                         device="cuda"))
    assert mk.ulp_diff(h, h2)[0].item() == 1


@requires_cuda
@pytest.mark.kernels_ci
def test_compare_pass_and_fail():
    def gen(seed):
        g = torch.Generator(device="cuda").manual_seed(seed)
        return (torch.randn(256, 512, device="cuda", generator=g),)

    ref = lambda x: torch.relu(x) * 2.0  # noqa: E731
    same = mk.compare(ref, ref, gen, trials=3)
    assert same["pass"] and same["worst_trial"]["max_ulp"] == 0

    close = mk.compare(ref, lambda x: (torch.relu(x) * 2.0), gen, trials=3)
    assert close["pass"]

    broken = mk.compare(ref, lambda x: torch.relu(x) * 2.0 + 1e-3, gen,
                        trials=3)
    assert not broken["pass"]
    assert broken["worst_trial"]["max_ulp"] > 100
    assert broken["worst_trial"]["first_divergent_index"] >= 0
    assert len(broken["ulp_hist_log2"]) > 0


@requires_cuda
@pytest.mark.kernels_ci
def test_compare_reduction_order_band():
    """A long fp32 sum judged against its own accumulation-order spread:
    the shuffled reference bounds the legitimate band, and a candidate
    that IS a reordered sum passes while a corrupted one fails."""
    def gen(seed):
        g = torch.Generator(device="cuda").manual_seed(seed)
        return (torch.randn(8, 200_000, device="cuda", generator=g)
                * 100.0,)

    def shuffle(args, seed):
        g = torch.Generator(device="cuda").manual_seed(seed)
        p = torch.randperm(args[0].shape[1], device="cuda", generator=g)
        return (args[0][:, p],)

    ref = lambda x: x.sum(dim=1)  # noqa: E731
    cand = lambda x: x.flip(1).sum(dim=1)  # noqa: E731
    rep = mk.compare(ref, cand, gen, trials=3, order_shuffle=shuffle,
                     order_trials=6)
    assert rep["pass"], rep["worst_trial"]

    bad = lambda x: x.sum(dim=1) * (1 + 1e-4)  # noqa: E731
    rep2 = mk.compare(ref, bad, gen, trials=3, order_shuffle=shuffle,
                      order_trials=6)
    assert not rep2["pass"], rep2["worst_trial"]


@requires_cuda
@pytest.mark.kernels_ci
def test_fuzz_shapes_envelope():
    def make(shape, dtype=torch.float32):
        n = math.prod(shape)
        return torch.randn(*shape, device="cuda", dtype=dtype) if n \
            else torch.zeros(*shape, device="cuda", dtype=dtype)

    ref = lambda x: torch.relu(x.contiguous())  # noqa: E731
    good = mk.fuzz_shapes(torch.relu, ref, make,
                          variants=("contig", "transposed", "strided_rows",
                                    "offset_slice"))
    assert good["failed"] == 0, good
    assert not good["context_lost"]

    # an op that silently assumes contiguity must be caught
    cheat = lambda x: torch.relu(  # noqa: E731
        torch.as_strided(x, x.shape, [x.shape[-1], 1] if x.dim() == 2
                         else [1]))
    bad = mk.fuzz_shapes(cheat, ref, make,
                         variants=("contig", "transposed", "strided_rows",
                                   "offset_slice"))
    assert bad["failed"] > 0, bad


@requires_cuda
@pytest.mark.kernels_ci
def test_fuzz_multi_arg_dtypes_overflow():
    """Binary op across dtypes with the overflow_expand variant: past
    2**31 logical elements the case either runs correctly (large card) or
    reports skipped_vram (small card); silent absence is not allowed."""
    def make(shape, dtype=torch.float16):
        n = math.prod(shape)
        f = torch.randn if n else lambda *s, **k: torch.zeros(*s, **k)
        return (f(*shape, device="cuda", dtype=dtype),
                f(*shape, device="cuda", dtype=dtype))

    op = torch.add
    ref = lambda a, b: a.contiguous() + b.contiguous()  # noqa: E731
    rep = mk.fuzz_shapes(op, ref, make, shapes=[(1, 17), (128, 128)],
                         dtypes=(torch.float16, torch.float32))
    assert rep["failed"] == 0, [c for c in rep["cases"]
                                if c["status"] not in ("ok", "skipped_vram",
                                                       "skipped_empty")]
    ov = [c for c in rep["cases"] if c["variant"] == "overflow_expand"]
    assert len(ov) == 4
    assert all(c["status"] in ("ok", "skipped_vram") for c in ov), ov
    ran = [c for c in ov if c["status"] == "ok"]
    for c in ran:
        assert c["logical_elems"] > (1 << 31), c


@requires_cuda
@pytest.mark.kernels_ci
def test_sweep_orders_configurations():
    x = torch.randn(1 << 22, device="cuda")

    def factory(chunk):
        def run():
            for piece in x.split(chunk):
                piece.mul_(1.0000001)
        return run

    rows = mk.sweep(factory, {"chunk": [1 << 14, 1 << 20, 1 << 22]},
                    iters=12, warmup=3)
    assert len(rows) == 3
    assert rows[0]["median_ms"] <= rows[-1]["median_ms"]
    assert all("chunk" in r for r in rows)


@requires_cuda
@pytest.mark.kernels_ci
def test_sweep_check_gate_and_budget():
    x = torch.randn(1 << 18, device="cuda")

    def factory(scale):
        return lambda: x.mul(scale)

    rows = mk.sweep(factory, {"scale": [1.0, 2.0]}, iters=6, warmup=2,
                    check=lambda fn: bool(torch.isfinite(fn()).all()))
    assert all(r["check"] is True for r in rows)

    rows2 = mk.sweep(factory, {"scale": [1.0, 2.0, 3.0, 4.0]}, iters=6,
                     warmup=2, budget_s=0.0)
    assert any(r.get("status") == "not_run" for r in rows2)
    assert len(rows2) == 4  # skipped points are visible, not dropped


@requires_cuda
@pytest.mark.kernels_ci
def test_stamps_per_block():
    st = torch.zeros(3, 3, dtype=torch.int64)
    st[0] = torch.tensor([100, 300, 600])
    st[1] = torch.tensor([100, 500, 800])
    st[2] = torch.tensor([0, 0, 0])  # a block that never stamped
    rep = mk.read_stamps(st.cuda(), labels=["load", "compute"],
                         clock_ghz=1.0)
    assert rep["blocks_reporting"] == 2 and rep["blocks_total"] == 3
    load = rep["phases"][0]
    assert load["ticks_min"] == 200 and load["ticks_max"] == 400
    assert load["ticks_median"] in (200, 300, 400)


@requires_cuda
@pytest.mark.kernels_ci
def test_stamps_roundtrip():
    st = mk.alloc_stamps(4)
    base = 1_000_000
    ticks = [base, base + 2_000_000, base + 3_000_000, base + 3_500_000]
    st.copy_(torch.tensor(ticks, dtype=torch.int64))
    rep = mk.read_stamps(st, labels=["load", "compute", "store"],
                         clock_ghz=2.0)
    assert [p["phase"] for p in rep["phases"]] == ["load", "compute",
                                                   "store"]
    assert rep["phases"][0]["ticks"] == 2_000_000
    assert abs(rep["phases"][0]["us"] - 1000.0) < 1e-6
    assert rep["phases"][2]["ticks"] == 500_000


@requires_cuda
@pytest.mark.kernels_ci
def test_chase_repeatability():
    """The instrument itself must be stable: back-to-back shared-memory
    chases agree within 20%."""
    ring = mk.ops  # ensure ops import path is alive
    import metakernel as _m
    r = _m._sattolo_ring(4096, "cuda")
    out = torch.zeros(2, dtype=torch.int64, device="cuda")
    mk.ops.mk_chase_shared(r, 1 << 14, out)
    torch.cuda.synchronize()
    t1 = int(out[0].item())
    mk.ops.mk_chase_shared(r, 1 << 14, out)
    torch.cuda.synchronize()
    t2 = int(out[0].item())
    assert abs(t1 - t2) / max(t1, t2) < 0.2, (t1, t2)
    # the ring is a single cycle: 2^14 hops through 4096 slots ends where
    # modular arithmetic says it must, proving every hop was taken
    assert 0 <= int(out[1].item()) < 4096