"""EMA Lightning — tiny, fast and accurate Turkish text-to-speech on ZeroGPU.

The model authors' own pip package (`ema-lightning`) provides the whole
inference path: the text frontend (normalizer-tr), the 5.6M-parameter DiT
acoustic model and the 3M-parameter decoder, both read from the Apache-2.0
checkpoint at canberkkkkkk/ema-lightning. 48 kHz speech in four ODE steps.
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")  # conv decoder hits transient alloc spikes

import spaces  # MUST come before torch / any CUDA-touching import

import time

import gradio as gr
import numpy as np
import torch
from ema_lightning import EMA

MAX_CHARS = 1200  # keeps the whole request well inside a few seconds of GPU

# Module-scope load: ZeroGPU intercepts .to("cuda") and packs the ~34 MB of
# weights to disk, streaming them into VRAM on the first @spaces.GPU call.
tts = EMA()
SAMPLE_RATE = 48000
# The package normally probes the GPU for its best batch size on first use
# (benching batches up to 128). On ZeroGPU the worker's VRAM is shared with
# the platform and the probe has OOM'd on allocator state; a fixed batch of
# one sentence per pass is plenty for a single-user demo of an 87x-real-time
# model, so pin it and skip the probe entirely.
tts._batch_size = 1

# The package loads checkpoints as .to(device).requires_grad_(False), but the
# ZeroGPU hijack snapshots requires_grad at .to() time, so the decoder's
# weights arrive in the worker flagged as grad-tracking. Playhead decodes in a
# background thread (engine.decode is the one stage without @torch.no_grad()),
# which makes the audio require grad and kills .numpy(). Wrap it.
_engine_decode = tts._engine.decode


@torch.no_grad()
def _decode_nograd(*args, **kwargs):
    return _engine_decode(*args, **kwargs)


tts._engine.decode = _decode_nograd


@spaces.GPU(duration=10)
def synthesize(text: str, speed: float = 1.0, seed: int = 0, sample_rate: int = 48000):
    """Synthesize Turkish speech from text with EMA Lightning.

    Args:
        text: Turkish text to read. Digits, dates, times, currency and
            percentages are expanded to spoken Turkish automatically; any
            written Turkish is accepted and text never raises.
        speed: Speaking rate, from 0.25 (four times slower) to 4 (four
            times faster). 1.0 is the natural pace.
        seed: Fixes the noise draw for a reproducible result. The voice is
            fixed regardless of seed.
        sample_rate: Output audio rate: 48000, 24000, 16000 or 8000 Hz.

    Returns:
        A tuple of (sample_rate, waveform) for Gradio's Audio component.
    """
    text = (text or "").strip()
    if not text:
        raise gr.Error("Lütfen okunacak bir metin girin. / Please enter some text.")
    if len(text) > MAX_CHARS:
        raise gr.Error(
            f"Metin çok uzun ({len(text)} karakter). Lütfen {MAX_CHARS} karakterin "
            f"altına indirin. / Text too long ({len(text)} chars); keep it under "
            f"{MAX_CHARS}."
        )
    speed = float(speed)
    if not 0.25 <= speed <= 4:
        speed = 1.0
    seed = max(0, int(seed))
    sample_rate = int(sample_rate) if int(sample_rate) in (48000, 24000, 16000, 8000) else 48000
    t0 = time.perf_counter()
    speech = tts.say(text, speed=speed, seed=seed, sample_rate=sample_rate)
    print(
        f"[ema-lightning] {len(text)} chars, speed={speed}, seed={seed} -> "
        f"{time.perf_counter() - t0:.2f}s synth, {speech.duration:.2f}s audio",
        flush=True,
    )
    return (speech.sample_rate, np.asarray(speech.audio, dtype=np.float32))


CSS = """
#col-container { max-width: 900px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    gr.Markdown(
        "# ⚡ EMA Lightning — Turkish Text-to-Speech\n"
        "An 8.6M-parameter (34 MB) flow-matching TTS that reads Turkish "
        "characters directly and synthesizes 48 kHz speech in four ODE steps — "
        "**0.92% WER on "
        "[Freya-TR-Eval](https://huggingface.co/datasets/freyavoice/freya-tr-eval)**, "
        "the lowest of every system measured, including ElevenLabs v4, Gemini 3.8 "
        "and the 2.38B-parameter Trendyol-TTS. One fixed voice, no cloning, "
        "Apache-2.0.\n\n"
        "[Model card](https://huggingface.co/canberkkkkkk/ema-lightning) · "
        "[PyPI package](https://pypi.org/project/ema-lightning/) · "
        "Numbers, dates, times and currency are expanded to spoken Turkish "
        "before the model reads them."
    )

    with gr.Column(elem_id="col-container"):
        text_input = gr.Textbox(
            label="Turkish text",
            placeholder="Merhaba, size nasıl yardımcı olabilirim?",
            lines=4,
            value="Merhaba, size nasıl yardımcı olabilirim?",
        )
        run_btn = gr.Button("Generate Speech", variant="primary")
        audio_output = gr.Audio(label="Generated speech", type="numpy")

        with gr.Accordion("Advanced settings", open=False):
            speed_slider = gr.Slider(
                minimum=0.25, maximum=4.0, value=1.0, step=0.05,
                label="Speed",
                info="Speaking rate. 1.0 is the natural pace.",
            )
            seed_input = gr.Number(
                label="Seed",
                value=0, precision=0,
                info="Fixes the noise draw for a reproducible result. "
                     "The voice is fixed regardless of seed.",
            )
            rate_dropdown = gr.Dropdown(
                choices=[48000, 24000, 16000, 8000],
                value=48000,
                label="Sample rate (Hz)",
                info="The model generates at 48 kHz and resamples cleanly.",
            )

        gr.Examples(
            examples=[
                ["Model sadece sekiz 8.6 milyon parametre ve çok hızlı.", 1.0, 0, 48000],
                ["15 Ekim 2026 Çarşamba günü saat 14:30'da, 1.250.000 TL "
                 "tutarındaki ödemeniz hesabınıza yatırılacak.", 1.0, 0, 48000],
                ["Sabahın erken saatlerinde liman henüz uyanmamıştı. "
                 "Balıkçılar ağlarını sessizce topluyor, martılar ise "
                 "teknelerin etrafında dönerek şanslarını deniyordu.", 1.0, 0, 48000],
            ],
            inputs=[text_input, speed_slider, seed_input, rate_dropdown],
            outputs=[audio_output],
            fn=synthesize,
            cache_examples=True,
            cache_mode="lazy",
        )

    run_btn.click(
        fn=synthesize,
        inputs=[text_input, speed_slider, seed_input, rate_dropdown],
        outputs=[audio_output],
        api_name="synthesize",
    )

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)