"""Step 5: score every take automatically (no listening).

usage: qc [--shard i/n] [--no-mos] [--f0] [--jobs jobs.json --out out.json]
Default jobs: every work/takes/<take>_s<seed>.wav of the current segments; results are cached in
work/qc/<i>of<n>.json by take name, so re-runs and lexicon changes only score new takes. Run n shards
in parallel to use the GPU fully (each process needs ~8-13 GB RAM and ~5 GB VRAM).

Per take:
- wer, hyp — GigaAM-v3 (RNN-T) on <=20 s pieces cut at silences, against the text;
- stress_ok / stress_n — meter2 (shipped with the ru-stress model) agrees with the RUAccent marks;
- sim — WavLM-SV x-vector cosine to the EN narrator centroid (timbre);
- mos — UTMOSv2 naturalness; cps — pace, chars/s; maxgap — longest inner pause; clip — clipped samples;
- f0 — pitch spread in semitones (only with --f0: pyin is slow and runs on the CPU).
"""
import glob
import json
import os
import re
import sys
import tempfile

import jiwer
import librosa
import numpy as np
import soundfile as sf
import torch

from .config import load

VOW = "аеёиоуыэюя"


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


def arg(name, default=None):
    return sys.argv[sys.argv.index(name) + 1] if name in sys.argv else default


