#!/usr/bin/env python3
"""mesm-tts text-to-speech with onnxruntime: ten English voices, 2.3-6.5 MB each, CPU only, no GPU or PyTorch.

    pip install -r requirements.txt
    python tts.py "Hello there. How are you?" --voice hadley --out hello.wav
    python tts.py --list

    from tts import MesmTTS
    tts = MesmTTS("miles")                 # downloads the voice's ONNX file on first use if it is not next to this script
    audio = tts.synth("Good evening.")       # float32 numpy array, 24 kHz mono
    tts.save("Good evening.", "evening.wav")

Pipeline (the same one the mesmer.tools browser engine runs, see web/):
  text -> sentences -> misaki G2P (US English, eSpeak NG fallback for words outside its lexicon; exactly the front end the
  models were trained on) -> phoneme ids (tokenizer.json, '$' = 0 at both ends, at most 512 per call) -> ONNX model
  (input_ids int64 [1, T], speed float32 [1] -> waveform float32 [1, samples] @ 24 kHz) -> natural sentence joins (Phonon
  voices) -> per-voice level (gain_db) with a soft knee above 0.9.
The models draw their own noise, so two runs of the same text differ slightly in the noisy parts.
"""
import argparse, json, os, re, sys, time, wave

import numpy as np
import onnxruntime as ort

REPO_ID = "mesmertech/mesm-tts"
HERE = os.path.dirname(os.path.abspath(__file__))
SR = 24000
MAX_IDS = 512


def _file(path, repo_id=REPO_ID):
    """A repo file: next to this script if present (a clone or snapshot of the repo), else from the Hugging Face Hub."""
    local = os.path.join(HERE, path)
    if os.path.exists(local):
        return local
    from huggingface_hub import hf_hub_download
    return hf_hub_download(repo_id, path)


# ---- text: sentences, as web/v2/paradee.mjs splitSentences ------------------------------------------------------------
# Split after . ! ? … (and any closing quotes / brackets) followed by a capital, digit, quote or bracket (the unit the models
# were trained on), except after titles, common abbreviations and initials ("Dr. Smith", "J. R. R. Tolkien", "U.S. Army").
_ABBREV = re.compile(r'(?:^|[\s"“(])(?:Mr|Mrs|Ms|Dr|St|Jr|Sr|Prof|Rev|Gen|Gov|Sen|Rep|Capt|Lt|Col|Sgt|Adm|Mt|Ft|No|Nos|vs|etc|'
                     r'Inc|Ltd|Co|Corp|Bros|Jan|Feb|Mar|Apr|Jun|Jul|Aug|Sep|Sept|Oct|Nov|Dec|[A-Z]|(?:[A-Za-z]\.)+[A-Za-z])\.$')
_END = re.compile(r'[.!?…]["”’\')\]]*$')
_BREAK = re.compile(r'\s+(?=[A-Z"“(\d])')


def split_sentences(text):
    t = re.sub(r"\s+", " ", text).strip()
    pieces, start = [], 0
    for m in _BREAK.finditer(t):
        if _END.search(t, 0, m.start()):
            pieces.append(t[start:m.start()]); start = m.end()
    pieces.append(t[start:])
    out = []
    for p in pieces:
        if out and _ABBREV.search(out[-1]):
            out[-1] += " " + p
        else:
            out.append(p)
    return [s.strip() for s in out if s.strip()]


def split_phonemes(ps, limit=MAX_IDS - 12):
    """A phoneme string longer than the model's 512-id window is cut at the last clause punctuation, else the last word break."""
    out = []
    while len(ps) > limit:
        head = ps[:limit]
        k = max(head.rfind(p) for p in (", ", "; ", ": ", "— ", "… "))
        k = k + 1 if k > limit / 3 else head.rfind(" ")
        if k <= 0:
            k = len(head)
        out.append(ps[:k].strip()); ps = ps[k:].strip()
    if ps:
        out.append(ps)
    return out


