"""Personal names that reach the translation from a photograph (reviewer 2, point 2; reviewer 1, point 1).

Name recall (capitalized words of the reference translation other than a stop list, as in analyze.py) for the
4B translator on reference transliterations and on the crop recognizer's output; by recognizer CER and for the most
confidently read photographs (mean token NLL of the recognizer output, as in Appendix C). Writes results/names_e2e.json.
"""
import json
import os
import re
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import metrics  # noqa: E402

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
P = f"{ROOT}/preds"
KEEP = set(json.load(open(f"{ROOT}/data/leakfree.json"))["keep_test"])
STOP = {"May", "God", "Christ", "The", "This", "These", "He", "She", "They", "His", "Her", "And", "In", "I", "A",
        "Holy", "Mary", "Lord", "Jesus", "Amen", "Saint", "St", "Gud", "Han", "Hon", "De", "Och", "Den", "Denna",
        "Detta", "Dessa", "Kristus"}


def names(t):
    return {w for w in re.findall(r"\b[A-ZÁÉÍÓÚÝÞÆØÖÄÅ][\w'-]+", t or "") if w not in STOP}


def load(n):
    return {r["id"]: r for r in metrics.read(f"{P}/{n}.jsonl") if (r.get("row_id") or r["id"]) in KEEP}


def recall(rows):
    hit = tot = 0
    for r in rows:
        ref = names(r["tgt"])
        tot += len(ref)
        hit += len(ref & names(r.get("pred")))
    return dict(n=len(rows), names=tot, recall=100 * hit / max(tot, 1))


def precision(rows):
    """Share of names in the output that are names of the reference (the rest are wrong or invented names)."""
    hit = tot = 0
    for r in rows:
        out = names(r.get("pred"))
        tot += len(out)
        hit += len(out & names(r["tgt"]))
    return 100 * hit / max(tot, 1)


ocr = load("ocr_test_ocr4b_cropcut")
fo = load("e2e_ocr4b_cropcut_mt4b_fromocr")
fg = load("e2e_ocr4b_mixcut_mt4b_fromgold")
ids = [i for i in ocr if i in fo and i in fg]
cer = {i: 100 * metrics.cer_parts(ocr[i]["target"], ocr[i]["pred"])[0] / metrics.cer_parts(ocr[i]["target"], ocr[i]["pred"])[1] for i in ids}
res = {"all": {"from_reference": recall([fg[i] for i in ids]), "from_photo": recall([fo[i] for i in ids]),
               "precision_from_reference": precision([fg[i] for i in ids]), "precision_from_photo": precision([fo[i] for i in ids])}}
for lo, hi in ((0, 25), (25, 50), (50, 75), (75, 1e9)):
    s = [i for i in ids if lo <= cer[i] < hi]
    res[f"cer_{lo}_{hi if hi < 1e9 else 'inf'}"] = {"from_reference": recall([fg[i] for i in s]), "from_photo": recall([fo[i] for i in s]),
                                                     "precision_from_photo": precision([fo[i] for i in s])}
sc = {r["id"]: r for r in metrics.read(f"{P}/score_test_ocr4b_cropcut.jsonl")}
key = next(k for k in ("nll", "mean_nll", "score") if k in next(iter(sc.values())))
order = sorted(ids, key=lambda i: (sc[i][key] if i in sc and sc[i][key] is not None else 1e9))
for f in (0.1, 0.2, 0.3, 0.5):
    s = order[: round(f * len(order))]
    res[f"top_{f}"] = {"from_reference": recall([fg[i] for i in s]), "from_photo": recall([fo[i] for i in s]),
                        "precision_from_photo": precision([fo[i] for i in s])}
json.dump(res, open(f"{ROOT}/results/names_e2e.json", "w"), indent=1)
for k, v in res.items():
    print(f"{k:12s} " + " ".join(f"{a}={b['recall']:.1f}(n={b['n']},names={b['names']})" if isinstance(b, dict) else f"{a}={b:.1f}" for a, b in v.items()))