"""Score predictions against ground truth: normalized edit distance and AURC.

AURC (area under the risk-coverage curve) follows Geifman & El-Yaniv (2017):
sort predictions by confidence descending, then average the running mean risk
(normalized edit distance) over all coverage levels. Lower is better. The
competition score is the mean of the date AURC and the locality AURC.

Usage:
  python evaluate.py --predictions val_predictions.csv --truth-csv ../train.csv
"""

import argparse
import csv

from common import normalized_edit_distance


def aurc(risks: list[float], confidences: list[float]) -> float:
    order = sorted(range(len(risks)), key=lambda i: -confidences[i])
    # Average risk within equal-confidence groups so the result does not
    # depend on input order among ties.
    smoothed: list[float] = []
    i = 0
    while i < len(order):
        j = i
        while j < len(order) and confidences[order[j]] == confidences[order[i]]:
            j += 1
        group = [risks[idx] for idx in order[i:j]]
        smoothed.extend([sum(group) / len(group)] * len(group))
        i = j
    total = 0.0
    running = 0.0
    for k, risk in enumerate(smoothed, start=1):
        running += risk
        total += running / k
    return total / len(smoothed)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--predictions", required=True)
    parser.add_argument("--truth-csv", required=True)
    args = parser.parse_args()

    with open(args.truth_csv, encoding="utf-8") as f:
        truth = {row["image_file"]: row for row in csv.DictReader(f)}
    with open(args.predictions, encoding="utf-8") as f:
        preds = list(csv.DictReader(f))

    preds = [p for p in preds if p["image_file"] in truth]
    if not preds:
        raise SystemExit("no overlapping image_file between predictions and truth")

    report = {}
    for field in ("verbatimDate", "verbatimLocality"):
        risks, confidences, exact = [], [], 0
        for p in preds:
            t = truth[p["image_file"]][field].strip()
            risk = normalized_edit_distance(p[field], t)
            risks.append(risk)
            confidences.append(float(p[f"{field}_confidence"]))
            exact += risk == 0.0
        report[field] = {
            "n": len(risks),
            "exact_match": round(exact / len(risks), 4),
            "mean_ned": round(sum(risks) / len(risks), 4),
            "aurc": round(aurc(risks, confidences), 4),
        }

    for field, metrics in report.items():
        print(
            f"{field}: n={metrics['n']} exact={metrics['exact_match']} "
            f"mean_ned={metrics['mean_ned']} aurc={metrics['aurc']}"
        )
    score = (report["verbatimDate"]["aurc"] + report["verbatimLocality"]["aurc"]) / 2
    print(f"competition score (mean AURC, lower is better): {round(score, 4)}")


if __name__ == "__main__":
    main()