"""Metrics with 95% percentile-bootstrap confidence intervals.

  OCR: CER and WER (corpus-level: total edits / total reference length), exact match
  MT : chrF++ (sacrebleu CHRF word_order=2) and BLEU, corpus-level
Paired bootstrap p-values compare two systems on the same items.

Resampling unit: inscription clusters (all items of a cluster are drawn together), so that several images or
pairs of one inscription do not count as independent evidence. BOOT_UNIT=item restores item-level resampling.
Items without a "cluster" field (synthetic images) form their own cluster.
"""
import os
import json
import random
import sys

import sacrebleu
from rapidfuzz.distance import Levenshtein

B = int(os.environ.get("BOOT_B", 1000))
UNIT = os.environ.get("BOOT_UNIT", "cluster")


def _groups(rows):
    """Index groups that are resampled together."""
    if UNIT == "item":
        return [[i] for i in range(len(rows))]
    by = {}
    for i, r in enumerate(rows):
        c = r.get("cluster")
        key = ("c", c) if c is not None else ("i", i)
        by.setdefault(key, []).append(i)
    return list(by.values())


def _draw(groups, rng):
    """One bootstrap replicate: as many clusters as there are, drawn with replacement."""
    idx = []
    for _ in range(len(groups)):
        idx.extend(groups[rng.randrange(len(groups))])
    return idx


def read(path):
    return [json.loads(l) for l in open(path) if l.strip()]


def cer_parts(ref, hyp):
    return Levenshtein.distance(ref, hyp), max(len(ref), 1)


def wer_parts(ref, hyp):
    r, h = ref.split(), hyp.split()
    return Levenshtein.distance(r, h), max(len(r), 1)


def ocr_scores(rows, ref_key="target", hyp_key="pred", seed=0):
    cp = [cer_parts(r[ref_key], r[hyp_key]) for r in rows]
    wp = [wer_parts(r[ref_key], r[hyp_key]) for r in rows]
    em = [float(r[ref_key] == r[hyp_key]) for r in rows]

    def agg(idx):
        return (100 * sum(cp[i][0] for i in idx) / sum(cp[i][1] for i in idx),
                100 * sum(wp[i][0] for i in idx) / sum(wp[i][1] for i in idx),
                100 * sum(em[i] for i in idx) / len(idx))

    n = len(rows)
    point = agg(range(n))
    rng = random.Random(seed)
    g = _groups(rows)
    boots = [agg(_draw(g, rng)) for _ in range(B)]
    ci = []
    for k in range(3):
        v = sorted(b[k] for b in boots)
        ci.append((v[int(0.025 * B)], v[int(0.975 * B)]))
    return dict(n=n, cer=point[0], cer_ci=ci[0], wer=point[1], wer_ci=ci[1], exact=point[2])


def mt_stats(hyps, refs):
    chrf = sacrebleu.CHRF(word_order=2)
    bleu = sacrebleu.BLEU()
    return chrf, bleu


def mt_scores(rows, hyp_key="pred", ref_key="tgt", seed=0, boot=True):
    hyps = [r.get(hyp_key) or "" for r in rows]
    refs = [r[ref_key] for r in rows]
    chrf = sacrebleu.CHRF(word_order=2)
    bleu = sacrebleu.BLEU()
    out = dict(n=len(rows), chrfpp=chrf.corpus_score(hyps, [refs]).score, bleu=bleu.corpus_score(hyps, [refs]).score,
               empty=sum(1 for h in hyps if not h.strip()))
    if boot and len(rows) > 1:
        # sentence-level statistics make the bootstrap cheap: resample and re-aggregate
        cs = [chrf._extract_corpus_statistics([h], [[r]]) for h, r in zip(hyps, refs)]
        bs = [bleu._extract_corpus_statistics([h], [[r]]) for h, r in zip(hyps, refs)]
        rng = random.Random(seed)
        g = _groups(rows)
        vc, vb = [], []
        for _ in range(B):
            idx = _draw(g, rng)
            vc.append(chrf._compute_score_from_stats([sum(x) for x in zip(*[cs[i][0] for i in idx])]).score)
            vb.append(bleu._compute_score_from_stats([sum(x) for x in zip(*[bs[i][0] for i in idx])]).score)
        vc.sort()
        vb.sort()
        out["chrfpp_ci"] = (vc[int(0.025 * B)], vc[int(0.975 * B)])
        out["bleu_ci"] = (vb[int(0.025 * B)], vb[int(0.975 * B)])
    return out


def paired_bootstrap_chrf(rows_a, rows_b, hyp_key="pred", ref_key="tgt", seed=0):
    """p-value for H0: system A is not better than B on chrF++ (aligned by id)."""
    b_by = {r["id"]: r for r in rows_b}
    common = [r for r in rows_a if r["id"] in b_by]
    chrf = sacrebleu.CHRF(word_order=2)
    sa = [chrf._extract_corpus_statistics([r.get(hyp_key) or ""], [[r[ref_key]]])[0] for r in common]
    sb = [chrf._extract_corpus_statistics([b_by[r["id"]].get(hyp_key) or ""], [[r[ref_key]]])[0] for r in common]
    n = len(common)

    def score(stats, idx):
        return chrf._compute_score_from_stats([sum(x) for x in zip(*[stats[i] for i in idx])]).score

    delta = score(sa, range(n)) - score(sb, range(n))
    rng = random.Random(seed)
    g = _groups(common)
    worse = 0
    for _ in range(B):
        idx = _draw(g, rng)
        if score(sa, idx) - score(sb, idx) <= 0:
            worse += 1
    return dict(n=n, delta=delta, p=worse / B, n_clusters=len(g))


if __name__ == "__main__":
    task, path = sys.argv[1], sys.argv[2]
    rows = read(path)
    print(json.dumps(ocr_scores(rows) if task == "ocr" else mt_scores(rows), indent=1))