"""
evaluate.py

Reproduce the retrieval metrics on the synthetic conversation.
"""

import numpy as np
from collections import Counter

from hv_context import (
    HVContext, generate_queries, DEFAULT_K,
)


def evaluate_retrieval(selector, queries, k=DEFAULT_K):
    precisions, recalls = [], []
    topic_counts = Counter(selector.topics)
    for true_topic, query in queries:
        idx, _ = selector.retrieve(query, k)
        retrieved = [selector.topics[i] for i in idx]
        n_rel = sum(1 for t in retrieved if t == true_topic)
        precisions.append(n_rel / len(retrieved))
        recalls.append(n_rel / topic_counts[true_topic])
    return float(np.mean(precisions)), float(np.mean(recalls))


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


def evaluate_recent(selector, queries, k=DEFAULT_K):
    n = len(selector.topics)
    idx = list(range(n - k, n))
    topics = [selector.topics[i] for i in idx]
    topic_counts = Counter(selector.topics)
    precisions, recalls = [], []
    for true_topic, _ in queries:
        n_rel = sum(1 for t in topics if t == true_topic)
        precisions.append(n_rel / k)
        recalls.append(n_rel / topic_counts[true_topic])
    return float(np.mean(precisions)), float(np.mean(recalls))


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

    print("Generating held-out queries...")
    queries = generate_queries(n_per_topic=4, seed=99)

    prec_hv, rec_hv = evaluate_retrieval(model, queries, DEFAULT_K)
    prec_r, rec_r = evaluate_random(model, queries, DEFAULT_K)
    prec_rec, rec_rec = evaluate_recent(model, queries, DEFAULT_K)

    print()
    print(f"{'method':<20} {'precision@5':>12} {'recall@5':>10}")
    print("-" * 46)
    print(f"{'random':<20} {prec_r:>12.3f} {rec_r:>10.3f}")
    print(f"{'most recent':<20} {prec_rec:>12.3f} {rec_rec:>10.3f}")
    print(f"{'HV context':<20} {prec_hv:>12.3f} {rec_hv:>10.3f}")
    print()
    print(f"HV / random: {prec_hv / max(prec_r, 1e-9):.1f}x")


if __name__ == "__main__":
    main()