"""ADR-185 P1 — AETHER binding tests, incl. the §4.1 bit-for-bit parity gate.

The parity test compares the binding's embedding to a committed golden VECTOR
produced by the native-Rust reference (`tests/aether_parity.rs`), within a
numerical tolerance. It is NOT a byte-hash: the embedding is f32 with
transcendental ops, so exact bytes are not reproducible across the CPU
architectures this project ships wheels for. A mismatch beyond tolerance is a
release blocker, not a warning.
"""

from __future__ import annotations

import json
import math
import struct
import tempfile
from pathlib import Path

import pytest

from wifi_densepose import aether

GOLDEN = Path(__file__).parent / "golden"

# Cross-architecture f32 parity tolerance. The AETHER embedding is pure f32 with
# transcendental ops that differ in the last bits across CPUs/libm, so exact
# byte equality is not portable across the wheels this project builds. 1e-4 is
# ~100x the observed cross-arch drift on unit-normed values and ~100x smaller
# than any real algorithm change. Combined atol+rtol so both small and larger
# components are bounded. (ADR-185 §4.1.)
PARITY_ATOL = 1e-4
PARITY_RTOL = 1e-4


def load_input() -> list[list[float]]:
    return json.loads((GOLDEN / "aether_input.json").read_text())


def assert_embedding_matches_golden(embedding: list[float], golden_name: str) -> None:
    """Assert `embedding` matches the committed golden vector within tolerance."""
    golden = json.loads((GOLDEN / golden_name).read_text())
    assert len(embedding) == len(golden), (
        f"{golden_name}: length {len(embedding)} != golden {len(golden)}"
    )
    worst = max(
        (abs(a - b), i, a, b)
        for i, (a, b) in enumerate(zip(embedding, golden))
        if abs(a - b) > PARITY_ATOL + PARITY_RTOL * abs(b)
    ) if any(
        abs(a - b) > PARITY_ATOL + PARITY_RTOL * abs(b)
        for a, b in zip(embedding, golden)
    ) else None
    assert worst is None, (
        f"{golden_name}: element {worst[1]} diverged from native golden beyond "
        f"tolerance (got {worst[2]}, golden {worst[3]}, |Δ|={worst[0]:.3e}) — "
        "a real regression, not cross-arch f32 drift."
    )


def build_extractor() -> aether.EmbeddingExtractor:
    # Must match the native-Rust reference construction exactly.
    cfg = aether.AetherConfig(d_model=64, d_proj=128, temperature=0.07, normalize=True)
    return aether.EmbeddingExtractor(n_subcarriers=56, config=cfg)


def _formula_weights(n: int) -> list[float]:
    # Byte-identical to aether_weights_parity.rs (k/65536 is exact in f32+f64).
    return [((i * 1103515245 + 12345) % 65536) / 65536.0 - 0.5 for i in range(n)]


def _write_weight_file(path: Path, weights: list[float]) -> None:
    # AETHER weight format: b"AETHERW1" + u32 count + LE f32 payload.
    with open(path, "wb") as f:
        f.write(b"AETHERW1")
        f.write(struct.pack("<I", len(weights)))
        f.write(b"".join(struct.pack("<f", w) for w in weights))


def test_config_roundtrips_fields() -> None:
    cfg = aether.AetherConfig(d_model=64, d_proj=128, temperature=0.07, normalize=True)
    assert cfg.d_model == 64
    assert cfg.d_proj == 128
    assert abs(cfg.temperature - 0.07) < 1e-6
    assert cfg.normalize is True


def test_embedding_shape_and_unit_norm() -> None:
    emb = build_extractor().embed(load_input())
    assert len(emb) == 128
    norm = math.sqrt(sum(x * x for x in emb))
    assert abs(norm - 1.0) < 1e-4, f"expected unit-norm embedding, got {norm}"


def test_binding_matches_native_golden_within_tolerance() -> None:
    """The release-blocking §4.1 gate: binding output == native Rust reference.

    Compares to a committed golden VECTOR within a numerical tolerance, not a
    SHA-256 of the raw f32 bytes. The embedding is pure f32 and uses
    transcendental ops (ln/sqrt/cos in the Gaussian init), which are NOT
    bit-reproducible across CPU architectures or libm implementations. A
    byte-hash therefore only ever matched the one arch that generated it, and
    failed on every other wheel this project builds (aarch64, macOS-arm). The
    tolerance below (1e-4) is orders of magnitude larger than cross-arch f32
    drift yet far tighter than any real algorithm change, which moves
    unit-normed elements by ~1e-2 or more. See ADR-185 §4.1.
    """
    emb = build_extractor().embed(load_input())
    assert_embedding_matches_golden(emb, "aether_embedding.json")


