#!/usr/bin/env python3
"""Score run.py outputs. Usage: score.py OUT_DIR

Moment hit@1: the first returned moment (a timestamp, or the midpoint of a clip)
falls inside the gold span widened by TOL seconds on each side; a moment only
counts when the test video was also the top-ranked item. Item@1: the test video
was the top-ranked item in the library. Answer correct: the returned answer
text matches the query's answer_pattern (case-insensitive regex).
"""
import glob
import json
import os
import re
import sys

TOL = 15.0
HERE = os.path.dirname(os.path.abspath(__file__))
QUERIES = {q["id"]: q for q in json.load(open(os.path.join(HERE, "queries.json")))["queries"]}


def moment_start(m):
    return (m[0] + m[1]) / 2 if isinstance(m, list) else m


def in_gold(sec, gold):
    return gold[0] - TOL <= sec <= gold[1] + TOL


def main(out):
    table = {}
    for path in sorted(glob.glob(os.path.join(out, "*.json"))):
        if path.endswith("results.json"):
            continue
        d = json.load(open(path))
        n = len(QUERIES)
        item = hit1 = hit3 = ans = 0
        n_ans = n_mom = 0
        lats = []
        for qid, q in QUERIES.items():
            r = d["queries"].get(qid, {})
            if r.get("item_rank") == 1:
                item += 1
            ms = r.get("moments") or []
            if ms:
                n_mom += 1
            if ms and r.get("item_rank") == 1:
                if in_gold(moment_start(ms[0]), q["gold"]):
                    hit1 += 1
                if any(in_gold(moment_start(m), q["gold"]) for m in ms[:3]):
                    hit3 += 1
            if r.get("answer"):
                n_ans += 1
                if re.search(q["answer_pattern"], r["answer"], re.I):
                    ans += 1
            if r.get("latency_s") is not None:
                lats.append(r["latency_s"])
        table[d["provider"]] = {
            "queries": n,
            "item_at_1": item,
            "moment_hit_at_1": hit1 if n_mom else None,
            "moment_hit_at_3": hit3 if n_mom else None,
            "answers_correct": ans if n_ans else None,
            "answers_given": n_ans,
            "mean_search_latency_s": round(sum(lats) / len(lats), 2) if lats else None,
            "library_items": d.get("library_items"),
            "version": d.get("version"),
        }
    with open(os.path.join(out, "results.json"), "w") as f:
        json.dump({"tolerance_s": TOL, "providers": table}, f, indent=1)
    print(f"{'provider':12} {'item@1':>7} {'mom@1':>6} {'mom@3':>6} {'answers':>8} {'lat s':>6}")
    for p, r in table.items():
        f = lambda v: "-" if v is None else str(v)
        print(
            f"{p:12} {r['item_at_1']:>4}/{r['queries']:<2} {f(r['moment_hit_at_1']):>6} "
            f"{f(r['moment_hit_at_3']):>6} {f(r['answers_correct']):>4}/{r['answers_given']:<3} "
            f"{f(r['mean_search_latency_s']):>6}"
        )


if __name__ == "__main__":
    main(sys.argv[1])