"""Step 2: reference clips of the EN narrator for the voice clone, and the narrator's measurements.

Candidates: whole notes of 7-16 s and first paragraphs of 7-15 s cut at the pause nearest to where the
paragraph should end (by its share of characters). Each needs its exact EN text (ICL cloning), so a cut
that misses the pause leaks: the clone then starts Russian takes with the English tail ("we were a
harvest...") — refbench catches that by WER.
Writes work/refs/<id>.wav (id = "<EN sound name>__all|__p1"), work/refs.json (text, length, similarity
to the narrator centroid, F0 spread), work/centroid.npy (mean WavLM x-vector of the narrator, the QC
target for the timbre) and work/en_pauses.json (pauses between paragraphs and phrases in the EN originals).
"""
import json
import re

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

from .config import load


def clean(s):
    return re.sub(r"<[^>]+>", "", s or "")


def main():
    cfg = load()
    notes = json.load(open(cfg.w("notes.json"), encoding="utf-8"))
    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()

    def embed(y24):
        y = librosa.resample(y24, orig_sr=24000, target_sr=16000)[:16000 * 30]
        x = fe(y, 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().numpy()

    cands, embs, para_gaps, sent_gaps = [], [], [], []
    for n in notes:
        y, sr = sf.read(cfg.w("en24", f"{n['k']}.wav"), dtype="float32")
        embs.append(embed(y[: sr * 20]))
        iv = librosa.effects.split(y, top_db=40, frame_length=1024, hop_length=240)
        if len(iv) < 2:
            continue
        gaps = [((iv[i][1] + iv[i + 1][0]) / 2 / sr, (iv[i + 1][0] - iv[i][1]) / sr) for i in range(len(iv) - 1)]
        a0, a1 = iv[0][0] / sr, iv[-1][1] / sr
        paras = [clean(p).strip() for p in re.split(r"\n\s*\n", n["en"] or "") if clean(p).strip()]
        full = " ".join(paras)
        if not full:
            continue
        if 7 <= a1 - a0 <= 16:
            cands.append(dict(id=f"{n['name']}__all", k=n["k"], a=a0, b=a1, text=re.sub(r"\s+", " ", full)))
        if len(paras) >= 2:
            pos = a0 + (a1 - a0) * len(paras[0]) / len(full)
            g = min(gaps, key=lambda g: abs(g[0] - pos))
            if abs(g[0] - pos) < 0.12 * (a1 - a0) and g[1] > 0.35:
                para_gaps.append(g[1])
                if 7 <= g[0] - a0 <= 15:
                    cands.append(dict(id=f"{n['name']}__p1", k=n["k"], a=a0, b=g[0] - g[1] / 2 + 0.1,
                                      text=re.sub(r"\s+", " ", paras[0])))
        sent_gaps += [g[1] for g in gaps if g[1] > 0.15]
    cen = np.mean(embs, 0)
    cen /= np.linalg.norm(cen)
    np.save(cfg.w("centroid.npy"), cen)
    out = []
    for c in cands:
        y, sr = sf.read(cfg.w("en24", f"{c['k']}.wav"), dtype="float32")
        seg = y[max(0, int((c["a"] - 0.15) * sr)): int((c["b"] + 0.1) * sr)]
        sf.write(cfg.w("refs", f"{c['id']}.wav"), seg, sr)
        f0, vf, _ = librosa.pyin(librosa.resample(seg, orig_sr=24000, target_sr=16000), fmin=60, fmax=400, sr=16000)
        f0 = f0[vf]
        c.update(sec=round(len(seg) / sr, 2), sim=round(float(embed(seg) @ cen), 3),
                 f0=round(float(np.std(12 * np.log2(f0 / np.median(f0)))) if len(f0) > 10 else 0.0, 2))
        out.append(c)
    out.sort(key=lambda c: -c["sim"])
    json.dump(out, open(cfg.w("refs.json"), "w", encoding="utf-8"), ensure_ascii=False, indent=1)
    pg, sg = np.array(para_gaps), np.array(sent_gaps)
    pauses = dict(paragraph_median=float(np.median(pg)), gap_median=float(np.median(sg)),
                  gap_p75=float(np.percentile(sg, 75)))
    json.dump(pauses, open(cfg.w("en_pauses.json"), "w"), indent=1)
    print("candidates", len(out), "pauses", pauses)
    want = cfg["voice"]["reference"]
    print("configured reference", want, "found" if any(c["id"] == want for c in out) else "NOT FOUND")


if __name__ == "__main__":
    main()