def test_embedding_is_deterministic() -> None:
    ext = build_extractor()
    inp = load_input()
    assert ext.embed(inp) == ext.embed(inp)


def test_cosine_similarity_self_is_one() -> None:
    v = [0.1 * i - 0.5 for i in range(32)]
    assert abs(aether.cosine_similarity(v, v) - 1.0) < 1e-5


def test_cosine_similarity_orthogonal_is_zero() -> None:
    a = [1.0, 0.0, 0.0, 0.0]
    b = [0.0, 1.0, 0.0, 0.0]
    assert abs(aether.cosine_similarity(a, b)) < 1e-6


def test_info_nce_loss_identical_batch_is_log_n() -> None:
    # Identical embeddings → all similarities equal → loss == ln(N).
    emb = [[1.0, 0.0, 0.0]] * 4
    loss = aether.info_nce_loss(emb, emb, 0.07)
    assert abs(loss - math.log(4)) < 0.1


def test_augment_pair_preserves_shape_and_differs() -> None:
    window = load_input()
    view_a, view_b = aether.CsiAugmenter().augment_pair(window, seed=42)
    assert len(view_a) == len(window)
    assert len(view_b) == len(window)
    assert len(view_a[0]) == len(window[0])
    differs = any(
        abs(x - y) > 1e-6
        for ra, rb in zip(view_a, view_b)
        for x, y in zip(ra, rb)
    )
    assert differs, "augment_pair should return two distinct views"


def test_missing_feature_message_names_the_real_fix() -> None:
    # The guard fires only on a from-source build without the feature. Its
    # message must name the real fix — rebuild with the feature — and must NOT
    # tell users to `pip install [aether]`, which is an empty extra that cannot
    # add compiled code to a built wheel.
    src = (Path(aether.__file__)).read_text()
    assert "--features aether" in src
    assert "pip install wifi-densepose[aether]" not in src


# ─── Weight loading (ADR-185 §13.a) ──────────────────────────────────

def test_load_weights_is_used_and_matches_native_golden() -> None:
    ext = build_extractor()
    baseline = ext.embed(load_input())  # random Xavier init

    weights = _formula_weights(ext.param_count)
    with tempfile.TemporaryDirectory() as d:
        wpath = Path(d) / "weights.bin"
        _write_weight_file(wpath, weights)
        ext.load_weights(str(wpath))

    loaded = ext.embed(load_input())

    # (1) The loaded weights actually take effect (not a silent no-op).
    assert any(abs(a - b) > 1e-6 for a, b in zip(baseline, loaded)), (
        "load_weights had no effect — embedding still equals the random-init baseline"
    )
    # (2) Matches the native-Rust reference that loaded the same weights, within
    #     tolerance. See test_binding_matches_native_golden_within_tolerance for
    #     why this is a tolerance compare and not a byte-hash.
    assert_embedding_matches_golden(loaded, "aether_loaded_embedding.json")


def test_save_then_load_weights_round_trips() -> None:
    ext = build_extractor()
    inp = load_input()
    with tempfile.TemporaryDirectory() as d:
        wpath = Path(d) / "roundtrip.bin"
        ext.save_weights(str(wpath))          # serialize current (random) weights
        emb_before = ext.embed(inp)
        ext2 = build_extractor()
        ext2.load_weights(str(wpath))         # load them into a fresh extractor
    assert ext2.embed(inp) == emb_before


def test_load_weights_rejects_bad_magic() -> None:
    ext = build_extractor()
    with tempfile.TemporaryDirectory() as d:
        wpath = Path(d) / "bad.bin"
        wpath.write_bytes(b"NOTAETHER" + b"\x00" * 8)
        with pytest.raises(ValueError):
            ext.load_weights(str(wpath))


def test_load_weights_rejects_wrong_param_count() -> None:
    ext = build_extractor()
    with tempfile.TemporaryDirectory() as d:
        wpath = Path(d) / "short.bin"
        _write_weight_file(wpath, [0.1, 0.2, 0.3])  # far too few params
        with pytest.raises(ValueError):
            ext.load_weights(str(wpath))