"""The compiled kernels (``rust/simurg-core``, plan §18 P9) must equal the pure-Python paths bit for bit."""

from __future__ import annotations

import datetime as dt
import math
from dataclasses import fields, replace

import numpy as np
import pytest
from tests.conftest import GSA_BPM, GSA_EGT, GSA_IDATS

import simurg.formats.bpm as bpm_mod
import simurg.genotype.baf_lrr as baf_mod
import simurg.genotype.gencall as gencall_mod
import simurg.signal.normalization.fit as fit_mod
from simurg._core import core
from simurg.core.errors import FormatError
from simurg.formats.bpm import read_bpm
from simurg.synth.files import SynthLocus, encode_bpm

pytestmark = pytest.mark.skipif(core is None, reason="simurg_core extension not installed")

MODULES = (bpm_mod, gencall_mod, fit_mod, baf_mod)


@pytest.fixture
def pure(monkeypatch):  # type: ignore[no-untyped-def]
    """Context switch: ``with pure(): ...`` runs the pure-Python paths."""

    class _Pure:
        def __enter__(self) -> None:
            for m in MODULES:
                monkeypatch.setattr(m, "core", None)

        def __exit__(self, *exc: object) -> None:
            for m in MODULES:
                monkeypatch.setattr(m, "core", core)

    return _Pure


def _bits(a: np.ndarray) -> bytes:
    return np.ascontiguousarray(a).tobytes()


def _edge_values() -> np.ndarray:
    neg_nan = np.frombuffer(np.uint64(0xFFF8000000000000).tobytes(), np.float64)[0]
    payload_nan = np.frombuffer(np.uint64(0x7FF8000000001234).tobytes(), np.float64)[0]
    tiny = np.finfo(np.float64).tiny
    special = [0.0, -0.0, 1.0, -1.0, math.inf, -math.inf, math.nan, neg_nan, payload_nan, tiny, tiny / 4, -tiny,
               1e-320, 709.78, 709.79, -745.2, -746.0, 1e300, -1e300, 2.0**-1074, 0.5, 10.0, 1e-5]  # fmt: skip
    rng = np.random.default_rng(7)
    rand = np.concatenate(
        [rng.normal(0, 3, 20_000), rng.uniform(-800, 800, 5_000), np.exp(rng.uniform(-700, 700, 5_000))]
    )
    return np.concatenate([np.array(special, np.float64), rand, -rand[:1000]])


@pytest.mark.parametrize("fn", [math.exp, math.log10, math.log2, gencall_mod._pow(1.71), gencall_mod._pow(0.35)])
def test_libm_kernels_bit_identical(fn, pure) -> None:  # type: ignore[no-untyped-def]
    x = _edge_values()
    fast = gencall_mod._libm(fn, x)
    with pure():
        ref = gencall_mod._libm(fn, x)
    assert fast.dtype == ref.dtype == np.float64
    assert _bits(fast) == _bits(ref)


def test_pow_kernel_rejects_integer_exponent() -> None:
    assert core is not None
    with pytest.raises(ValueError):
        core.pow(np.ones(3), 2.0)


