import numpy as np
import networkx as nx
import pytest

from hieromem_core.memory import Memory, BatchSymbolicSearchInput
from hieromem_core.graph_merge import hierarchical_merge
from hieromem_core.stability import GlobalStabilityMonitor
from hieromem_core.retrieval import DualRetrievalInput


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 test_multimodal_encoder_is_deterministic() -> None:
    from hieromem_core.multimodal import MultiModalEncoder
    encoder = MultiModalEncoder()
    payload = {"text": ["hello world"], "image": ["img_A"], "audio": ["wav_A"]}
    first = encoder.encode(payload)
    second = encoder.encode(payload)
    expect(bool(np.isclose(sum(first.modality_weights.values()), 1.0)),
           "modality weights should sum to 1")
    expect(bool(np.allclose(first.vector, second.vector)),
           "encoder outputs should be deterministic")


def test_hierarchical_merge_preserves_dag() -> None:
    g1 = nx.DiGraph()
    g1.add_edge("a", "b", relation="knows")
    g2 = nx.DiGraph()
    g2.add_edge("a", "c", relation="mentors")
    merged = hierarchical_merge([g1, g2], weights=[0.7, 0.3])
    expect(bool(nx.is_directed_acyclic_graph(merged)), "merged graph should be a DAG")
    expect(bool(merged.has_edge("a", "b")), "edge from graph g1 should be retained")
    expect(bool(merged.has_edge("a", "c")), "edge from graph g2 should be retained")


def test_hierarchical_merge_only_removes_cyclic_edges() -> None:
    g1 = nx.DiGraph()
    g1.add_edges_from(
        [
            ("c1", "c2"),
            ("c2", "c3"),
            ("c3", "c1"),
            ("root", "leaf"),
        ]
    )
    g2 = nx.DiGraph()
    g2.add_edges_from([("x", "y"), ("y", "x")])

    merged = hierarchical_merge([g1, g2])

    expect(bool(nx.is_directed_acyclic_graph(merged)), "merged graph should remain acyclic")
    # Acyclic edges should be preserved.
    expect(merged.has_edge("root", "leaf"), "acyclic edge should remain")

    remaining_cycle_edges = {
        ("c1", "c2"),
        ("c2", "c3"),
        ("c3", "c1"),
        ("x", "y"),
        ("y", "x"),
    } & set(merged.edges())
    # Exactly one edge should be removed from each original cycle.
    expect(len(remaining_cycle_edges) == 3,
           "one edge should be removed from each original cycle")
    expect(bool(("c1", "c2") not in merged.edges()),
           "cycle edge should be removed")
    expect(bool(("x", "y") not in merged.edges()),
           "cycle edge should be removed")


def test_memory_dual_retrieve_pipeline() -> None:
    memory = Memory(Nmax=4)

    g = nx.DiGraph()
    g.add_edge("user", "goal")
    memory.multi_modal_insert(
        {"text": ["alpha narrative"], "image": ["img1"]}, graph=g)
    memory.multi_modal_insert(
        {"text": ["beta narrative"], "audio": ["clip"]}, graph=g)

    query = memory.multi_modal_encoder.encode({"text": ["alpha"]}).vector
    result = memory.dual_retrieve(query, query_graph=g, k=2)

    expect(len(result.episodes) <= 2, "dual retrieval should respect k limit")
    expect(result.weights.size == len(result.episodes),
           "weights should align with episodes")
    expect(bool(set(result.dense_scores).issuperset(
        {ep.id for ep in result.episodes})),
           "dense scores should include episode ids")
    expect(bool(set(result.symbolic_scores).issuperset(
        {ep.id for ep in result.episodes})),
           "symbolic scores should include episode ids")
    stats = memory.stats()
    expect("global_stability" in stats, "stability stats should be present")
    expect(0.0 <= stats["summary_fidelity"] <= 1.0,
           "summary_fidelity should be bounded")
    expect(0.0 <= stats["summary_distortion"] <= 2.0,
           "summary_distortion should be bounded")


def test_memory_batch_dual_retrieve_supports_multiple_queries() -> None:
    memory = Memory(Nmax=5)

    g = nx.DiGraph()
    g.add_edge("user", "goal")
    memory.multi_modal_insert({"text": ["alpha narrative"]}, graph=g)
    memory.multi_modal_insert({"text": ["beta narrative"]}, graph=g)

    first_vec = memory.multi_modal_encoder.encode({"text": ["alpha"]}).vector
    second_vec = memory.multi_modal_encoder.encode({"text": ["beta"]}).vector

    batch = [
        DualRetrievalInput(
            vector=first_vec,
            query_graph=g,
            k=1),
        DualRetrievalInput(
            vector=second_vec,
            query_graph=g,
            alpha=0.7,
            beta=0.3,
            temperature=0.9),
    ]

    results = memory.batch_dual_retrieve(batch)

    expect(len(results) == 2, "should return one result per query")
    expect(all(result.iterations >= 0 for result in results),
           "iterations should be non-negative")
    expect(all(result.weights.size <= 1 for result in results[:1]),
           "weights should be at most length one for first query")
    expect(all(result.episodes for result in results),
           "each result should contain episodes")


def test_memory_batch_symbolic_search_returns_results_per_query() -> None:
    memory = Memory()

    g = nx.DiGraph()
    g.add_edge("agent", "objective", relation="pursues")
    memory.multi_modal_insert({"text": ["agent pursues objective"]}, graph=g)

    pattern = nx.DiGraph()
    pattern.add_edge("agent", "objective", relation="pursues")

    queries = [
        BatchSymbolicSearchInput(path=["agent", "objective"], limit=3),
        BatchSymbolicSearchInput(pattern=pattern),
    ]

    results = memory.batch_symbolic_search(queries)

    expect(len(results) == 2, "should return a result list per query")
    expect(all(isinstance(result, list) for result in results),
           "each result should be a list")
    expect(bool(results[0]), "first query should yield results")
    expect(bool(results[1]), "second query should yield results")


def test_stability_monitor_snapshot() -> None:
    monitor = GlobalStabilityMonitor(window=4)
    monitor.observe_retrieval([0.6, 0.4])
    monitor.observe_summary(0.8, merged=2, total=3, distortion=0.2)
    monitor.observe_trust([0.1, 0.5])
    snapshot = monitor.stability()
    expect(snapshot.summary_fidelity == pytest.approx(0.8),
           "summary fidelity should match observed value")
    expect(snapshot.summary_distortion == pytest.approx(0.2),
           "summary distortion should match observed value")
    expect(snapshot.summary_distortion_bound >= 0.0,
           "distortion bound should be non-negative")
    expect(snapshot.trust_floor == pytest.approx(0.1),
           "trust floor should reflect observed trust")
    expect(0.0 <= snapshot.contraction_factor <= 1.0,
           "contraction factor should be bounded")
    expect(isinstance(snapshot.lyapunov_warning, bool),
           "Lyapunov warning flag should be boolean")


def test_stabilise_weights_resets_per_key() -> None:
    monitor = GlobalStabilityMonitor(window=2, smoothing=0.1)
    monitor.stabilise_weights([0.9, 0.1], key="red-cube")
    blue_weights = monitor.stabilise_weights([0.1, 0.9], key="blue-sphere")

    expect(int(np.argmax(blue_weights)) == 1,
           "weights should be stabilised per key")
