"""Decode a real model's axes and check they read as their own concepts.

Needs the Qwen3.5-4B weights and its published lens, so it is skipped unless
`ULENS_TEST_NETWORK=1`:

    ULENS_TEST_NETWORK=1 pytest tests/test_decode_real.py

The assertions are deliberately about *shape of outcome*, not exact word lists: a decode is
a ranking over a quarter of a million tokens and a tokenizer revision can reorder its tail.
What must hold is that each axis returns its own concept far more than a random direction
does, and that the random direction returns nothing.
"""
import os

import numpy as np
import pytest

import ulens

pytestmark = pytest.mark.skipif(os.environ.get("ULENS_TEST_NETWORK") != "1",
                                reason="set ULENS_TEST_NETWORK=1 to decode a real model")

MODEL = "qwen3.5-4b"
K = 10


@pytest.fixture(scope="module")
def decoded():
    block = ulens.models.decode_block(MODEL)
    cone = ulens.basis.get(MODEL, block=block)
    voc = ulens.decode.Vocabulary(MODEL, block)
    return cone, voc, voc.decode(cone, k=K)


def stem_hits(words, anchors, n=5) -> int:
    """How many decoded words share a 5-character stem with any anchor of their own axis."""
    low = {a.lower() for a in anchors}
    return sum(any(w[:n] == a[:n] or a.startswith(w) or w.startswith(a) for a in low)
               for w in (x.lower() for x in words))


def test_decode_block_is_not_the_readout_block(decoded):
    cone, voc, _ = decoded
    assert cone.block == voc.block == ulens.models.decode_block(MODEL) == 23
    assert ulens.models.get(MODEL).readout_block == 20


def test_every_axis_reads_as_its_own_concept(decoded):
    cone, _, dec = decoded
    for d in dec:
        n = stem_hits(d.words("+"), cone.axis(d.name).anchors)
        assert n >= 3, f"{d.name} decoded to {d.words('+')}"


def test_a_random_direction_decodes_to_nothing(decoded):
    cone, voc, _dec = decoded
    rng = np.random.default_rng(0)
    anchors = [w for a in cone.axes for w in a.anchors]
    for _ in range(3):
        v = rng.normal(size=cone.hidden_size).astype(np.float32)
        assert stem_hits([h.word for h in voc.top(v, k=K)], anchors) <= 1


def test_each_axis_beats_the_other_axes_anchors(decoded):
    """An axis must match its own anchors more than another axis's."""
    cone, _, dec = decoded
    for d in dec:
        own = stem_hits(d.words("+"), cone.axis(d.name).anchors)
        for other in cone.names:
            if other == d.name:
                continue
            assert own >= stem_hits(d.words("+"), cone.axis(other).anchors), \
                f"{d.name} matches {other}'s anchors at least as well"


def test_the_negative_end_is_not_the_positive_end(decoded):
    _, _, dec = decoded
    for d in dec:
        assert not set(d.words("+")) & set(d.words("-"))
        assert all(h.score > 0 for h in d.positive)
        assert all(h.score < 0 for h in d.negative)


def test_norms_cache_is_tiny_next_to_the_matrix_it_replaces(decoded):
    _, voc, _ = decoded
    assert voc.cache_path.exists()
    on_disk = voc.cache_path.stat().st_size
    materialised = voc.n_tokens * voc.spec.hidden_size * 2        # fp16
    assert on_disk * 100 < materialised
    assert len(voc.norms) == voc.n_tokens
    assert np.all(voc.norms > 0)


def test_scores_are_bounded_like_cosines(decoded):
    _, voc, _ = decoded
    rng = np.random.default_rng(1)
    s = voc.scores(rng.normal(size=voc.spec.hidden_size).astype(np.float32))
    assert s.shape == (voc.n_tokens,)
    assert s.min() >= -1.001 and s.max() <= 1.001