import json, random
from pathlib import Path
from experiment import build_graph, name_match, distances, units_for, render


def run():
    names, desc, edges = build_graph(n=48, seed=23)
    queries = [
        f"Explain {names[i]} and {names[j]}" for i, j in
        [(1, 18), (4, 27), (8, 35), (12, 43), (20, 31), (2, 40), (15, 29), (6, 22)]
    ]
    rows = []
    for qi, q in enumerate(queries):
        seeds = name_match(q, names)
        d = distances(names, edges, seeds, 2)
        units = units_for(names, desc, edges, d)
        # Relevant evidence is the complete seed-local one-hop region.
        gold = {u.ident for u in units if u.hop <= 1}
        gold_len = sum(u.length for u in units if u.hop <= 1)
        D = sum(u.length for u in units)
        # A budget that fits the relevant region but not generally the full region.
        budget = gold_len
        row = {"query": q, "D": D, "gold_units": len(gold), "budget": budget}
        for mode in ("anchored", "global", "random"):
            _, chosen, used = render(units, budget, mode, random.Random(700 + qi))
            got = {u.ident for u in chosen}
            row[mode] = {"recall": len(got & gold) / len(gold), "units": len(chosen), "used": used}
        rows.append(row)
    avg = {m: sum(r[m]["recall"] for r in rows) / len(rows) for m in ("anchored", "global", "random")}
    exact = {m: sum(r[m]["recall"] == 1.0 for r in rows) for m in avg}
    report = {"queries": len(rows), "rows": rows, "average_relevant_recall": avg,
              "queries_with_full_relevant_recall": exact}
    Path("mini_results.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == "__main__":
    run()