def main():
    cfg = load()
    sh = tuple(int(x) for x in arg("--shard", "0/1").split("/"))
    use_mos, use_f0 = "--no-mos" not in sys.argv and cfg.get("qc", "mos", True), "--f0" in sys.argv
    if arg("--jobs"):
        jobs, out_path = json.load(open(arg("--jobs"), encoding="utf-8")), arg("--out")
    else:
        segs = {s["take"]: s for s in json.load(open(cfg.w("segs.json"), encoding="utf-8"))}
        jobs = []
        for f in sorted(glob.glob(str(cfg.w("takes", "x").parent / "*_s*.wav"))):
            stem = os.path.basename(f)[:-4]
            take = stem.rsplit("_s", 1)[0]
            if take in segs:
                jobs.append(dict(id=stem, text=segs[take]["text"], plus=segs[take]["plus"], wav=f))
        out_path = str(cfg.w("qc", f"{sh[0]}of{sh[1]}.json"))
    done = {}
    for f in glob.glob(str(cfg.w("qc", "x").parent / "*.json")) + ([out_path] if os.path.exists(out_path) else []):
        try:
            done.update(json.load(open(f, encoding="utf-8")))
        except Exception:
            pass
    jobs = [j for j in jobs[sh[0]::sh[1]] if j["id"] not in done]
    print("takes to score", len(jobs), flush=True)
    if not jobs:
        return
    res = json.load(open(out_path, encoding="utf-8")) if os.path.exists(out_path) else {}

    from huggingface_hub import snapshot_download
    model = cfg.get("tts", "model", "siriusfreak/qwen3-tts-12hz-1.7b-ru-stress-cf")  # meter2 ships with it
    meter = os.path.join(model if os.path.isdir(model) else snapshot_download(model, allow_patterns=["meter/*"]),
                         "meter")
    os.environ.setdefault("METER2_HEAD", os.path.join(meter, "head.pt"))
    os.environ.setdefault("METER2_DEVICE", "cuda")
    sys.path.insert(0, meter)
    from meter2 import judge, vowel_pos, words_in

    def marks(plus, text):
        toks = [m.group(0) for m in re.finditer(r"[А-Яа-яЁё+]+(?:-[А-Яа-яЁё+]+)*", plus)]
        if len(toks) != len(words_in(text)):
            return {}
        return {wi: sum(1 for ch in w[:w.index("+")] if ch.lower() in VOW) for wi, w in enumerate(toks)
                if "+" in w and len(vowel_pos(w.replace("+", ""))) > 1}

    import gigaam
    asr = gigaam.load_model("v3_rnnt")
    tmp = os.path.join(tempfile.gettempdir(), f"ensvoice_qc_{os.getpid()}.wav")  # one per process

    def transcribe(y, sr):
        iv = librosa.effects.split(y, top_db=35, frame_length=1024, hop_length=256)
        pieces, start, cur = [], None, None
        for a, b in iv:
            if cur is None:
                start, cur = a, b
            elif (b - start) / sr <= 20:
                cur = b
            else:
                pieces.append((start, cur)); start, cur = a, b
        if cur is not None:
            pieces.append((start, cur))
        out = []
        for a, b in pieces:
            sf.write(tmp, y[max(0, a - 2400):b + 2400], sr)
            r = asr.transcribe(tmp)
            out.append(r if isinstance(r, str) else getattr(r, "text", str(r)))
        return " ".join(out), iv

    from transformers import AutoFeatureExtractor, WavLMForXVector
    sv_id = cfg.get("qc", "wavlm", "microsoft/wavlm-base-plus-sv")
    fe = AutoFeatureExtractor.from_pretrained(sv_id)
    sv = WavLMForXVector.from_pretrained(sv_id).cuda().eval()
    cen = torch.tensor(np.load(cfg.w("centroid.npy")))

    def embed(y16):
        if len(y16) < 32000:
            y16 = np.pad(y16, (0, 32000 - len(y16)))
        x = fe(y16[:16000 * 30], sampling_rate=16000, return_tensors="pt")
        with torch.no_grad():
            e = sv(**{k: v.cuda() for k, v in x.items()}).embeddings
        return torch.nn.functional.normalize(e, dim=-1)[0].cpu()

    if use_mos:
        import utmosv2
        mos_model = utmosv2.create_model(pretrained=True)

    for i, j in enumerate(jobs):
        try:
            y, sr = sf.read(j["wav"], dtype="float32")
            if y.ndim > 1:
                y = y.mean(1)
            dur = len(y) / sr
            hyp, iv = transcribe(y, sr)
            wer = jiwer.wer(norm(j["text"]), norm(hyp)) if norm(hyp) else 1.0
            J = judge(y, sr, j["text"])
            items = [(wi, v) for wi, v in marks(j["plus"], j["text"]).items() if wi < len(J) and J[wi]["best"] is not None]
            y16 = librosa.resample(y, orig_sr=sr, target_sr=16000)
            gaps = [(iv[k + 1][0] - iv[k][1]) / sr for k in range(len(iv) - 1)]
            r = dict(dur=round(dur, 2), cps=round(len(j["text"]) / dur, 1), wer=round(wer, 3), hyp=hyp,
                     stress_ok=sum(1 for wi, v in items if J[wi]["best"] == v), stress_n=len(items),
                     bad=[J[wi]["word"] for wi, v in items if J[wi]["best"] != v][:10],
                     sim=round(float(embed(y16) @ cen), 3), maxgap=round(max(gaps), 2) if gaps else 0,
                     clip=int(np.sum(np.abs(y) > 0.99)))
            if use_f0:
                f0, vf, _ = librosa.pyin(y16, fmin=60, fmax=400, sr=16000, frame_length=1024)
                f0 = f0[vf]
                r["f0"] = round(float(np.std(12 * np.log2(f0 / np.median(f0)))) if len(f0) > 10 else 0.0, 2)
            if use_mos:
                r["mos"] = round(float(mos_model.predict(input_path=j["wav"])), 3)
        except Exception as ex:  # a broken take simply loses
            r = dict(error=str(ex)[:200], wer=1.0, cps=0, stress_ok=0, stress_n=1, sim=0, mos=0)
        res[j["id"]] = r
        print(f"{i + 1}/{len(jobs)}", j["id"], r.get("wer"), f'{r["stress_ok"]}/{r["stress_n"]}', r.get("sim"),
              r.get("mos"), flush=True)
        if i % 20 == 0:
            json.dump(res, open(out_path, "w", encoding="utf-8"), ensure_ascii=False)
    json.dump(res, open(out_path, "w", encoding="utf-8"), ensure_ascii=False)
    print("QC_DONE", len(res))


if __name__ == "__main__":
    main()