"""End-to-end TTS inference using the 4 CoreML components.

Pipeline (mirrors supertonic.core.Supertonic):
    text -> tokenize
        -> duration_predictor -> frame count
        -> text_encoder       -> text embedding
        -> sample noisy latent ~ N(0, I)
        -> vector_estimator x 8 (flow-matching ODE step, runs on ANE)
        -> vocoder -> 44.1 kHz waveform

All four mlpackages are static-shape buckets at T=L=320. The driver pads
inputs to that bucket and trims outputs.

The supertonic-3 model truncates long prompts at its content limit
(~13.7s natural; CoreML's bucket-leak extends this to ~16.7s but still
short for long inputs). The `--auto-pad` mode does a two-pass synthesis
(once unpadded to find the natural endpoint, once with a long filler
sentence appended that gives the model more frames to render the full
original prompt), then trims at the silence gap between original and
appended content. Recommended for prompts longer than ~5s.

Usage:
    python inference.py --text "Hello, world." --voice F1 --lang en
    python inference.py --text "<longer prompt>" --voice F5 --lang en --auto-pad
"""
from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

import coremltools as ct
import numpy as np
import soundfile as sf

HERE = Path(__file__).parent
T_BUCKET = 320
L_BUCKET = 320
SAMPLE_RATE = 44_100
LATENT_DIM = 24
CHUNK_COMPRESS_FACTOR = 6
BASE_CHUNK_SIZE = 512
DEFAULT_TOTAL_STEPS = 8
DEFAULT_SPEED = 1.05
DEFAULT_AUTO_PAD = " And with that, the gentle silence wrapped itself around the room."


def _pad(arr: np.ndarray, axis: int, target: int) -> np.ndarray:
    if arr.shape[axis] >= target:
        return arr
    pad = [(0, 0)] * arr.ndim
    pad[axis] = (0, target - arr.shape[axis])
    return np.pad(arr, pad)


def _load_voice(name: str) -> tuple[np.ndarray, np.ndarray]:
    j = json.loads((HERE / "voice_styles" / f"{name}.json").read_text())
    def r(part): return np.array(part["data"], dtype=np.float32).reshape(*part["dims"])
    return r(j["style_ttl"]), r(j["style_dp"])


def _load_tokenizer(indexer_path: Path):
    """Reuse the official supertonic UnicodeProcessor (handles the 31
    languages, abbreviation expansion, punctuation rules, etc.).
    Install with: pip install supertonic
    """
    try:
        from supertonic.core import UnicodeProcessor
    except ImportError as e:
        raise RuntimeError(
            "supertonic package is required for tokenization. "
            "Install with: pip install supertonic"
        ) from e
    return UnicodeProcessor(str(indexer_path))


def _last_loud_window(audio: np.ndarray, thresh: float = 0.025, win_s: float = 0.05) -> int:
    win = int(win_s * SAMPLE_RATE)
    n = len(audio) // win
    rms = np.sqrt(np.mean(audio[: n * win].reshape(n, win) ** 2, axis=1))
    loud = np.where(rms > thresh)[0]
    return int(loud[-1]) if len(loud) else 0


def trim_padded(unpad: np.ndarray, padded: np.ndarray) -> np.ndarray:
    """Trim padded synthesis at the longest clean silence between original
    prompt and appended suffix. Tail-pad with 0.5 s of true silence."""
    win = int(0.05 * SAMPLE_RATE)
    n = len(padded) // win
    rms = np.sqrt(np.mean(padded[: n * win].reshape(n, win) ** 2, axis=1))
    floor = _last_loud_window(unpad)
    ceil_ = _last_loud_window(padded) + 1
    candidates = []
    j = floor
    while j < ceil_ - 1:
        if rms[j] < 0.025 and rms[j + 1] < 0.025:
            start = j; total = 0.0; cnt = 0
            while j < ceil_ and rms[j] < 0.025:
                total += float(rms[j]); cnt += 1; j += 1
            candidates.append((start, cnt, total / max(cnt, 1)))
        else:
            j += 1
    if not candidates:
        return padded[: ceil_ * win]
    start_win, length, avg = max(candidates, key=lambda c: (c[1], -c[0]))
    end_samples = start_win * win
    out = padded[:end_samples].copy()
    fade = min(int(0.06 * SAMPLE_RATE), len(out))
    out[-fade:] *= np.linspace(1.0, 0.0, fade, dtype=np.float32)
    return np.concatenate([out, np.zeros(int(0.5 * SAMPLE_RATE), dtype=np.float32)])