def test_closest_binned_identical(pure) -> None:  # type: ignore[no-untyped-def]
    rng = np.random.default_rng(3)
    for trial in range(20):
        n = int(rng.integers(0, 3000))
        raw_x = rng.gamma(2.0, 400.0, n).astype(np.float32)
        raw_y = rng.gamma(1.0, 60.0 if trial % 2 else 3000.0, n).astype(np.float32)
        if trial == 5 and n:
            raw_x[: n // 3] = raw_x[0]  # ties: first minimum must win
            raw_y[: n // 3] = raw_y[0]
        if trial == 7 and n:
            raw_x[n // 2] = np.nan
        m = int(rng.integers(1, 400))
        axis = np.linspace(0, float(np.float32(rng.uniform(100, 5000))), m).astype(np.float32)
        zeros = np.zeros(m, np.float32)
        ax, ay = (axis, zeros) if trial % 3 else (zeros, axis)
        fast = fit_mod._closest_binned(raw_x, raw_y, ax, ay)
        with pure():
            ref = fit_mod._closest_binned(raw_x, raw_y, ax, ay)
        assert fast.dtype == ref.dtype
        assert np.array_equal(fast, ref), trial


def _loci() -> list[SynthLocus]:
    return [
        SynthLocus("rs1", "1", 100, "[A/G]", 11, 0, 0, norm_id=1, source_seq="ACGT[A/G]TTGA", probe_a="ACGTA"),
        SynthLocus("rs2", "2", 200, "[A/T]", 12, 13, 1, ilmn_strand="BOT", ref_strand="-", probe_a="AC", probe_b="AT"),
        SynthLocus("rs3", "X", 300, "[C/G]", 14, 15, 2, norm_id=99),
        SynthLocus("rsü", "Y", 400, "[T/C]", 16, 0, 0, norm_id=2, intensity_only=1),
        SynthLocus("rs5", "MT", 500, "[A/C]", 17, 0, 0, norm_id=99, source_seq="A" * 200 + "[A/C]"),
    ]


def _manifest_dict(m) -> dict:  # type: ignore[no-untyped-def]
    out = {}
    for f in fields(m):
        v = getattr(m, f.name)
        out[f.name] = (v.dtype.str, v.shape, v.tobytes()) if isinstance(v, np.ndarray) else v
    return out


@pytest.mark.parametrize("entry_version", [4, 6, 7, 8])
@pytest.mark.parametrize("order", [None, [4, 3, 2, 1, 0]])
def test_bpm_parser_identical(tmp_path, pure, entry_version, order) -> None:  # type: ignore[no-untyped-def]
    p = tmp_path / "m.bpm"
    p.write_bytes(encode_bpm(_loci(), entry_version=entry_version, entry_order=order))
    fast = _manifest_dict(read_bpm(p))
    with pure():
        ref = _manifest_dict(read_bpm(p))
    assert fast == ref


def test_bpm_parser_errors_identical(tmp_path, pure) -> None:  # type: ignore[no-untyped-def]
    good = encode_bpm(_loci())
    cases = {
        "truncated": good[: len(good) - 40],
        "bad_assay": encode_bpm([SynthLocus("rs1", address_b=0, assay_type=1)]),
        "bad_map": encode_bpm([SynthLocus("rs1", pos=12345)]).replace(b"\x0512345", b"\x0512x45"),
    }
    for label, data in cases.items():
        p = tmp_path / f"{label}.bpm"
        p.write_bytes(data)
        with pytest.raises(FormatError) as fast:
            read_bpm(p)
        with pure(), pytest.raises(FormatError) as ref:
            read_bpm(p)
        assert (fast.value.code, str(fast.value)) == (ref.value.code, str(ref.value)), label


@pytest.mark.data
@pytest.mark.skipif(not GSA_BPM.exists(), reason="GSA demo data not fetched")
def test_genotype_sample_identical(pure) -> None:  # type: ignore[no-untyped-def]
    from simurg.formats.egt import read_egt
    from simurg.formats.gtc import encode_gtc
    from simurg.formats.idat import read_idat
    from simurg.genotype.clusters import align_clusters
    from simurg.genotype.sample import genotype_sample, manifest_context, params_for_preset
    from simurg.signal.raw_xy import LocusMapCache

    def run() -> bytes:
        m = read_bpm(GSA_BPM)
        ctx = manifest_context(m, align_clusters(read_egt(GSA_EGT), m))
        ctx = replace(ctx, manifest=m.lite(), cache=LocusMapCache(m.lite()))
        stem = next(iter(sorted(GSA_IDATS.rglob("*_Grn.idat")))).as_posix()[: -len("_Grn.idat")]
        g = genotype_sample(read_idat(stem + "_Grn.idat"), read_idat(stem + "_Red.idat"), ctx,
                            params_for_preset("current"), dt.datetime(2026, 1, 1))  # fmt: skip
        return encode_gtc(g)

    fast = run()
    with pure():
        ref = run()
    assert fast == ref


def test_closest_exact_identical(pure) -> None:  # type: ignore[no-untyped-def]
    rng = np.random.default_rng(11)
    for trial in range(10):
        n = int(rng.integers(1, 1500))
        xr = rng.normal(500, 200, n).astype(np.float32)
        yr = rng.normal(300, 100, n).astype(np.float32)
        if trial == 2:
            xr[n // 2 :] = xr[0]
            yr[n // 2 :] = yr[0]
        if trial == 3:
            xr[n // 3] = np.nan
        m = int(rng.integers(1, 300))
        x = rng.normal(500, 200, m).astype(np.float32)
        y = rng.normal(300, 100, m).astype(np.float32)
        fast = fit_mod._closest_slow(xr, yr, x, y)
        with pure():
            ref = fit_mod._closest_slow(xr, yr, x, y)
        assert np.array_equal(fast, ref), trial


def _synthetic_clusters(n: int, rng: np.random.Generator):  # type: ignore[no-untyped-def]
    from simurg.genotype.clusters import ClusterArrays

    tm = np.sort(rng.uniform(0, 1, (n, 3)), axis=1).astype(np.float32)
    td = rng.uniform(0.0, 0.08, (n, 3)).astype(np.float32)
    rm = rng.uniform(0.2, 3, (n, 3)).astype(np.float32)
    rd = rng.uniform(0.0, 0.3, (n, 3)).astype(np.float32)
    counts = rng.integers(0, 200, (n, 3)).astype(np.int32)
    counts[rng.random((n, 3)) < 0.2] = 0
    counts[: n // 50] = 0  # no samples at all
    td[rng.random((n, 3)) < 0.02] = 0.0
    tm[rng.random((n, 3)) < 0.01] = np.nan
    rm[rng.random((n, 3)) < 0.01] = np.nan
    thr = rng.uniform(0, 0.3, n).astype(np.float32)
    thr[rng.random(n) < 0.3] = np.nan
    gt = rng.uniform(0, 1, n).astype(np.float32)
    gt[rng.random(n) < 0.02] = np.nan
    return ClusterArrays(
        present=rng.random(n) > 0.01, n=counts, theta_mean=tm, theta_dev=td, r_mean=rm, r_dev=rd,
        intensity_threshold=thr, gentrain=gt, mean_r=rng.uniform(0.2, 3, n).astype(np.float32),
    )  # fmt: skip


def test_call_genotypes_and_baf_lrr_identical(pure) -> None:  # type: ignore[no-untyped-def]
    from simurg.genotype.baf_lrr import baf_lrr
    from simurg.genotype.gencall import call_genotypes

    rng = np.random.default_rng(5)
    n = 60_000
    c = _synthetic_clusters(n, rng)
    theta = rng.uniform(-0.05, 1.05, n).astype(np.float32)
    theta[: n // 3] = c.theta_mean[: n // 3, rng.integers(0, 3)]  # exactly on a cluster centre
    r = rng.gamma(2.0, 0.6, n).astype(np.float32)
    raw_zero = rng.random(n) < 0.01
    theta[raw_zero] = np.nan
    r[raw_zero] = np.nan
    r[rng.random(n) < 0.01] = 0.0
    intensity_only = rng.random(n) < 0.02
    for cutoff in (0.15, 0.0):
        fast = call_genotypes(theta, r, raw_zero, c, intensity_only, cutoff)
        with pure():
            ref = call_genotypes(theta, r, raw_zero, c, intensity_only, cutoff)
        assert _bits(fast.genotypes) == _bits(ref.genotypes)
        assert _bits(fast.scores) == _bits(ref.scores)
        assert (fast.num_calls, fast.num_intensity_only) == (ref.num_calls, ref.num_intensity_only)
    fb, fl = baf_lrr(theta, r, raw_zero, c, intensity_only)
    with pure():
        rb, rl = baf_lrr(theta, r, raw_zero, c, intensity_only)
    assert _bits(fb) == _bits(rb)
    assert _bits(fl) == _bits(rl)