"""Compositionality regression suite (ported from the Phase 1 checks).

These are the headline properties that make the HRR morpheme representation
useful. Each HRR measurement is paired against a random-vector control of the
same dimension, asserting the HRR signal is meaningfully stronger.
"""
import numpy as np

import mlx.core as mx

from morph_hrr import MorphemeTokenizer, unbind, cosine_similarity

DIM = 2048
# A tokenizer is reused across tests; construction is cheap and deterministic.
TOK = MorphemeTokenizer(dim=DIM, seed=0)

# Random-vector control: a fresh tokenizer whose "word vectors" are i.i.d. random
# unit vectors (no composition at all). Expected to show ~0 structure.


def _control():
    rng = np.random.default_rng(99)
    return normalize(mx.array(rng.normal(size=(DIM,)).astype(np.float32)))


def normalize(v):
    n = mx.sqrt(mx.sum(v * v) + 1e-12)
    return v / n


def _rand_vec(tag: str) -> mx.array:
    """Deterministic random unit vector keyed by tag (for the control family)."""
    rng = np.random.default_rng(abs(hash(tag)) % (2**32))
    return normalize(mx.array(rng.normal(size=(DIM,)).astype(np.float32)))


# -- 1. Prefix recovery: unbind recovers the prefix filler from the holistic vec -

def test_prefix_recovery_beats_control():
    word = TOK.word_vector("unhappy")
    recovered = unbind(word, TOK.prefix_role)
    target = TOK.bytes_vector("un")
    hrr_cos = float(cosine_similarity(target, recovered))

    # Control: unbinding a random vector with the prefix role gives noise.
    ctrl = unbind(_rand_vec("control-prefix"), TOK.prefix_role)
    ctrl_cos = abs(float(cosine_similarity(target, ctrl)))

    assert hrr_cos > 0.50, f"prefix recovery cosine {hrr_cos} too low"
    assert hrr_cos > ctrl_cos + 0.40, (
        f"HRR prefix recovery {hrr_cos} not clearly above control {ctrl_cos}"
    )


# -- 2. Shared-root clustering: unhappy ~ happy (same root "happy") ------------

def test_shared_root_clusters():
    unhappy = TOK.word_vector("unhappy")
    happy = TOK.word_vector("happy")
    hrr_cos = float(cosine_similarity(unhappy, happy))

    # Control: two unrelated random words should be ~orthogonal.
    ctrl_cos = abs(float(cosine_similarity(_rand_vec("a"), _rand_vec("b"))))

    assert hrr_cos > 0.20, f"shared-root cosine {hrr_cos} too low"
    assert hrr_cos > ctrl_cos + 0.15


# -- 3. Shared-suffix clustering: running ~ walking (same suffix "ing") --------

def test_shared_suffix_clusters():
    running = TOK.word_vector("running")
    walking = TOK.word_vector("walking")
    hrr_cos = float(cosine_similarity(running, walking))

    ctrl_cos = abs(float(cosine_similarity(_rand_vec("c"), _rand_vec("d"))))

    assert hrr_cos > 0.15, f"shared-suffix cosine {hrr_cos} too low"
    assert hrr_cos > ctrl_cos + 0.10


# -- 4. OOV root family: a built-from-pieces vector still neighbors its root ---

def test_oov_root_family():
    # "unkind" is in-vocab via segmentation; its root filler is "kind".
    unkind = TOK.word_vector("unkind")
    kind = TOK.word_vector("kind")
    recovered_root = unbind(unkind, TOK.root_role)
    hrr_cos = float(cosine_similarity(TOK.bytes_vector("kind"), recovered_root))

    ctrl = unbind(_rand_vec("control-root"), TOK.root_role)
    ctrl_cos = abs(float(cosine_similarity(TOK.bytes_vector("kind"), ctrl)))

    assert hrr_cos > 0.50, f"OOV root recovery cosine {hrr_cos} too low"
    assert hrr_cos > ctrl_cos + 0.40


# -- 5. Roles are near-orthogonal (composition uses distinct slots) -----------

def test_roles_are_near_orthogonal():
    pairs = [
        (TOK.prefix_role, TOK.root_role),
        (TOK.prefix_role, TOK.suffix_role),
        (TOK.root_role, TOK.suffix_role),
    ]
    for a, b in pairs:
        cos = abs(float(cosine_similarity(a, b)))
        assert cos < 0.10, f"roles not orthogonal: cosine {cos}"