class Supertonic3CoreML:
    def __init__(self, quant: str = "fp16"):
        d = HERE / quant
        self.dp = ct.models.MLModel(str(d / "duration_predictor.mlpackage"))
        self.te = ct.models.MLModel(str(d / "text_encoder.mlpackage"))
        self.ve = ct.models.MLModel(str(d / "vector_estimator.mlpackage"))
        self.voc = ct.models.MLModel(str(d / "vocoder.mlpackage"))
        self.tok = _load_tokenizer(HERE / "unicode_indexer.json")

    def _synth(self, text: str, voice: str, lang: str, seed: int,
               total_steps: int, speed: float, full_bucket: bool) -> np.ndarray:
        text_ids, text_mask = self.tok([text], lang)
        text_ids = text_ids.astype(np.int64); text_mask = text_mask.astype(np.float32)
        style_ttl, style_dp = _load_voice(voice)
        text_ids_p = _pad(text_ids.astype(np.int32), 1, T_BUCKET)
        text_mask_p = _pad(text_mask, 2, T_BUCKET)
        dur = float(self.dp.predict({"text_ids": text_ids_p, "style_dp": style_dp,
                                     "text_mask": text_mask_p})["duration"][0]) / speed
        text_emb = self.te.predict({"text_ids": text_ids_p, "style_ttl": style_ttl,
                                    "text_mask": text_mask_p})["text_emb"]
        L_real = max(1, min(L_BUCKET, (int(dur * SAMPLE_RATE) + BASE_CHUNK_SIZE * CHUNK_COMPRESS_FACTOR - 1) // (BASE_CHUNK_SIZE * CHUNK_COMPRESS_FACTOR)))
        np.random.seed(seed)
        xt = (np.random.randn(1, LATENT_DIM * CHUNK_COMPRESS_FACTOR, L_real)).astype(np.float32)
        latent_mask = np.ones((1, 1, L_real), dtype=np.float32)
        xt = xt * latent_mask
        xt = _pad(xt, 2, L_BUCKET)
        latent_mask = _pad(latent_mask, 2, L_BUCKET)
        total_step_arr = np.array([float(total_steps)], dtype=np.float32)
        for step in range(total_steps):
            xt = self.ve.predict({
                "noisy_latent": xt, "text_emb": text_emb, "style_ttl": style_ttl,
                "text_mask": text_mask_p, "latent_mask": latent_mask,
                "current_step": np.array([float(step)], dtype=np.float32),
                "total_step": total_step_arr,
            })["denoised_latent"]
        wav = self.voc.predict({"latent": xt})["wav_tts"][0]
        if full_bucket:
            return wav
        return wav[: L_real * CHUNK_COMPRESS_FACTOR * BASE_CHUNK_SIZE]

    def synthesize(self, text: str, voice: str = "F1", lang: str = "en", seed: int = 0,
                   total_steps: int = DEFAULT_TOTAL_STEPS, speed: float = DEFAULT_SPEED,
                   auto_pad: str | None = DEFAULT_AUTO_PAD) -> np.ndarray:
        """Synthesize speech. With ``auto_pad`` set, runs the 2-pass auto-pad
        flow for full content rendering on longer prompts."""
        if auto_pad is None:
            return self._synth(text, voice, lang, seed, total_steps, speed, full_bucket=False)
        unpad_audio = self._synth(text, voice, lang, seed, total_steps, speed, full_bucket=True)
        pad_audio = self._synth(text + auto_pad, voice, lang, seed, total_steps, speed, full_bucket=True)
        return trim_padded(unpad_audio, pad_audio)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--text", required=True, help="Text to synthesize")
    ap.add_argument("--voice", default="F1", choices=[f"F{i}" for i in range(1, 6)] + [f"M{i}" for i in range(1, 6)])
    ap.add_argument("--lang", default="en")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--total-steps", type=int, default=DEFAULT_TOTAL_STEPS)
    ap.add_argument("--auto-pad", nargs="?", const=DEFAULT_AUTO_PAD, default=None,
                    help="2-pass synthesis with filler suffix + auto-trim (recommended).")
    ap.add_argument("--quant", default="fp16", choices=["fp16"])
    ap.add_argument("--out", default="out.wav")
    args = ap.parse_args()

    t0 = time.time()
    tts = Supertonic3CoreML(quant=args.quant)
    print(f"Loaded models in {time.time() - t0:.2f}s")

    t0 = time.time()
    audio = tts.synthesize(args.text, voice=args.voice, lang=args.lang, seed=args.seed,
                           total_steps=args.total_steps, auto_pad=args.auto_pad)
    dur = len(audio) / SAMPLE_RATE
    sf.write(args.out, audio, SAMPLE_RATE)
    print(f"Synthesized {dur:.2f}s of audio in {time.time() - t0:.2f}s -> {args.out}")
    return 0


if __name__ == "__main__":
    sys.exit(main())