"""
test_learned_retrieval.py
-------------------------
Test the PyTorch-based LearnedRetriever (Neural Ranker).
"""
import numpy as np
from hieromem_core.memory import Memory
from hieromem_core.retrieval.learned import LearnedRetriever, LearnedConfig
from hieromem_core.evaluation.framework import EvaluationExample
from hieromem_core.encoder import Encoder, EncoderConfig
from hieromem_core.persistence.in_memory_graph_store import InMemoryGraphStore


def get_stub_encoder():
    config = EncoderConfig(model_name="stub-hash", backend="stub")
    return Encoder(config)


def test_learned_retriever_training():
    """
    Test that LearnedRetriever:
    1. Initializes RankerNet.
    2. Trains on examples without error.
    3. Improves ranking (or at least changes weights) after training.
    """
    store = InMemoryGraphStore()
    memory = Memory(store=store, Nmax=20, encoder=get_stub_encoder())

    # 1. Insert Synthetic Data
    # 'Relevant' items: High dense score logic (we enforce vectors)
    # Target: Query [1, 0, ...] matches [1, 0, ...]

    # Ep1: Relevant (Target)
    vec1 = np.array([1.0, 0.0, 0.0, 0.0])
    ep1 = memory.insert(vec1, meta={"id": "ep1", "type": "target", "trust": 0.9})

    # Ep2: Irrelevant (Noise)
    vec2 = np.array([0.0, 1.0, 0.0, 0.0])
    ep2 = memory.insert(vec2, meta={"id": "ep2", "type": "noise", "trust": 0.1})

    # Ep3: Semi-Relevant
    vec3 = np.array([0.5, 0.5, 0.0, 0.0])
    memory.insert(vec3, meta={"id": "ep3", "type": "semi", "trust": 0.5})

    retriever = LearnedRetriever(memory)

    # Check initial prediction score
    # Artificially create a feature matrix to test forward pass
    feats = np.array([[1.0, 0.0, 1.0, 0.0, 1.0]], dtype=np.float32)
    initial_score = retriever.predict_scores(feats)
    assert initial_score.shape == (1,)  # nosec B101

    # 2. Create Training Examples
    # Query: [1, 0, ...] -> Relevant: ep1
    ex1 = EvaluationExample(
        query="test query",
        query_vector=np.array([1.0, 0.0, 0.0, 0.0]),
        relevant_ids={ep1.id}
    )

    # Query: [0, 1, ...] -> Relevant: ep2
    ex2 = EvaluationExample(
        query="noise query",
        query_vector=np.array([0.0, 1.0, 0.0, 0.0]),
        relevant_ids={ep2.id}
    )

    examples = [ex1, ex2] * 10  # Repeat to have enough steps

    # 3. Train
    config = LearnedConfig(epochs=20, learning_rate=0.01, batch_size=4)
    retriever.train(examples, config)

    # 4. Verify Retrieval
    # Query for Ep1
    q_vec = np.array([1.0, 0.0, 0.0, 0.0])
    cfg_test = LearnedConfig(k=3)
    result = retriever.retrieve(q_vec, config=cfg_test)

    print("Scores:")
    for i, ep in enumerate(result.episodes):
        print(f"  {ep.id}: {result.weights[i]:.4f}")

    # Ep1 should be top 1
    assert result.episodes[0].id == ep1.id  # nosec B101

    print("Training and Retrieval verification passed.")


if __name__ == "__main__":
    test_learned_retriever_training()