class Tokenizer:
    """tokenizer.json is a character map: drop characters outside its set, one id per character, '$' (0) at both ends."""

    def __init__(self, path):
        j = json.load(open(path, encoding="utf-8"))
        self.vocab = j["model"]["vocab"]
        self.strip = re.compile(j["normalizer"]["pattern"]["Regex"])

    def __call__(self, ps):
        s = self.strip.sub("", ps)
        return np.array(([0] + [self.vocab[c] for c in s if c in self.vocab] + [0])[:MAX_IDS], np.int64)


def make_g2p():
    """misaki's English G2P exactly as Kokoro's KPipeline(lang_code="a") builds it, which is what produced the training data."""
    from misaki import en
    try:
        from misaki import espeak
        fallback = espeak.EspeakFallback(british=False)
    except Exception as e:  # eSpeak NG missing: words outside misaki's lexicon (mostly rare names) are skipped
        print(f"note: no eSpeak NG fallback ({e}); words outside misaki's lexicon will be skipped", file=sys.stderr)
        fallback = None
    g2p = en.G2P(trf=False, british=False, fallback=fallback, unk="")
    return lambda sentence: g2p(sentence)[0]


# ---- audio: natural sentence joins (web/pv/flow.mjs = the reference pv_flow.join), level ---------------------------------
_FR = 240  # 10 ms frames
_FADE = np.linspace(0, 1, _FR, dtype=np.float32)


def _edges(y):
    """[first speech sample, end of the last speech frame]; speech = 10 ms frames within 35 dB of the loudest one."""
    n = len(y) // _FR
    db = 10 * np.log10((y[:n * _FR].reshape(n, _FR) ** 2).mean(1) + np.float32(1e-12))
    sp = np.where(db > db.max() - 35)[0]
    return sp[0] * _FR, (sp[-1] + 1) * _FR


def join(parts, gap_scale=1.0, lead_keep_s=0.02):
    """Sentences 2..n keep lead_keep_s of leading silence; sentences 1..n-1 keep gap_scale of their trailing silence."""
    if not parts:
        return np.zeros(0, np.float32)
    if gap_scale == 1.0:
        return np.concatenate(parts)
    lead, out = round(lead_keep_s * SR), []
    for i, y in enumerate(parts):
        y = y.copy()
        a, _ = _edges(y)
        if i > 0 and a > lead:
            y = y[a - lead:]; y[:_FR] *= _FADE
        _, b = _edges(y)
        tail = len(y) - b; keep = int(tail * gap_scale)
        if i < len(parts) - 1 and keep < tail:
            y = y[:b + keep]; y[-_FR:] *= _FADE[::-1]
        out.append(y)
    return np.concatenate(out)


def level(y, gain_db):
    """The voice's gain, then a soft knee above 0.9 (identity below; 0.9 + 0.1 tanh((|x| - 0.9) / 0.1) above): never clips."""
    y = y * np.float32(10 ** (gain_db / 20))
    a = np.abs(y)
    return np.where(a <= 0.9, y, np.sign(y) * (0.9 + 0.1 * np.tanh((a - 0.9) / 0.1))).astype(np.float32)


def write_wav(path, y, sr=SR):
    pcm = (np.clip(y, -1, 1) * 32767).round().astype("<i2")
    with wave.open(path, "wb") as w:
        w.setnchannels(1); w.setsampwidth(2); w.setframerate(sr); w.writeframes(pcm.tobytes())


# ---- the model ------------------------------------------------------------------------------------------------------------
def load_voices(repo_id=REPO_ID):
    return json.load(open(_file("voices.json", repo_id), encoding="utf-8"))["voices"]


