"""Phase 2 shared tools: romanisation for MMS forced alignment, chunked aligner, pause normalisation,
pause / F0 / Yoruba-tone metrics.  Used by prep2.py (data) and eval2.py (metrics)."""
import re, unicodedata
import numpy as np

# ------------------------------------------------------------------ text
_MAP = {"ọ": "o", "ẹ": "e", "ṣ": "s", "ị": "i", "ụ": "u", "ṅ": "n", "ń": "n", "ǹ": "n", "ḿ": "m", "ɓ": "b", "ɗ": "d",
        "ƙ": "k", "ƴ": "y", "ŋ": "n", "ɛ": "e", "ɔ": "o", "ß": "s", "æ": "ae", "œ": "oe", "ø": "o", "đ": "d", "ł": "l"}
TONE_MARK = {"\u0301": 2, "\u0300": 0, "\u0304": 1}          # acute=H, grave=L, macron=M
SENT_END = set(".?!")
CLAUSE = set(";:")
COMMA = set(",")
_VOW = set("aeiou")


def roman_char(c):
    """one NFC char -> (romanised latin string, tone or None)."""
    lc = c.lower()
    if lc in _MAP:
        return _MAP[lc], None
    d = unicodedata.normalize("NFD", lc)
    base = "".join(x for x in d if unicodedata.category(x) != "Mn")
    tone = None
    for x in d:
        if x in TONE_MARK:
            tone = TONE_MARK[x]
    base = "".join(_MAP.get(b, b) for b in base)
    base = "".join(b for b in base if ("a" <= b <= "z") or b == "'")
    return base, tone


def words_for_align(text):
    """text -> list of dict(word, rom, punct_after, chars=[(rom_char, orig_index, tone)]).
    Words that romanise to nothing (digits, symbols) are aligned with the '*' star token."""
    out = []
    for m in re.finditer(r"\S+", text):
        w = m.group(0)
        chars = []
        for k, c in enumerate(w):
            r, t = roman_char(c)
            for j, rc in enumerate(r):
                chars.append((rc, m.start() + k, t if j == len(r) - 1 else None))
        trail = re.sub(r"[\w'’]+", "", w[-3:]) if w else ""
        p = None
        if any(c in SENT_END for c in trail):
            p = "sent"
        elif any(c in CLAUSE for c in trail):
            p = "clause"
        elif any(c in COMMA for c in trail):
            p = "comma"
        rom = "".join(c[0] for c in chars)
        out.append(dict(word=w, rom=rom if rom else "*", punct=p, chars=chars))
    return out


# ------------------------------------------------------------------ aligner
class Aligner:
    """torchaudio MMS_FA (CC-BY-NC weights, used only as a data tool/judge) with chunked emissions for long audio."""

    def __init__(self, device="cuda"):
        import torch, torchaudio
        self.torch = torch
        self.F = torchaudio.functional
        b = torchaudio.pipelines.MMS_FA
        self.model = b.get_model().to(device).eval().half()
        self.dict = b.get_dict()
        self.dev = device
        self.rate = 320                         # samples per emission frame at 16 kHz (20 ms)

    def emission(self, y16, chunk_s=30, ctx_s=2):
        torch = self.torch
        sr = 16000
        x = torch.from_numpy(y16).float().to(self.dev)
        if len(y16) <= (chunk_s + 2 * ctx_s) * sr:
            with torch.inference_mode():
                e, _ = self.model(x[None].half())
            return e[0].float().log_softmax(-1)
        outs = []
        step, ctx = chunk_s * sr, ctx_s * sr
        nfr_total = len(y16) // self.rate
        for s in range(0, len(y16), step):
            a, b = max(0, s - ctx), min(len(y16), s + step + ctx)
            with torch.inference_mode():
                e, _ = self.model(x[a:b][None].half())
            e = e[0].float().log_softmax(-1)
            f0 = (s - a) // self.rate
            nfr = min(step, len(y16) - s) // self.rate
            outs.append(e[f0:f0 + nfr])
        e = torch.cat(outs, 0)
        return e[:max(nfr_total, 1)]

    def align(self, y16, words):
        """-> per word (start_s, end_s, score) and per char spans [(start_s, end_s)], or None on failure."""
        torch = self.torch
        em = self.emission(y16)
        toks, wid = [], []
        for i, w in enumerate(words):
            for c in w["rom"]:
                if c in self.dict:
                    toks.append(self.dict[c]); wid.append(i)
        if not toks or len(toks) * 1 > em.shape[0]:
            return None
        t = torch.tensor([toks], dtype=torch.int32, device=self.dev)
        try:
            ali, sc = self.F.forced_align(em[None], t, blank=0)
        except Exception:
            return None
        ali, sc = ali[0].cpu().numpy(), sc[0].exp().cpu().numpy()
        # token spans (merge repeats; blank=0)
        spans, i, k = [], 0, 0
        T = len(ali)
        while i < T:
            if ali[i] == 0:
                i += 1; continue
            j = i
            while j + 1 < T and ali[j + 1] == ali[i] and not (j + 1 < T and False):
                j += 1
            spans.append((i, j + 1, float(sc[i:j + 1].mean())))
            i = j + 1
        # repeated identical consecutive tokens (e.g. "aa") collapse in the path above only if no blank between;
        # forced_align always inserts a blank between identical tokens, so len(spans)==len(toks)
        if len(spans) != len(toks):
            return None
        r = self.rate / 16000
        W = [[None, None, []] for _ in words]
        chars = [[] for _ in words]
        for (a, b, s), w in zip(spans, wid):
            if W[w][0] is None:
                W[w][0] = a * r
            W[w][1] = b * r; W[w][2].append(s); chars[w].append((a * r, b * r))
        res = []
        for i, (a, b, s) in enumerate(W):
            if a is None:
                return None
            res.append(dict(start=a, end=b, score=float(np.mean(s)), chars=chars[i]))
        return res


