"""Recognizer choice on the validation photographs (Reviewer 3, comment 2): CER, rare-word recall and end-to-end
chrF++ (4B LoRA translator on the truncated recognizer output) for every real-trained recognizer, on the
leak-free validation images; paired cluster bootstrap of each recognizer against the crop recognizer.
Writes results/valsel.json."""
import json
import os
import statistics
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import metrics  # noqa: E402
from analyze import load, domain, rare_word_recall  # noqa: E402

ROOT = os.path.expanduser("~/runic_access")
SYSTEMS = [  # tag, recognizer, test-time input
    ("ocr4b_real", "real photographs only", "whole"),
    ("ocr4b_mix", "real + synthetic (mix)", "whole"),
    ("ocr4b_diffmix", "real + diffusion", "whole"),
    ("ocr4b_mix_hr", "mix, higher resolution", "whole"),
    ("ocr4b_mix_oncrop", "mix, applied to crops", "crop"),
    ("ocr4b_crop", "crop recognizer", "crop"),
]
DOMS = ("gr", "rdb_photo", "card")
out = {}
gold_mt = {r["id"]: r for r in load("e2eval_mt4b_fromgold")}
out["reference_input"] = {"all": metrics.mt_scores(list(gold_mt.values()))}
base = {r["id"]: r for r in load("e2eval_ocr4b_cropcut_mt4b_fromocr")}
for tag, name, inp in SYSTEMS:
    ocr = load(f"ocr_val_{tag}cut")
    e2e = load(f"e2eval_{tag}cut_mt4b_fromocr")
    if not ocr or not e2e:
        print("missing", tag)
        continue
    cers = [metrics.cer_parts(r["target"], r["pred"]) for r in ocr]
    d = dict(name=name, input=inp, n_images=len(ocr), ocr=metrics.ocr_scores(ocr),
             median_cer=statistics.median(100 * e / max(n, 1) for e, n in cers),
             rare_recall=rare_word_recall(ocr), e2e=metrics.mt_scores(e2e), by_domain={})
    for dom in DOMS:
        o = [r for r in ocr if domain(r) == dom]
        e = [r for r in e2e if domain(r) == dom]
        if o:
            d["by_domain"][dom] = dict(n=len(o), cer=metrics.ocr_scores(o)["cer"],
                                       chrfpp=metrics.mt_scores(e, boot=False)["chrfpp"],
                                       ref_chrfpp=metrics.mt_scores([gold_mt[r["id"]] for r in e], boot=False)["chrfpp"])
    if tag != "ocr4b_crop":
        ids = [r["id"] for r in e2e if r["id"] in base]
        a = [base[i] for i in ids]
        b = [{r["id"]: r for r in e2e}[i] for i in ids]
        d["paired_crop_minus_this"] = metrics.paired_bootstrap_chrf(a, b)
    d["retention_vs_reference"] = 100 * d["e2e"]["chrfpp"] / out["reference_input"]["all"]["chrfpp"]
    out[tag] = d
    print(f"{tag:18s} n={len(ocr)} CER {d['ocr']['cer']:.1f} [{d['ocr']['cer_ci'][0]:.1f},{d['ocr']['cer_ci'][1]:.1f}] "
          f"med {d['median_cer']:.1f} rare {d['rare_recall']:.1f} e2e {d['e2e']['chrfpp']:.1f} "
          f"[{d['e2e']['chrfpp_ci'][0]:.1f},{d['e2e']['chrfpp_ci'][1]:.1f}] "
          + " ".join(f"{k}:{v['chrfpp']:.1f}" for k, v in d["by_domain"].items())
          + (f" p(crop>this)={d['paired_crop_minus_this'].get('p')}" if "paired_crop_minus_this" in d else ""))
print("reference input", round(out["reference_input"]["all"]["chrfpp"], 1))
json.dump(out, open(f"{ROOT}/results/valsel.json", "w"), indent=1)