"""Choose the reference clip by numbers: 16 note paragraphs x 2 seeds in every candidate voice.

usage: refbench [ref_id ...]      render (tts env); no ids: the configured reference + the 8 best by sim
       refbench --report          table after `qc --jobs work/refbench/jobs.json --out work/refbench/qc.json`
Pick the reference with low WER and few bad takes (a leaking reference shows as many takes with WER > 0.1:
English at the start), then the highest MOS; put its id into voice.reference. Local backend only.
"""
import collections
import json
import os
import random
import re
import statistics as st
import sys

from .config import load


def report(cfg):
    q = json.load(open(cfg.w("refbench", "qc.json"), encoding="utf-8"))
    by = collections.defaultdict(list)
    for k, r in q.items():
        if "error" not in r:
            by[k.split("|")[0]].append(r)
    rows = []
    for v, rs in by.items():
        so = sum(r["stress_ok"] for r in rs) / max(1, sum(r["stress_n"] for r in rs))
        rows.append((st.mean(r.get("mos", 0) for r in rs), v, len(rs), st.mean(r["wer"] for r in rs),
                     sum(r["wer"] > .1 for r in rs), so, st.mean(r["sim"] for r in rs)))
    print(f"{'MOS':>5} {'WER':>6} {'bad':>4} {'stress':>7} {'SIM':>6}  reference")
    for mos, v, n, wer, bad, so, sim in sorted(rows, reverse=True):
        print(f"{mos:5.2f} {wer:6.3f} {bad:4d} {so:7.3f} {sim:6.3f}  {v}")


def main():
    cfg = load()
    if "--report" in sys.argv:
        return report(cfg)
    from .accent import acute_marks, plus_marks
    from .lexicon import post, pre
    from .render import reference
    from .tts.local import LocalTTS
    refs = json.load(open(cfg.w("refs.json"), encoding="utf-8"))
    ids = [a for a in sys.argv[1:] if not a.startswith("-")] or \
        list(dict.fromkeys([cfg["voice"]["reference"]] + [c["id"] for c in refs[:8]]))
    notes = json.load(open(cfg.w("notes.json"), encoding="utf-8"))
    tp = cfg.w("refbench", "texts.json")
    if tp.exists():
        sel = json.load(open(tp, encoding="utf-8"))
    else:
        random.seed(7)
        paras = [(n["k"], p.strip()) for n in notes for p in re.split(r"\n\s*\n", pre(n["text"])) if 150 <= len(p.strip()) <= 380]
        random.shuffle(paras)
        seen, sel = set(), []
        for k, p in paras:
            if k not in seen and len(sel) < 16:
                seen.add(k); sel.append(dict(id=k, text=p, plus=post(plus_marks(p))))
        json.dump(sel, open(tp, "w", encoding="utf-8"), ensure_ascii=False, indent=1)
    tts, jobs = None, []
    for rid in ids:
        wav, text = reference(cfg, rid)
        tts = tts or LocalTTS(cfg, wav, text)
        tts.set_voice(wav, text)
        for seed in (1000, 1001):
            items = [(acute_marks(s["plus"]), str(cfg.w("refbench", "takes", f"{rid}~{s['id']}~{seed}.wav")))
                     for s in sel]
            todo = [it for it in items if not os.path.exists(it[1])]
            if todo:
                tts.speak(todo, seed)
            jobs += [dict(id=f"{rid}|{s['id']}|{seed}", text=s["text"], plus=s["plus"], wav=it[1]) for s, it in zip(sel, items)]
    json.dump(jobs, open(cfg.w("refbench", "jobs.json"), "w", encoding="utf-8"), ensure_ascii=False)
    print("REFBENCH_RENDERED", len(jobs), "takes; next: python run.py qc --jobs", cfg.w("refbench", "jobs.json"),
          "--out", cfg.w("refbench", "qc.json"))


if __name__ == "__main__":
    main()