# ------------------------------------------------------------------ energy / pauses
def frame_db(y, sr, hop_s=0.01):
    hop = int(sr * hop_s)
    n = max(1, len(y) // hop)
    fr = y[: n * hop].reshape(n, hop)
    e = 10 * np.log10((fr.astype(np.float64) ** 2).mean(1) + 1e-10)
    return e


def silence_mask(y, sr, rel_db=-40.0, hop_s=0.01):
    e = frame_db(y, sr, hop_s)
    ref = np.percentile(e, 98)
    return e < ref + rel_db


def runs(mask):
    """-> list of (start, end) index runs where mask is True."""
    out, i, n = [], 0, len(mask)
    while i < n:
        if mask[i]:
            j = i
            while j < n and mask[j]:
                j += 1
            out.append((i, j)); i = j
        else:
            i += 1
    return out


PAUSE_TARGET = {           # (min, max) silence seconds after a word, by the punctuation that follows it
    "sent": (0.30, 0.60),
    "clause": (0.20, 0.50),
    "comma": (0.00, 0.40),
    None: (0.00, 0.20),
}
# profile B (Phase 2 v2): keeps natural phrase breaks where there is no punctuation (the real anchor pauses at
# ~14% of unpunctuated word boundaries), still removes hesitations; set P2PAUSE=B
import os as _os
if _os.environ.get("P2PAUSE") == "B":
    PAUSE_TARGET = {"sent": (0.35, 0.70), "clause": (0.25, 0.55), "comma": (0.00, 0.50), None: (0.00, 0.35)}
LEAD_S, TAIL_S = 0.12, 0.18


def normalise_pauses(y, sr, words, ali, seg_a, seg_b, sil=None, lo_s=0.0, hi_s=None):
    """Rebuild one segment [word seg_a .. seg_b] of y: edge silence fixed to LEAD_S/TAIL_S, every inter-word
    silence clamped to PAUSE_TARGET[punct] (only truly silent frames are removed; silence is inserted, as a
    copy of the quietest local frames, when a sentence/clause boundary has too little).  -> (y_out, pauses)"""
    hop = int(0.01 * sr)
    if sil is None:
        sil = silence_mask(y, sr)
    hi_s = len(y) / sr if hi_s is None else hi_s
    xf = int(0.005 * sr)
    pieces, pauses = [], []
    a0 = ali[seg_a]["start"]
    start = int(max(lo_s, a0 - LEAD_S, 0.0) * sr)
    cur = start
    for i in range(seg_a, seg_b):
        e_i, s_n = ali[i]["end"], ali[i + 1]["start"]
        f_a, f_b = int(e_i / 0.01), int(s_n / 0.01) + 1
        # longest silent run inside the gap (with 20 ms guard to the word edges)
        best = None
        for ra, rb in runs(sil[f_a:f_b]):
            if best is None or rb - ra > best[1] - best[0]:
                best = (ra, rb)
        lo, hi = PAUSE_TARGET[words[i]["punct"]]
        if best is None:
            cur_sil = 0.0; mid = int((e_i + s_n) / 2 * sr)
        else:
            ra, rb = (f_a + best[0]) * hop, (f_a + best[1]) * hop
            cur_sil = (rb - ra) / sr; mid = (ra + rb) // 2
        if cur_sil > hi and best is not None:            # shrink: cut the middle of the silent run
            cut = int((cur_sil - hi) * sr)
            c0 = mid - cut // 2
            pieces.append(y[cur:c0]); cur = c0 + cut
            pauses.append(hi)
        elif cur_sil < lo:                               # insert quiet frames at the quietest point of the gap
            add = int((lo - cur_sil) * sr)
            g0, g1 = int(e_i * sr), max(int(s_n * sr), int(e_i * sr) + hop)
            seg = y[g0:g1]
            if len(seg) >= hop:
                k = int(np.argmin(frame_db(seg, sr))) * hop
                ins_at = g0 + k
                q = y[ins_at:ins_at + hop]
            else:
                ins_at, q = mid, np.zeros(hop, dtype=y.dtype)
            filler = np.tile(q, add // hop + 1)[:add] * 0.5
            pieces.append(y[cur:ins_at]); pieces.append(filler); cur = ins_at
            pauses.append(lo)
        else:
            pauses.append(cur_sil)
    end = min(len(y), int(min(hi_s, ali[seg_b]["end"] + TAIL_S) * sr))
    pieces.append(y[cur:end])
    out = np.concatenate([p for p in pieces if len(p)]) if pieces else y[start:end]
    # pad edges with silence if the source had less than LEAD/TAIL available
    lead_have = a0 - start / sr
    if lead_have < LEAD_S - 0.02:
        out = np.concatenate([np.zeros(int((LEAD_S - lead_have) * sr), out.dtype), out])
    tail_have = end / sr - ali[seg_b]["end"]
    if tail_have < TAIL_S - 0.02:
        out = np.concatenate([out, np.zeros(int((TAIL_S - tail_have) * sr), out.dtype)])
    # tiny fades at the edges
    f = min(xf, len(out) // 4)
    if f > 0:
        out[:f] *= np.linspace(0, 1, f); out[-f:] *= np.linspace(1, 0, f)
    return out.astype(np.float32), pauses


# ------------------------------------------------------------------ metrics on audio
def pause_stats(y, sr, min_pause=0.15):
    """VAD (energy, -40 dB vs 98th pct) pauses: leading/trailing silence and internal pauses >= min_pause."""
    sil = silence_mask(y, sr)
    R = runs(sil)
    n = len(sil)
    lead = R[0][1] * 0.01 if R and R[0][0] == 0 else 0.0
    tail = (n - R[-1][0]) * 0.01 if R and R[-1][1] == n else 0.0
    internal = [(b - a) * 0.01 for a, b in R if a > 0 and b < n and (b - a) * 0.01 >= min_pause]
    speech = max(1e-3, n * 0.01 - lead - tail)
    return dict(lead=round(lead, 3), tail=round(tail, 3), pauses=[round(p, 3) for p in internal],
                n_pauses=len(internal), pause_rate=round(len(internal) / speech, 4),
                pause_frac=round(sum(internal) / speech, 4), speech_s=round(speech, 3))


def f0_track(y16):
    import pyworld as pw
    x = y16.astype(np.float64)
    f0, t = pw.dio(x, 16000, f0_floor=60, f0_ceil=500, frame_period=10)
    return pw.stonemask(x, f0, t, 16000)


def f0_stats(y16):
    f0 = f0_track(y16)
    v = f0[f0 > 0]
    if len(v) < 20:
        return dict(voiced=float(len(v)), f0_med=None, st_std=None, st_range=None)
    st = 12 * np.log2(v / np.median(v))
    d = np.diff(st)
    return dict(voiced_frac=round(len(v) / len(f0), 3), f0_med=round(float(np.median(v)), 1),
                st_std=round(float(st.std()), 3), st_range=round(float(np.percentile(st, 95) - np.percentile(st, 5)), 3),
                st_vel=round(float(np.abs(d[np.abs(d) < 3]).mean()), 4))


def tone_pairs(y16, words, ali):
    """Yoruba tone-direction accuracy: for consecutive tone-bearing vowels whose tones differ, is the sign of the
    median-F0 change (semitones) the same as the sign of the tone change?  -> (n_pairs, n_correct)"""
    f0 = f0_track(y16)
    vals = []
    for w, a in zip(words, ali):
        k = 0
        for (rc, oi, tone), (s, e) in zip([c for c in w["chars"] if c[0] in "abcdefghijklmnopqrstuvwxyz'"], a["chars"]):
            if rc in _VOW or (rc in "nm" and tone is not None):
                t = 1 if tone is None else tone
                fr = f0[int(s / 0.01): int(e / 0.01) + 1]
                fr = fr[fr > 0]
                vals.append((t, float(np.median(12 * np.log2(fr))) if len(fr) else None))
    n = c = 0
    for (t1, f1), (t2, f2) in zip(vals, vals[1:]):
        if t1 == t2 or f1 is None or f2 is None:
            continue
        n += 1
        c += int(np.sign(f2 - f1) == np.sign(t2 - t1))
    return n, c


def wasserstein1(a, b):
    """1-D earth mover's distance between two samples."""
    if not len(a) or not len(b):
        return None
    a, b = np.sort(a), np.sort(b)
    q = np.linspace(0, 1, 101)
    return float(np.mean(np.abs(np.quantile(a, q) - np.quantile(b, q))))