"""Step 6: pick the best take of every segment (work/pick.json, work/pick_report.json).

A take passes when it has no error, a pace of 7-19 chars/s and
- segments of up to 3 words (headings, names, signatures): CER <= 0.35 (ASR mangles lone names);
- longer ones: WER <= 0.12 or a single wrong word.
Among passing takes: score = 6*stress + 0.5*MOS + 2*SIM - 3*WER (stress first: a wrong stress is what
a listener hears). No passing take: the lowest WER, and the segment is listed as weak — `render --weak`
speaks it in voice.extra_seeds, then qc and pick again.
"""
import collections
import glob
import json
import re

import jiwer

from .config import load


def norm(t):
    return " ".join(re.sub(r"[^\w\s]", " ", t.lower().replace("ё", "е")).split())


def passes(s, r):
    if "error" in r or not 7 <= r["cps"] <= 19:
        return False
    if len(s["text"].split()) <= 3:
        h = norm(r.get("hyp", ""))
        return bool(h) and jiwer.cer(norm(s["text"]), h) <= 0.35
    return r["wer"] <= 0.12 or r["wer"] * len(norm(s["text"]).split()) <= 1.01


def score(r):
    return 6 * r["stress_ok"] / max(1, r["stress_n"]) + 0.5 * r.get("mos", 3) + 2 * r["sim"] - 3 * r["wer"]


def main():
    cfg = load()
    segs = json.load(open(cfg.w("segs.json"), encoding="utf-8"))
    qc = {}
    for f in sorted(glob.glob(str(cfg.w("qc", "x").parent / "*.json"))):
        qc.update(json.load(open(f, encoding="utf-8")))
    by = collections.defaultdict(dict)
    for stem, r in qc.items():
        take, seed = stem.rsplit("_s", 1)
        by[take][int(seed)] = r
    pick, weak, missing = {}, [], []
    for s in segs:
        takes = by.get(s["take"], {})
        if not takes:
            missing.append(s["id"]); continue
        ok = {sd: r for sd, r in takes.items() if passes(s, r)}
        if ok:
            pick[s["id"]] = max(ok, key=lambda sd: score(ok[sd]))
        else:
            pick[s["id"]] = min(takes, key=lambda sd: takes[sd]["wer"])
            weak.append(s["id"])
    json.dump(pick, open(cfg.w("pick.json"), "w"))
    take_of = {s["id"]: s["take"] for s in segs}
    P = [by[take_of[i]][sd] for i, sd in pick.items()]
    rep = dict(segments=len(segs), picked=len(pick), weak=weak, missing=missing,
               wer=round(sum(r["wer"] for r in P) / max(1, len(P)), 4),
               mos=round(sum(r.get("mos", 0) for r in P) / max(1, len(P)), 3),
               sim=round(sum(r["sim"] for r in P) / max(1, len(P)), 3),
               stress=round(sum(r["stress_ok"] for r in P) / max(1, sum(r["stress_n"] for r in P)), 4),
               seeds=dict(collections.Counter(pick.values())))
    json.dump(rep, open(cfg.w("pick_report.json"), "w"), indent=1)
    print({k: (len(v) if isinstance(v, list) else v) for k, v in rep.items()})
    if missing:
        print(f"{len(missing)} segments have no scored take: run render and qc first")


if __name__ == "__main__":
    main()