"""
evaluate.py

Reproduce the retrieval metrics.
"""

import numpy as np
from collections import Counter

from hv_temporal_context import HVTemporalContext, DEFAULT_K
from train_temporal_context import load_turns, _fill, TOPIC_TEMPLATES


def generate_queries(seed=99, n_per_topic=4):
    rng = np.random.default_rng(seed)
    queries = []
    for topic, tpl in TOPIC_TEMPLATES.items():
        for _ in range(n_per_topic):
            u = _fill(rng.choice(tpl["user"]), tpl["fills"], rng)
            queries.append((topic, u))
    rng.shuffle(queries)
    return queries


def evaluate_similarity(model, queries, k=DEFAULT_K):
    precisions = []
    for true_topic, query in queries:
        idx, _ = model.retrieve(query, k)
        topics = [model.topics[i] for i in idx]
        precisions.append(sum(1 for t in topics if t == true_topic) / len(topics))
    return float(np.mean(precisions))


def evaluate_temporal(model, queries, k=DEFAULT_K, lambda_surprise=0.5):
    precisions = []
    for true_topic, query in queries:
        idx, _ = model.retrieve_temporal(query, true_topic, k, lambda_surprise)
        topics = [model.topics[i] for i in idx]
        precisions.append(sum(1 for t in topics if t == true_topic) / len(topics))
    return float(np.mean(precisions))


def evaluate_random(model, queries, k=DEFAULT_K, seed=0):
    rng = np.random.default_rng(seed)
    precisions = []
    for true_topic, _ in queries:
        idx = rng.choice(len(model.topics), size=k, replace=False)
        topics = [model.topics[i] for i in idx]
        precisions.append(sum(1 for t in topics if t == true_topic) / len(topics))
    return float(np.mean(precisions))


def main(model_path="temporal_context_model"):
    print("Loading model...")
    model = HVTemporalContext.load(model_path)
    print(f"  {len(model.texts)} turns, "
          f"{len(model.vocab_index)} vocabulary")

    queries = generate_queries(seed=99, n_per_topic=4)

    p_rand = evaluate_random(model, queries)
    p_sim = evaluate_similarity(model, queries)
    p_temp = evaluate_temporal(model, queries)

    print()
    print(f"  {'method':<28} {'precision@5':>12}")
    print("  " + "-" * 42)
    print(f"  {'random':<28} {p_rand:>12.3f}")
    print(f"  {'similarity':<28} {p_sim:>12.3f}")
    print(f"  {'position-conditional':<28} {p_temp:>12.3f}")
    print()
    print(f"  gap over similarity:  {p_temp - p_sim:+.3f}")


if __name__ == "__main__":
    main()