class MesmTTS:
    """One voice. voice: a key of voices.json ("heart", "wren", "everly", ...). model: Heart's file ("v2f_lite", "v2f",
    "v2g"; default v2f_lite). fp32_compute: the *_noint8 twin (Heart only). threads: onnxruntime intra-op threads."""

    def __init__(self, voice="heart", model=None, fp32_compute=False, threads=1, repo_id=REPO_ID, g2p=None):
        voices = load_voices(repo_id)
        if voice not in voices:
            raise ValueError(f"unknown voice {voice!r}; one of: {', '.join(voices)}")
        v = self.info = voices[voice]
        if "models" in v:  # Heart: several files of one voice
            m = v["models"][model or v["default_model"]]
            onnx = m["onnx_fp32_compute"] if fp32_compute else m["onnx"]
        else:
            if model:
                raise ValueError(f"{voice} has a single model file")
            onnx = v["onnx"]
        self.gain_db, self.join = v["gain_db"], v.get("join")
        so = ort.SessionOptions()
        so.intra_op_num_threads, so.inter_op_num_threads = threads, 1
        self.session = ort.InferenceSession(_file(onnx, repo_id), so, providers=["CPUExecutionProvider"])
        self.tokenize = Tokenizer(_file("tokenizer.json", repo_id))
        self.g2p = g2p or make_g2p()

    def plan(self, text):
        """text -> [(sentence, phonemes, ids)], long sentences split to fit the model's window."""
        out = []
        for s in split_sentences(text):
            for ps in split_phonemes(self.g2p(s)):
                ids = self.tokenize(ps)
                if len(ids) > 2:
                    out.append((s, ps, ids))
        return out

    def synth_ids(self, ids, speed=1.0):
        """One window of phoneme ids -> raw float32 audio (no join, no gain)."""
        return self.session.run(None, {"input_ids": ids[None], "speed": np.array([speed], np.float32)})[0][0]

    def synth(self, text, speed=1.0):
        """Whole text -> float32 audio @ 24 kHz, joined and levelled as the browser plays it. speed > 1 talks faster."""
        parts = [self.synth_ids(ids, speed) for _, _, ids in self.plan(text)]
        j = self.join or {}
        return level(join(parts, j.get("gap_scale", 1.0), j.get("lead_keep_s", 0.02)), self.gain_db)

    def save(self, text, path, speed=1.0):
        y = self.synth(text, speed)
        write_wav(path, y)
        return y


def main():
    ap = argparse.ArgumentParser(description="mesm-tts text-to-speech (onnxruntime, CPU).")
    ap.add_argument("text", nargs="?", help="text to speak (or --file)")
    ap.add_argument("-v", "--voice", default="heart", help="voice id (see --list); default heart")
    ap.add_argument("-o", "--out", default="out.wav")
    ap.add_argument("-f", "--file", help="read the text from this file")
    ap.add_argument("--model", help="Heart only: v2f_lite (default), v2f or v2g")
    ap.add_argument("--fp32-compute", action="store_true", help="Heart only: the *_noint8 twin (no int8 matrix products)")
    ap.add_argument("--speed", type=float, default=1.0, help="speaking rate (1 = as trained; 1.2 = 20%% faster)")
    ap.add_argument("--threads", type=int, default=1, help="onnxruntime threads (default 1; the models are small)")
    ap.add_argument("--repo", default=REPO_ID, help="Hugging Face repo to fetch files from when they are not local")
    ap.add_argument("--list", action="store_true", help="list the voices and exit")
    a = ap.parse_args()
    if a.list:
        for k, v in load_voices(a.repo).items():
            print(f"{k:8s} {v['name']:7s} {v['gender']:6s} {v['accent'].upper():3s} {v['tag']:9s} from {v['teacher']}")
        return
    text = open(a.file, encoding="utf-8").read() if a.file else a.text
    if not text:
        ap.error("give the text to speak, or --file")
    t0 = time.perf_counter()
    tts = MesmTTS(a.voice, a.model, a.fp32_compute, a.threads, a.repo)
    t1 = time.perf_counter()
    plan = tts.plan(text)
    t2 = time.perf_counter()
    parts = [tts.synth_ids(ids, a.speed) for _, _, ids in plan]
    t3 = time.perf_counter()
    j = tts.join or {}
    y = level(join(parts, j.get("gap_scale", 1.0), j.get("lead_keep_s", 0.02)), tts.gain_db)
    write_wav(a.out, y)
    dur = len(y) / SR
    print(f"{a.out}: {dur:.1f} s of audio ({tts.info['name']}). load {1e3 * (t1 - t0):.0f} ms, phonemes {1e3 * (t2 - t1):.0f} ms, "
          f"synthesis {1e3 * (t3 - t2):.0f} ms = {dur / max(t3 - t2, 1e-9):.0f}x real time on {a.threads} thread(s)")


if __name__ == "__main__":
    main()