"""Real recognition errors vs synthetic corruption on the same transliterations (reviewer 3, point 4).

For every leak-free test photograph, the recognizer output (crop recognizer, loop-cut) and the synthetic corruptions of
the same reference transliteration (noise levels of E4) are translated by the same translator. Two comparisons:
  bins    photographs grouped by recognizer CER; chrF++ retention (output / reference input) for the real output
          and for each synthetic level on the same photographs, with the synthetic CER measured on those items
  matched for each photograph the synthetic level whose CER on that item is closest to the recognizer's CER
          (only photographs with a level within 10 CER points); corpus chrF++ real vs matched synthetic, cluster
          paired bootstrap
Writes results/matched_noise.json.
"""
import json
import os
import random
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import metrics  # noqa: E402
import sacrebleu  # 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"])
OCR_TAG = os.environ.get("OCR_TAG", "ocr4b_cropcut")
CHRF = sacrebleu.CHRF(word_order=2)


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


def cer(ref, hyp):
    d, n = metrics.cer_parts(ref, hyp)
    return 100 * d / n


def chrf(hyps, refs):
    return CHRF.corpus_score(hyps, [refs]).score


ocr = {r["id"]: r for r in load(f"ocr_test_{OCR_TAG}")}
res = {}
for tr in ("mt4b", "mt4b_noise"):
    real = {r["id"]: r for r in load(f"e2e_{OCR_TAG}_{tr}_fromocr")}
    gold = {r["id"]: r for r in load(f"e2e_ocr4b_mixcut_{tr}_fromgold")}
    noise = {}
    for r in load(f"noise_{tr}"):
        noise[(r["base_id"], r["level"])] = r
    levels = sorted({lv for _, lv in noise})
    photos = [i for i in ocr if i in real and i in gold and (ocr[i]["row_id"], 0.0) in noise]
    out = {"n_photos": len(photos), "bins": {}, "matched": {}}

    def retention(ids, hyp_of):
        g = chrf([gold[i]["pred"] for i in ids], [gold[i]["tgt"] for i in ids])
        h = chrf([hyp_of(i) for i in ids], [gold[i]["tgt"] for i in ids])
        return h, g, 100 * h / g

    for lo, hi in ((0, 25), (25, 50), (50, 75), (75, 1e9)):
        ids = [i for i in photos if lo <= cer(ocr[i]["target"], ocr[i]["pred"]) < hi]
        if not ids:
            continue
        rc = sum(metrics.cer_parts(ocr[i]["target"], ocr[i]["pred"])[0] for i in ids) / \
            sum(metrics.cer_parts(ocr[i]["target"], ocr[i]["pred"])[1] for i in ids) * 100
        h, g, ret = retention(ids, lambda i: real[i]["pred"])
        b = {"n": len(ids), "real": {"cer": rc, "chrfpp": h, "chrfpp_gold": g, "retention": ret}}
        for lv in levels:
            if lv == 0.0:
                continue
            rows = [noise[(ocr[i]["row_id"], lv)] for i in ids]
            sc = sum(metrics.cer_parts(r["src_clean"], r["src"])[0] for r in rows) / \
                sum(metrics.cer_parts(r["src_clean"], r["src"])[1] for r in rows) * 100
            hs, gs, rs = retention(ids, lambda i, lv=lv: noise[(ocr[i]["row_id"], lv)]["pred"])
            b[f"synthetic_{lv}"] = {"cer": sc, "chrfpp": hs, "retention": rs}
        out["bins"][f"{lo}-{hi if hi < 1e9 else 'inf'}"] = b

    # per-photo CER matching
    pairs = []
    for i in photos:
        rc = cer(ocr[i]["target"], ocr[i]["pred"])
        best = min((abs(cer(noise[(ocr[i]["row_id"], lv)]["src_clean"], noise[(ocr[i]["row_id"], lv)]["src"]) - rc), lv)
                   for lv in levels if lv > 0)
        if best[0] <= 10:
            pairs.append((i, best[1]))
    ids = [i for i, _ in pairs]
    lv_of = dict(pairs)
    syn = {i: noise[(ocr[i]["row_id"], lv_of[i])] for i in ids}
    real_cer = 100 * sum(metrics.cer_parts(ocr[i]["target"], ocr[i]["pred"])[0] for i in ids) / \
        sum(metrics.cer_parts(ocr[i]["target"], ocr[i]["pred"])[1] for i in ids)
    syn_cer = 100 * sum(metrics.cer_parts(syn[i]["src_clean"], syn[i]["src"])[0] for i in ids) / \
        sum(metrics.cer_parts(syn[i]["src_clean"], syn[i]["src"])[1] for i in ids)
    a = [dict(real[i], id=i) for i in ids]
    s = [dict(syn[i], id=i, tgt=real[i]["tgt"], cluster=real[i]["cluster"]) for i in ids]
    h_real, g, ret_real = retention(ids, lambda i: real[i]["pred"])
    h_syn, _, ret_syn = retention(ids, lambda i: syn[i]["pred"])
    out["matched"] = {"n": len(ids), "real_cer": real_cer, "synthetic_cer": syn_cer, "chrfpp_gold": g,
                      "real_chrfpp": h_real, "synthetic_chrfpp": h_syn, "real_retention": ret_real,
                      "synthetic_retention": ret_syn, "paired_real_gt_synthetic": metrics.paired_bootstrap_chrf(a, s)}
    res[tr] = out
json.dump(res, open(f"{ROOT}/results/matched_noise.json", "w"), indent=1)
for tr, out in res.items():
    print(f"=== {tr}: {out['n_photos']} photos")
    for k, b in out["bins"].items():
        print(f"  CER {k:7s} n={b['n']:3d} real CER {b['real']['cer']:.1f} ret {b['real']['retention']:.0f}% | " +
              " ".join(f"{lv[10:]}: CER {v['cer']:.1f} ret {v['retention']:.0f}%" for lv, v in b.items() if lv.startswith("synthetic")))
    m = out["matched"]
    print(f"  matched n={m['n']} CER real {m['real_cer']:.1f} syn {m['synthetic_cer']:.1f} | chrF gold {m['chrfpp_gold']:.1f} "
          f"real {m['real_chrfpp']:.1f} ({m['real_retention']:.0f}%) syn {m['synthetic_chrfpp']:.1f} ({m['synthetic_retention']:.0f}%) "
          f"p={m['paired_real_gt_synthetic']['p']:.3f}")