from __future__ import annotations
import networkx as nx
import numpy as np
import pytest

from hieromem_core.memory import Memory


def expect(condition: bool, message: str | None = None) -> None:
    """Fail the test when *condition* is false without using ``assert``."""

    if not condition:
        pytest.fail(message or "Expectation failed", pytrace=True)


def _insert_cluster(
        memory: Memory,
        center: np.ndarray,
        count: int,
        label: str) -> None:
    rng = np.random.default_rng(42)
    for _ in range(count):
        noise = rng.normal(scale=0.02, size=center.shape)
        memory.insert(center + noise, nx.DiGraph(), meta={"cluster": label})


def test_resonant_retrieve_converges_to_dominant_cluster() -> None:
    memory = Memory(Nmax=20, tau=0.3)

    cluster_a = np.array([1.0, 0.0], dtype=float)
    cluster_b = np.array([0.0, 1.0], dtype=float)

    _insert_cluster(memory, cluster_a, count=5, label="A")
    _insert_cluster(memory, cluster_b, count=2, label="B")

    query = np.array([0.6, 0.4], dtype=float)

    results = memory.resonant_retrieve(query, k=3)

    expect(all(ep.meta.get("cluster") == "A" for ep in results),
           "results should converge to dominant cluster")

    context_vec = memory.context
    if context_vec is None:
        pytest.fail("context vector should be available", pytrace=True)

    context = context_vec / (np.linalg.norm(context_vec) + 1e-9)
    target = cluster_a / np.linalg.norm(cluster_a)
    expect(bool(np.linalg.norm(context - target) < 0.16),
           "context should align with dominant cluster")


def test_resonant_retrieve_rejects_zero_query() -> None:
    memory = Memory(Nmax=5)
    memory.insert(np.array([1.0, 0.0], dtype=float), nx.DiGraph())

    with pytest.raises(ValueError):
        memory.resonant_retrieve(np.zeros(2, dtype=float))
