#!/usr/bin/env python3
"""Tests for the retrieval metrics. Pure python, no GPU, no model."""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from src.evaluate_retrieval import aggregate, dcg, score_ranking

PASS = FAIL = 0
def chk(c, l, d=""):
    global PASS, FAIL
    if c: PASS += 1; print(f"  ok   {l}")
    else:  FAIL += 1; print(f"  FAIL {l}  {d}")

P = {"p1", "p2"}

print("[perfect ranking]")
r = score_ranking([{"p1"}, {"x"}, {"y"}, {"z"}, {"w"}], P)
chk(r["p1"] == 1.0, "P@1 = 1")
chk(r["mrr"] == 1.0, "MRR = 1")
chk(r["hit"] == 1.0, "Hit = 1")
chk(r["ndcg"] == 1.0, "nDCG = 1", r["ndcg"])
chk(r["zero_hit"] == 0.0, "zero_hit = 0")

print("\n[positive at rank 3]")
r = score_ranking([{"x"}, {"y"}, {"p2"}, {"z"}, {"w"}], P)
chk(r["p1"] == 0.0, "P@1 = 0")
chk(abs(r["mrr"] - 1/3) < 1e-9, "MRR = 1/3", r["mrr"])
chk(r["hit"] == 1.0, "still a hit")
chk(0 < r["ndcg"] < 1, "nDCG between 0 and 1", r["ndcg"])

print("\n[complete miss]")
r = score_ranking([{"a"}, {"b"}, {"c"}, {"d"}, {"e"}], P)
chk(r["hit"] == 0.0 and r["mrr"] == 0.0, "no hit, no reciprocal rank")
chk(r["zero_hit"] == 1.0, "zero_hit = 1  <- the metric that exposes failures")
chk(r["ndcg"] == 0.0, "nDCG = 0")

print("\n[a chunk spanning several passages counts if ANY is positive]")
r = score_ranking([{"x", "y", "p2"}], P)
chk(r["p1"] == 1.0, "multi-block chunk hits", r)

print("\n[edge cases]")
chk(score_ranking([], P)["hit"] == 0.0, "empty ranking")
chk(score_ranking([{"p1"}], set())["hit"] == 0.0, "no positives -> no hit")
chk(dcg([]) == 0.0, "dcg of empty list")
chk(abs(dcg([1, 1]) - (1 + 1/1.5849625)) < 1e-6, "dcg discount correct", dcg([1, 1]))

print("\n[aggregate]")
qs = [score_ranking([{"p1"}], P), score_ranking([{"z"}], P)]
a = aggregate(qs)
chk(a["n_queries"] == 2, "counts queries")
chk(a["hit"] == 0.5, "mean hit = 0.5", a["hit"])
chk(a["zero_hit"] == 0.5, "zero_hit is 1 - hit", a["zero_hit"])
chk("ndcg_std" in a, "reports per-query variance")
chk(aggregate([]) == {}, "empty input")

print("\n[k truncation]")
far = [{"x"}] * 5 + [{"p1"}]
chk(score_ranking(far, P, k=5)["hit"] == 0.0, "positive beyond k is not a hit")
chk(score_ranking(far, P, k=6)["hit"] == 1.0, "found when k is large enough")


print("\n[context-budget scoring — the fair comparison]")
from src.evaluate_retrieval import score_at_budget

# strategy A: 2 big chunks (100 words each). positive is 2nd.
a = score_at_budget([{"x"}, {"p1"}], [100, 100], P, budget_words=400)
# strategy B: 8 small chunks (25 words). positive is 6th.
b = score_at_budget([{"x"}]*5 + [{"p1"}] + [{"y"}]*2, [25]*8, P, budget_words=400)
chk(a["b_hit"] == 1.0, "A finds it within budget")
chk(b["b_hit"] == 1.0, "B finds it within budget too — both get 400 words", b)
chk(b["b_chunks_used"] > a["b_chunks_used"], "B uses more chunks for the same words",
    (a["b_chunks_used"], b["b_chunks_used"]))

print("\n[budget actually binds]")
tight = score_at_budget([{"x"}, {"p1"}], [100, 100], P, budget_words=150)
chk(tight["b_hit"] == 0.0, "second chunk exceeds 150w -> not retrieved", tight)
chk(tight["b_words_used"] == 100.0, "used exactly the first chunk", tight["b_words_used"])

print("\n[a single oversized chunk is still taken]")
big = score_at_budget([{"p1"}], [900], P, budget_words=400)
chk(big["b_hit"] == 1.0, "never return an empty context", big)
chk(big["b_chunks_used"] == 1.0, "exactly one chunk")

print("\n[fixed-k rewards emitting fewer chunks; budget does not]")
# few-chunk strategy: k=5 sees all 3 of its chunks
few = score_ranking([{"x"}, {"p1"}, {"y"}], P, k=5)
# many-chunk strategy: positive sits at rank 7, outside k=5
many = score_ranking([{"x"}]*6 + [{"p1"}], P, k=5)
chk(few["hit"] == 1.0 and many["hit"] == 0.0, "fixed-k penalises the finer strategy")
few_b = score_at_budget([{"x"}, {"p1"}, {"y"}], [130]*3, P, 400)
many_b = score_at_budget([{"x"}]*6 + [{"p1"}], [55]*7, P, 400)
chk(few_b["b_hit"] == 1.0 and many_b["b_hit"] == 1.0,
    "equal word budget lets both find it", (few_b["b_hit"], many_b["b_hit"]))

print("\n[aggregate carries the budget keys through]")
agg = aggregate([a, b])
chk("b_ndcg" in agg and "b_hit" in agg, "budget metrics aggregated", list(agg))
chk("b_words_used" in agg, "word usage reported")

print(f"\n{'='*50}\n  {PASS} passed, {FAIL} failed\n{'='*50}")
sys.exit(1 if FAIL else 0)