"""Gradio demo for VAE-based unsupervised phoneme alignment.

Shows what the aligner does: given speech and its text, it places every
phoneme (and every sub-phoneme state) on the time axis.  The result is
presented as a playable, clickable alignment ribbon rather than as files
to download.

Run locally from the repository root:

    uv run --with gradio python demo/app.py

The same file is the entry point of the Hugging Face Space (see
demo/README.md).  Everything the demo needs beyond the installed
``vae_speech_align`` package -- model checkpoints, per-language configs,
phoneme inventories and the sample recordings -- is taken from the
repository checkout when the demo runs inside one, and downloaded from
GitHub otherwise.
"""

from __future__ import annotations

import base64
import hashlib
import io
import json
import logging
import os
import threading
import time
import urllib.request
import wave
from dataclasses import dataclass
from functools import lru_cache
from pathlib import Path
from typing import Callable

import gradio as gr
import numpy as np
import torch
from omegaconf import OmegaConf

from vae_speech_align.config import (
    AlignmentExpConfig,
    AlignmentImplementation,
)
from vae_speech_align.g2p.kana import (
    alphabet_to_kana_table,
    kana_to_phoneme_table,
)
from vae_speech_align.model.aco_feat_extractor import (
    create_aco_feat_extractor,
)
from vae_speech_align.model.model import Model

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s [%(name)s] %(message)s"
)
_logger = logging.getLogger("vae_speech_align.demo")
if os.environ.get("VSA_DEMO_DEBUG"):
    # Logs the text that reaches each handler, which is what tells a
    # front-end problem apart from one in the browser.
    _logger.setLevel(logging.DEBUG)

# --------------------------------------------------------------------------
# Where the demo's assets come from
# --------------------------------------------------------------------------

GITHUB_REPO = "CyberAgentAILab/vae_speech_align"
MODEL_RELEASE_TAG = "models-v3"

# Configs, phoneme inventories and the sample recordings sit next to
# this file, so they are part of the Space and are served from the
# application directory like any other file it owns.
ASSETS = Path(__file__).resolve().parent / "assets"
REPO_ROOT = Path(__file__).resolve().parent.parent
CACHE_DIR = Path(
    os.environ.get(
        "VSA_DEMO_CACHE", Path.home() / ".cache" / "vae_speech_align_demo"
    )
)

WAV_SAMPLE_RATE = 16000
MAX_DURATION_SEC = 20.0
MIN_DURATION_SEC = 0.2


def _download(url: str, dest: Path) -> None:
    dest.parent.mkdir(parents=True, exist_ok=True)
    tmp = dest.with_suffix(dest.suffix + ".part")
    _logger.info(f"Downloading {url}")
    request = urllib.request.urlopen(url, timeout=60)
    with request as response, open(tmp, "wb") as f:
        while chunk := response.read(1 << 20):
            f.write(chunk)
    tmp.replace(dest)


def model_file(lang_key: str, asset: str, sha256: str) -> Path:
    """Return the pretrained checkpoint, downloading it if necessary.

    The checkpoints live in GitHub Releases (they are not in git); the
    SHA-256 is verified exactly as ``run_align.sh`` does.
    """
    for candidate in (
        REPO_ROOT / "example" / "inference" / lang_key / "model.safetensors",
        CACHE_DIR / asset,
    ):
        if candidate.exists():
            return candidate

    dest = CACHE_DIR / asset
    _download(
        f"https://github.com/{GITHUB_REPO}/releases/download/"
        f"{MODEL_RELEASE_TAG}/{asset}",
        dest,
    )

    digest = hashlib.sha256(dest.read_bytes()).hexdigest()
    if digest != sha256:
        dest.unlink()
        raise RuntimeError(
            f"{asset} does not match the expected SHA-256 checksum; "
            "the download may be corrupted"
        )
    return dest


# --------------------------------------------------------------------------
# Languages
# --------------------------------------------------------------------------


@dataclass(frozen=True)
class Language:
    key: str
    label: str
    model_asset: str
    model_sha256: str
    ssl_model: str
    text_placeholder: str

    @property
    def asset_dir(self) -> Path:
        return ASSETS / self.key

    @property
    def config_path(self) -> Path:
        return self.asset_dir / "config.yaml"

    @property
    def phoneme_list_path(self) -> Path:
        return self.asset_dir / "all_phonemes.txt"


LANGUAGES: dict[str, Language] = {
    "ja": Language(
        key="ja",
        label="Japanese",
        model_asset="model_ja.safetensors",
        model_sha256=(
            "2ab637d8d26ec47fa914c2692e1a16e7fed0db350979f84d8d45a6f32e3c0234"
        ),
        ssl_model="yky-h/japanese-hubert-base",
        text_placeholder="アカリオケシテ",
    ),
    "en": Language(
        key="en",
        label="English",
        model_asset="model_en.safetensors",
        model_sha256=(
            "dc0371d94c06a9dde42285d4d8769c13d5d4f0b52a742b753e38267386ba04c8"
        ),
        ssl_model="facebook/wav2vec2-large-xlsr-53",
        text_placeholder="Turn off the lights.",
    ),
}

DEFAULT_UI = "en"


# --------------------------------------------------------------------------
# Model
# --------------------------------------------------------------------------

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")


@dataclass
class Aligner:
    """Everything needed to align one utterance in one language."""

    language: Language
    model: Model
    aco_feat_extractor: torch.nn.Module
    phonemes: list[str]
    frame_shift: float
    states_per_token: int
    aco_upsample_scale: int

    @property
    def phoneme_index(self) -> dict[str, int]:
        return {p: i for i, p in enumerate(self.phonemes)}


_ALIGNERS: dict[str, Aligner] = {}
_ALIGNER_LOCK = threading.Lock()


def load_aligner(lang_key: str) -> Aligner:
    """Build (or return the cached) aligner for one language.

    The lock keeps the background preloader and an early request from
    loading the same model twice.
    """
    with _ALIGNER_LOCK:
        if lang_key not in _ALIGNERS:
            _ALIGNERS[lang_key] = _build_aligner(lang_key)
        return _ALIGNERS[lang_key]


def _build_aligner(lang_key: str) -> Aligner:
    language = LANGUAGES[lang_key]

    conf: AlignmentExpConfig = OmegaConf.merge(  # type: ignore[assignment]
        OmegaConf.structured(AlignmentExpConfig),
        OmegaConf.load(language.config_path),
    )

    phonemes = [
        line.strip()
        for line in language.phoneme_list_path.read_text().splitlines()
        if line.strip()
    ]
    if len(phonemes) != conf.features.linguistic.num_tokens:
        raise RuntimeError(
            f"{language.key}: phoneme list has {len(phonemes)} entries "
            f"but the config declares "
            f"{conf.features.linguistic.num_tokens}"
        )

    model = Model(
        conf.model,
        conf.features,
        alignment_impl=AlignmentImplementation.NUMBA,
    )
    from safetensors.torch import load_file

    model.load_state_dict(
        load_file(
            model_file(
                language.key, language.model_asset, language.model_sha256
            )
        )
    )
    model.eval()
    model.to(DEVICE)

    aco_feat_extractor = create_aco_feat_extractor(
        conf.features.acoustic_feature_extractor
    )
    aco_feat_extractor.to(DEVICE)

    # One aligned frame lasts 1 / (feature rate x acoustic upsampling) seconds.
    frame_shift = 1.0 / (
        aco_feat_extractor.frame_rate * model.acoustic_encoder.upsample_scale
    )

    _logger.info(
        f"Loaded {language.label} aligner on {DEVICE} "
        f"(frame shift {frame_shift * 1000:.1f} ms)"
    )

    return Aligner(
        language=language,
        model=model,
        aco_feat_extractor=aco_feat_extractor,
        phonemes=phonemes,
        frame_shift=frame_shift,
        states_per_token=model.linguistic_encoder.upsample_scale,
        aco_upsample_scale=model.acoustic_encoder.upsample_scale,
    )


_G2P_LOCK = threading.Lock()


@lru_cache(maxsize=1)
def _english_g2p() -> Callable[[str], str]:
    """Build the English front-end once (it downloads NLTK data)."""
    with _G2P_LOCK:
        from vae_speech_align.g2p.english import EnglishG2p

        return EnglishG2p().text_to_phoneme


def text_to_phonemes(lang_key: str, text: str) -> str:
    """Run the language's text front-end."""
    text = (text or "").strip()
    if not text:
        return ""
    if lang_key == "ja":
        from vae_speech_align.g2p.kana import kana_text_to_phoneme

        return kana_text_to_phoneme(text)
    return _english_g2p()(text)


# Every character the kana tables can read. Anything else is dropped
# without a word, which silently shortens the phoneme sequence.
_KANA_CHARS = frozenset(
    char for key in kana_to_phoneme_table for char in key
) | frozenset(alphabet_to_kana_table)


def front_end_notice(ui: str, lang_key: str, text: str) -> str:
    """Name the characters the front-end will ignore, if any."""
    if lang_key != "ja":
        return ""

    ignored = sorted(
        {c for c in (text or "") if c not in _KANA_CHARS and not c.isspace()}
    )
    if not ignored:
        return ""

    return ui_text(ui)["notice"].format(
        chars=" ".join(f"`{c}`" for c in ignored)
    )


# --------------------------------------------------------------------------
# Alignment
# --------------------------------------------------------------------------


@dataclass
class Segment:
    label: str
    start: float
    end: float


@dataclass
class AlignmentResult:
    tokens: list[Segment]
    states: list[Segment]
    log_likelihood: float
    gamma: np.ndarray  # [T, K] posterior over the aligned grid
    path: np.ndarray  # [T, K] Viterbi path indicator
    frame_shift: float
    duration: float


def _prepare_wave(audio: tuple[int, np.ndarray]) -> np.ndarray:
    """Gradio audio -> mono float32 waveform at 16 kHz."""
    sample_rate, data = audio

    data = np.asarray(data)
    scale = (
        float(np.iinfo(data.dtype).max)
        if np.issubdtype(data.dtype, np.integer)
        else 1.0
    )

    wave_np = data.astype(np.float32)
    if wave_np.ndim > 1:
        wave_np = wave_np.mean(axis=1)
    wave_np /= scale

    # No level normalisation: the command-line pipeline feeds the decoded
    # waveform to the SSL model as is, and the demo must match it.
    wave_t = torch.from_numpy(wave_np)
    if sample_rate != WAV_SAMPLE_RATE:
        import torchaudio.functional as AF

        wave_t = AF.resample(wave_t, sample_rate, WAV_SAMPLE_RATE)

    return wave_t.numpy().astype(np.float32)


def align(
    ui: str, lang_key: str, wave_np: np.ndarray, phonemes: list[str]
) -> AlignmentResult:
    strings = ui_text(ui)
    aligner = load_aligner(lang_key)

    index = aligner.phoneme_index
    unknown = sorted({p for p in phonemes if p not in index})
    if unknown:
        raise gr.Error(
            strings["err_unknown"].format(
                language=strings["language_choices"][lang_key],
                tokens=", ".join(unknown),
            )
        )

    # 0 is the padding index, so tokens are 1-indexed (see dataset.Collate).
    x = torch.LongTensor([[1 + index[p] for p in phonemes]]).to(DEVICE)
    x_lengths = torch.LongTensor([len(phonemes)]).to(DEVICE)

    wav = torch.from_numpy(wave_np).unsqueeze(0).to(DEVICE)
    wav_lengths = torch.LongTensor([wave_np.shape[0]]).to(DEVICE)

    model = aligner.model
    with torch.no_grad():
        y, y_lengths = aligner.aco_feat_extractor(wav, wav_lengths)

        num_states = int(x_lengths.item()) * aligner.states_per_token
        num_frames = int(y_lengths.item()) * aligner.aco_upsample_scale
        if num_states > num_frames:
            raise gr.Error(
                strings["err_needs_audio"].format(
                    needed=num_states * aligner.frame_shift,
                    tokens=len(phonemes),
                    per=aligner.states_per_token,
                    duration=wave_np.shape[0] / WAV_SAMPLE_RATE,
                )
            )

        x_out = model.forward_x(x, x_lengths)
        y_out = model.forward_y(y, y_lengths)

        viterbi = model.calc_viterbi(x_out, y_out)
        gamma = model.calc_gamma(x_out, y_out)

    # The DP matrices are padded with one entry on each side of both axes.
    path = viterbi.path[0, 1:-1, 1:-1].cpu().numpy()
    gamma_np = gamma[0, 1:-1, 1:-1].cpu().numpy()

    state_durations = path.sum(axis=0)  # frames per state
    token_durations = state_durations.reshape(
        -1, aligner.states_per_token
    ).sum(axis=1)

    frame_shift = aligner.frame_shift

    def _segments(durations: np.ndarray, labels: list[str]) -> list[Segment]:
        boundaries = np.concatenate(
            [[0.0], np.cumsum(durations) * frame_shift]
        )
        return [
            Segment(label, float(boundaries[i]), float(boundaries[i + 1]))
            for i, label in enumerate(labels)
        ]

    state_labels = [
        f"{p}{i + 1}"
        for p in phonemes
        for i in range(aligner.states_per_token)
    ]

    return AlignmentResult(
        tokens=_segments(token_durations, phonemes),
        states=_segments(state_durations, state_labels),
        log_likelihood=float(viterbi.log_likelihoods.sum().item()),
        gamma=gamma_np,
        path=path,
        frame_shift=frame_shift,
        duration=wave_np.shape[0] / WAV_SAMPLE_RATE,
    )


# --------------------------------------------------------------------------
# Rendering
# --------------------------------------------------------------------------

PIXELS_PER_SECOND = 150
MIN_TRACK_PX = 720
MAX_TRACK_PX = 12000
IMAGE_HEIGHT_PX = 190
IMAGE_SCALE = 2  # render at 2x for crisp display on high-DPI screens

PANEL_BG = "#12161f"
WAVE_COLOR = "#6ea8ff"
BOUNDARY_COLOR = "#ffd166"


def _wav_data_uri(wave_np: np.ndarray) -> str:
    clipped = np.clip(wave_np, -1.0, 1.0)
    pcm = (clipped * 32767.0).astype("<i2")

    buffer = io.BytesIO()
    with wave.open(buffer, "wb") as wav_file:
        wav_file.setnchannels(1)
        wav_file.setsampwidth(2)
        wav_file.setframerate(WAV_SAMPLE_RATE)
        wav_file.writeframes(pcm.tobytes())

    encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
    return f"data:audio/wav;base64,{encoded}"


def _figure_data_uri(figure, fmt: str = "png") -> str:
    """Render a figure to a self-contained data URI.

    The spectrogram is stored as JPEG: it is photographic, and a PNG of
    the same panel is roughly ten times larger to ship to the browser.
    """
    from matplotlib.backends.backend_agg import FigureCanvasAgg

    FigureCanvasAgg(figure)
    buffer = io.BytesIO()
    figure.savefig(
        buffer,
        format=fmt,
        facecolor=figure.get_facecolor(),
        pil_kwargs={"quality": 88} if fmt == "jpg" else None,
    )
    mime = "jpeg" if fmt == "jpg" else fmt
    encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
    return f"data:image/{mime};base64,{encoded}"


def _track_width(duration: float) -> int:
    return int(
        min(
            MAX_TRACK_PX,
            max(MIN_TRACK_PX, round(duration * PIXELS_PER_SECOND)),
        )
    )


def _spectrogram_uri(
    wave_np: np.ndarray,
    duration: float,
    boundaries: list[float],
    width_px: int,
) -> str:
    """Waveform over spectrogram, drawn edge to edge so that the x axis
    maps exactly onto the ribbon's time axis."""
    from matplotlib.figure import Figure

    dpi = 100
    figure = Figure(
        figsize=(
            width_px * IMAGE_SCALE / dpi,
            IMAGE_HEIGHT_PX * IMAGE_SCALE / dpi,
        ),
        dpi=dpi,
        facecolor=PANEL_BG,
    )
    grid = figure.add_gridspec(
        2,
        1,
        height_ratios=[1.0, 2.2],
        hspace=0.0,
        left=0,
        right=1,
        top=1,
        bottom=0,
    )

    ax_wave = figure.add_subplot(grid[0])
    ax_wave.set_facecolor(PANEL_BG)
    times = np.arange(wave_np.shape[0]) / WAV_SAMPLE_RATE
    ax_wave.plot(times, wave_np, linewidth=0.5 * IMAGE_SCALE, color=WAVE_COLOR)
    amplitude = max(float(np.abs(wave_np).max()), 1e-3)
    ax_wave.set_ylim(-amplitude * 1.05, amplitude * 1.05)

    ax_spec = figure.add_subplot(grid[1])
    ax_spec.set_facecolor(PANEL_BG)
    with np.errstate(divide="ignore"):
        # Digital silence gives empty bins; matplotlib takes log10 of them.
        ax_spec.specgram(
            wave_np,
            NFFT=512,
            Fs=WAV_SAMPLE_RATE,
            noverlap=512 - 80,
            cmap="magma",
            xextent=(0.0, duration),
        )
    ax_spec.set_ylim(0, 8000)

    for axis in (ax_wave, ax_spec):
        axis.set_xlim(0, duration)
        axis.axis("off")
        for boundary in boundaries:
            axis.axvline(
                boundary,
                color=BOUNDARY_COLOR,
                linewidth=0.7 * IMAGE_SCALE,
                alpha=0.75,
            )

    return _figure_data_uri(figure, fmt="jpg")


def _ruler_ticks(duration: float, width_px: int) -> list[float]:
    for step in (0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0):
        if step / max(duration, 1e-6) * width_px >= 70:
            return list(np.arange(0.0, duration + 1e-9, step))
    return [0.0, duration]


def ribbon_labels(ui: str) -> dict[str, str]:
    """Strings the ribbon script renders on the client."""
    strings = ui_text(ui)
    return {
        "empty": strings["ribbon_empty"],
        "hint": strings["ribbon_hint"],
        "speed": strings["speed"],
        "loop": strings["loop"],
    }


def encode_payload(payload: dict) -> str:
    """Serialise a ribbon payload for embedding in the page."""
    text = json.dumps(payload)
    return text.replace("<", "\\u003c").replace("&", "\\u0026")


def ribbon_empty(ui: str) -> str:
    """Payload for the ribbon before anything has been aligned."""
    return encode_payload({"labels": ribbon_labels(ui)})


def relabel_payload(ui: str, payload: str) -> str:
    """Re-render an existing ribbon in another interface language."""
    if not payload:
        return ribbon_empty(ui)
    data = json.loads(payload)
    data["labels"] = ribbon_labels(ui)
    return encode_payload(data)


def ribbon_payload(
    ui: str, result: AlignmentResult, wave_np: np.ndarray
) -> str:
    width_px = _track_width(result.duration)
    boundaries = [segment.end for segment in result.tokens[:-1]]

    def _serialise(segments: list[Segment]) -> list[dict]:
        return [
            {
                "label": segment.label,
                "start": round(segment.start, 4),
                "end": round(segment.end, 4),
            }
            for segment in segments
        ]

    payload = {
        "duration": round(result.duration, 4),
        "width": width_px,
        "image": _spectrogram_uri(
            wave_np, result.duration, boundaries, width_px
        ),
        "audio": _wav_data_uri(wave_np),
        "tokens": _serialise(result.tokens),
        "states": _serialise(result.states),
        "ticks": [
            round(t, 3) for t in _ruler_ticks(result.duration, width_px)
        ],
        "labels": ribbon_labels(ui),
    }
    return encode_payload(payload)


def gamma_figure(result: AlignmentResult, phonemes: list[str]):
    """The forward-sum posterior with the decoded Viterbi path on top."""
    from matplotlib.figure import Figure

    gamma = result.gamma
    num_frames = min(
        gamma.shape[0], int(round(result.duration / result.frame_shift))
    )
    gamma = gamma[:num_frames]
    path = result.path[:num_frames]

    figure = Figure(figsize=(11, 4.2), dpi=110, facecolor=PANEL_BG)
    axis = figure.add_subplot(111)
    axis.set_facecolor(PANEL_BG)

    extent = (0.0, num_frames * result.frame_shift, 0.0, float(gamma.shape[1]))
    axis.imshow(
        gamma.T,
        origin="lower",
        aspect="auto",
        cmap="viridis",
        extent=extent,
        vmin=0.0,
        vmax=1.0,
        interpolation="nearest",
    )

    frames, states = np.nonzero(path)
    axis.plot(
        (frames + 0.5) * result.frame_shift,
        states + 0.5,
        color=BOUNDARY_COLOR,
        linewidth=1.2,
        label="Viterbi path",
    )

    states_per_token = len(result.states) // max(len(phonemes), 1)
    ticks = [(i + 0.5) * states_per_token for i in range(len(phonemes))]
    if len(phonemes) <= 60:
        axis.set_yticks(ticks)
        axis.set_yticklabels(phonemes, fontsize=6.5)
    else:
        step = len(phonemes) // 40 + 1
        axis.set_yticks(ticks[::step])
        axis.set_yticklabels(phonemes[::step], fontsize=6.5)

    axis.set_xlabel("time (s)", color="#c9d1d9", fontsize=9)
    axis.set_ylabel("phoneme state", color="#c9d1d9", fontsize=9)
    axis.tick_params(colors="#c9d1d9", labelsize=8)
    for spine in axis.spines.values():
        spine.set_color("#30363d")
    axis.legend(
        loc="upper left",
        fontsize=8,
        facecolor=PANEL_BG,
        edgecolor="#30363d",
        labelcolor="#c9d1d9",
    )
    figure.tight_layout()
    return figure


def table_rows(result: AlignmentResult) -> list[list]:
    return [
        [
            i + 1,
            segment.label,
            round(segment.start, 3),
            round(segment.end, 3),
            round((segment.end - segment.start) * 1000),
        ]
        for i, segment in enumerate(result.tokens)
    ]


# --------------------------------------------------------------------------
# The alignment ribbon (a self-contained HTML component)
# --------------------------------------------------------------------------

RIBBON_CSS = """
.vsa-wrap { display: flex; flex-direction: column; gap: 8px; }
.vsa-scroll {
  overflow-x: auto; overflow-y: hidden;
  background: #12161f; border: 1px solid #263041; border-radius: 10px;
}
.vsa-track { position: relative; height: 252px; min-width: 100%; }
.vsa-image {
  display: block; width: 100%; height: 190px;
  user-select: none; -webkit-user-drag: none; cursor: crosshair;
}
.vsa-tier { position: relative; width: 100%; }
.vsa-tier-tokens { height: 28px; border-top: 1px solid #263041; }
.vsa-tier-states { height: 16px; }
.vsa-ruler { position: relative; height: 18px; border-top: 1px solid #263041; }
.vsa-seg {
  position: absolute; top: 0; bottom: 0;
  box-sizing: border-box; border-left: 1px solid rgba(255, 209, 102, 0.5);
  display: flex; align-items: center; justify-content: center;
  overflow: hidden; cursor: pointer; color: #e6edf3;
  font-size: 12px; line-height: 1; white-space: nowrap;
}
.vsa-seg:last-child { border-right: 1px solid rgba(255, 209, 102, 0.5); }
.vsa-seg.vsa-alt { background: rgba(110, 168, 255, 0.08); }
.vsa-seg.vsa-pause { color: #6b7684; background: rgba(255, 255, 255, 0.03); }
.vsa-seg:hover { background: rgba(110, 168, 255, 0.3); }
.vsa-seg.vsa-on { background: rgba(255, 209, 102, 0.4); color: #fff; }
.vsa-state {
  position: absolute; top: 2px; bottom: 2px; box-sizing: border-box;
  border-left: 1px solid rgba(126, 231, 199, 0.45);
  background: rgba(126, 231, 199, 0.1);
}
.vsa-state.vsa-on { background: rgba(126, 231, 199, 0.55); }
.vsa-tick {
  position: absolute; top: 0; bottom: 0;
  border-left: 1px solid #30363d; padding-left: 4px;
  color: #8b949e; font-size: 10px; line-height: 18px; white-space: nowrap;
}
.vsa-playhead {
  position: absolute; top: 0; height: 234px; width: 2px;
  background: #ff6b6b; pointer-events: none; left: 0;
  box-shadow: 0 0 6px rgba(255, 107, 107, 0.8);
}
.vsa-bar {
  display: flex; align-items: center; gap: 12px; flex-wrap: wrap;
  font-size: 13px; color: var(--body-text-color, #c9d1d9);
}
.vsa-bar audio { height: 34px; }
.vsa-now { font-variant-numeric: tabular-nums; }
.vsa-now b { font-size: 15px; color: var(--body-text-color, #e6edf3); }
.vsa-hint { font-size: 12px; opacity: 0.7; }
.vsa-empty {
  border: 1px dashed #30363d; border-radius: 10px; padding: 28px;
  text-align: center; opacity: 0.65; font-size: 14px;
}
.vsa-ctl {
  background: rgba(127, 127, 127, 0.12); color: inherit;
  border: 1px solid rgba(127, 127, 127, 0.45); border-radius: 6px;
  padding: 5px 10px; font-size: 13px; line-height: 1; cursor: pointer;
}
.vsa-ctl:hover { background: rgba(110, 168, 255, 0.22); }
.vsa-ctl.vsa-active {
  background: rgba(110, 168, 255, 0.35);
  border-color: rgba(110, 168, 255, 0.85);
}
.vsa-group { display: flex; align-items: center; gap: 4px; }
.vsa-glabel { font-size: 12px; opacity: 0.7; margin-right: 2px; }
"""

RIBBON_JS = r"""
const root = element.querySelector('.vsa-root');
const PAUSES = ['sil', 'pau'];

const div = (cls) => {
  const el = document.createElement('div');
  el.className = cls;
  return el;
};
const pct = (x) => (x * 100).toFixed(4) + '%';

const readData = () => {
  let raw = '';
  try { raw = props.value || ''; } catch (e) { raw = ''; }
  if (!raw) {
    const node = element.querySelector('.vsa-data');
    raw = node ? node.textContent : '';
  }
  raw = (raw || '').trim();
  if (!raw) return null;
  try { return JSON.parse(raw); } catch (e) { return null; }
};

let playing = null;

const render = () => {
  const data = readData();
  if (playing) {
    playing.pause();
    playing = null;
  }
  root.innerHTML = '';

  // Every visible string comes from the server, so the ribbon follows
  // the interface language.
  const labels = (data && data.labels) || {};

  if (!data || !data.tokens || !data.tokens.length) {
    const empty = div('vsa-empty');
    empty.textContent = labels.empty || '';
    root.appendChild(empty);
    return;
  }

  const dur = data.duration;
  const scroll = div('vsa-scroll');
  const track = div('vsa-track');
  track.style.width = data.width + 'px';

  const image = document.createElement('img');
  image.className = 'vsa-image';
  image.src = data.image;
  image.draggable = false;
  track.appendChild(image);

  const stateTier = div('vsa-tier vsa-tier-states');
  const stateEls = data.states.map((seg) => {
    const el = div('vsa-state');
    el.style.left = pct(seg.start / dur);
    el.style.width = pct((seg.end - seg.start) / dur);
    el.title = seg.label;
    stateTier.appendChild(el);
    return el;
  });

  const tokenTier = div('vsa-tier vsa-tier-tokens');
  const tokenEls = data.tokens.map((seg, i) => {
    const isPause = PAUSES.indexOf(seg.label) >= 0;
    const el = div(
      'vsa-seg' + (isPause ? ' vsa-pause' : i % 2 ? ' vsa-alt' : '')
    );
    el.style.left = pct(seg.start / dur);
    el.style.width = pct((seg.end - seg.start) / dur);
    el.title =
      seg.label + '  ' + seg.start.toFixed(3) + '–' + seg.end.toFixed(3) +
      ' s  (' + Math.round((seg.end - seg.start) * 1000) + ' ms)';
    const label = document.createElement('span');
    label.textContent = seg.label;
    el.appendChild(label);
    el.addEventListener('click', (event) => {
      event.stopPropagation();
      selected = i;
      playSegment(seg);
    });
    tokenTier.appendChild(el);
    return el;
  });

  const ruler = div('vsa-ruler');
  data.ticks.forEach((t) => {
    const tick = div('vsa-tick');
    tick.style.left = pct(t / dur);
    tick.textContent = t.toFixed(t < 10 ? 1 : 0) + 's';
    ruler.appendChild(tick);
  });

  const playhead = div('vsa-playhead');

  track.appendChild(tokenTier);
  track.appendChild(stateTier);
  track.appendChild(ruler);
  track.appendChild(playhead);
  scroll.appendChild(track);

  const audio = document.createElement('audio');
  playing = audio;
  audio.controls = true;
  audio.preload = 'auto';
  audio.src = data.audio;

  const now = div('vsa-now');
  now.innerHTML = '<b>—</b>';

  const bar = div('vsa-bar');
  bar.appendChild(audio);

  let loopSegment = false;
  const loopButton = document.createElement('button');
  loopButton.type = 'button';
  loopButton.className = 'vsa-ctl';
  loopButton.textContent = labels.loop || 'loop phoneme';
  loopButton.addEventListener('click', () => {
    loopSegment = !loopSegment;
    loopButton.classList.toggle('vsa-active', loopSegment);
  });

  // Buttons rather than a dropdown: a <select> is easy to miss and
  // fiddly to hit on the page.
  const SPEEDS = [['1×', 1], ['0.5×', 0.5], ['0.25×', 0.25]];
  const speedGroup = div('vsa-group');
  const speedLabel = document.createElement('span');
  speedLabel.className = 'vsa-glabel';
  speedLabel.textContent = labels.speed || 'speed';
  speedGroup.appendChild(speedLabel);

  const speedButtons = SPEEDS.map(([text, rate]) => {
    const button = document.createElement('button');
    button.type = 'button';
    button.className = 'vsa-ctl' + (rate === 1 ? ' vsa-active' : '');
    button.textContent = text;
    button.addEventListener('click', () => {
      audio.playbackRate = rate;
      speedButtons.forEach((other, i) => {
        other.classList.toggle('vsa-active', SPEEDS[i][1] === rate);
      });
    });
    speedGroup.appendChild(button);
    return button;
  });

  bar.appendChild(speedGroup);
  bar.appendChild(loopButton);
  bar.appendChild(now);

  const hint = div('vsa-hint');
  hint.textContent = labels.hint || '';

  root.appendChild(scroll);
  root.appendChild(bar);
  root.appendChild(hint);

  let selected = -1;
  let stopAt = null;
  let current = -1;
  let currentState = -1;
  let frame = null;

  const findAt = (segments, t, from) => {
    if (from >= 0 && t >= segments[from].start && t < segments[from].end) {
      return from;
    }
    for (let i = 0; i < segments.length; i++) {
      if (t >= segments[i].start && t < segments[i].end) return i;
    }
    return -1;
  };

  const playSegment = (seg) => {
    stopAt = seg.end;
    audio.currentTime = seg.start;
    audio.play();
  };

  const update = () => {
    const t = audio.currentTime;

    if (stopAt !== null && t >= stopAt - 0.004) {
      if (loopSegment && selected >= 0) {
        audio.currentTime = data.tokens[selected].start;
      } else {
        stopAt = null;
        audio.pause();
      }
    }

    // The track stretches when the viewport is wider than the ribbon,
    // so measure it rather than trusting the requested width.
    const width = track.clientWidth || data.width;
    const x = Math.min(Math.max(t / dur, 0), 1) * width;
    playhead.style.left = x + 'px';

    const index = findAt(data.tokens, t, current);
    if (index !== current) {
      if (current >= 0) tokenEls[current].classList.remove('vsa-on');
      if (index >= 0) {
        tokenEls[index].classList.add('vsa-on');
        const seg = data.tokens[index];
        now.innerHTML =
          '<b>' + seg.label + '</b> &nbsp;' + seg.start.toFixed(2) + '–' +
          seg.end.toFixed(2) + ' s &nbsp;(' +
          Math.round((seg.end - seg.start) * 1000) + ' ms)';
      } else {
        now.innerHTML = '<b>—</b>';
      }
      current = index;
    }

    const stateIndex = findAt(data.states, t, currentState);
    if (stateIndex !== currentState) {
      if (currentState >= 0) stateEls[currentState].classList.remove('vsa-on');
      if (stateIndex >= 0) stateEls[stateIndex].classList.add('vsa-on');
      currentState = stateIndex;
    }

    const left = scroll.scrollLeft;
    const view = scroll.clientWidth;
    if (x < left + 40 || x > left + view - 40) {
      scroll.scrollLeft = Math.max(0, x - view / 2);
    }
  };

  const tick = () => {
    update();
    frame = audio.paused ? null : requestAnimationFrame(tick);
  };

  audio.addEventListener('play', () => {
    if (frame === null) frame = requestAnimationFrame(tick);
  });
  audio.addEventListener('pause', () => {
    if (frame !== null) cancelAnimationFrame(frame);
    frame = null;
    update();
  });
  audio.addEventListener('seeked', update);
  audio.addEventListener('loadedmetadata', update);

  image.addEventListener('click', (event) => {
    const rect = image.getBoundingClientRect();
    const time = ((event.clientX - rect.left) / rect.width) * dur;
    selected = -1;
    stopAt = null;
    audio.currentTime = Math.min(Math.max(time, 0), dur);
    audio.play();
  });
};

try {
  watch('value', render);
} catch (e) {
  const node = element.querySelector('.vsa-data');
  if (node) new MutationObserver(render).observe(node, {
    childList: true, characterData: true, subtree: true,
  });
}

render();
"""

RIBBON_TEMPLATE = (
    '<div class="vsa-wrap">'
    '<div class="vsa-root"></div>'
    '<div class="vsa-data" style="display:none">${value}</div>'
    "</div>"
)


# --------------------------------------------------------------------------
# Sample recordings
# --------------------------------------------------------------------------

SAMPLE_IDS = {
    "ja": [f"sample_ja0{i}" for i in range(1, 6)],
    "en": [f"sample_en0{i}" for i in range(1, 6)],
}


def build_examples() -> list[list]:
    examples = []
    for key, language in LANGUAGES.items():
        for base in SAMPLE_IDS[key]:
            wav = language.asset_dir / "wav" / f"{base}.wav"
            text = (language.asset_dir / "text" / f"{base}.txt").read_text()
            examples.append([key, str(wav), text.strip()])
    return examples


# --------------------------------------------------------------------------
# Event handlers
# --------------------------------------------------------------------------


def on_language_change(ui: str, lang_key: str):
    """Model language changed: relabel the text box, clear the phonemes."""
    strings = ui_text(ui)
    return (
        gr.update(
            label=strings["text_label"][lang_key],
            placeholder=LANGUAGES[lang_key].text_placeholder,
            info=strings["text_hint"][lang_key],
        ),
        "",
        "",
        "",
    )


def fill_phonemes(
    ui: str, lang_key: str, text: str, phoneme_text: str, source_text: str
):
    """Refresh the phoneme box when its source text has changed."""
    text = (text or "").strip()
    _logger.debug("fill_phonemes: text=%r (%d chars)", text, len(text))

    if not text:
        return phoneme_text, source_text, ""

    notice = front_end_notice(ui, lang_key, text)
    if text == source_text and (phoneme_text or "").strip():
        return phoneme_text, source_text, notice

    try:
        return text_to_phonemes(lang_key, text), text, notice
    except Exception as error:  # surfaced to the user as a toast
        raise gr.Error(ui_text(ui)["err_front_end"].format(error=error))


def example_phonemes(lang_key: str, _audio, text: str):
    """Fill the phoneme box for a bundled sample.

    The samples never trip the ignored-character notice, so it is
    cleared rather than recomputed for an interface language this
    callback does not receive.
    """
    text = (text or "").strip()
    return text_to_phonemes(lang_key, text), text, ""


def _resolve_phonemes(
    ui: str, lang_key: str, text: str, phoneme_text: str, source_text: str
) -> tuple[str, str]:
    """Return the phoneme string to align, regenerating it if stale."""
    if text and (text != source_text or not phoneme_text):
        return text_to_phonemes(lang_key, text), text
    return phoneme_text, source_text


def _check_duration(ui: str, duration: float) -> None:
    strings = ui_text(ui)
    if duration < MIN_DURATION_SEC:
        raise gr.Error(strings["err_too_short"])
    if duration > MAX_DURATION_SEC:
        raise gr.Error(
            strings["err_too_long"].format(
                duration=duration, limit=MAX_DURATION_SEC
            )
        )


def _summary(
    ui: str, result: AlignmentResult, language: Language, elapsed: float
) -> str:
    return ui_text(ui)["summary"].format(
        tokens=len(result.tokens),
        states=len(result.states),
        per=len(result.states) // len(result.tokens),
        duration=result.duration,
        shift=result.frame_shift * 1000,
        likelihood=result.log_likelihood,
        elapsed=elapsed,
        rtf=elapsed / result.duration,
        device=DEVICE.type,
        ssl=language.ssl_model,
    )


def run_alignment(
    ui: str,
    lang_key: str,
    audio: tuple[int, np.ndarray] | None,
    text: str,
    phoneme_text: str,
    source_text: str,
):
    strings = ui_text(ui)
    if audio is None:
        raise gr.Error(strings["err_no_audio"])

    language = LANGUAGES[lang_key]
    text = (text or "").strip()
    phoneme_text = (phoneme_text or "").strip()
    _logger.debug(
        "run_alignment: text=%r phoneme_box=%r source_state=%r",
        text,
        phoneme_text,
        source_text,
    )

    phoneme_text, source_text = _resolve_phonemes(
        ui, lang_key, text, phoneme_text, source_text
    )
    if not phoneme_text:
        raise gr.Error(strings["err_no_text"])

    phonemes = phoneme_text.split()
    _logger.info(
        "run_alignment: %d characters of %s text -> %d phonemes",
        len(text),
        lang_key,
        len(phonemes),
    )

    wave_np = _prepare_wave(audio)
    _check_duration(ui, wave_np.shape[0] / WAV_SAMPLE_RATE)

    started = time.perf_counter()
    result = align(ui, lang_key, wave_np, phonemes)
    elapsed = time.perf_counter() - started

    payload = ribbon_payload(ui, result, wave_np)
    return (
        payload,
        _summary(ui, result, language, elapsed),
        table_rows(result),
        gamma_figure(result, phonemes),
        phoneme_text,
        source_text,
        front_end_notice(ui, lang_key, text),
        payload,
    )


# --------------------------------------------------------------------------
# UI
# --------------------------------------------------------------------------

_REPO_URL = f"https://github.com/{GITHUB_REPO}"
_PAPER_URL = "https://doi.org/10.21437/Interspeech.2024-1127"

INTRO_EN = f"""
# Phoneme and speech aligner

*Unsupervised phoneme alignment.*
[Paper]({_PAPER_URL}) · [Code]({_REPO_URL})

> **Unofficial demo.** Built by a third party, with no connection to the
> paper's author or to the `CyberAgentAILab/vae_speech_align` repository.
> It uses their published code and pretrained models.

Give it **speech** and **the words spoken in it**. It marks where every
phoneme starts and ends.
"""

INTRO_JA = f"""
# 音声と音素のアライナー

*教師なし音素アライメント*
[論文]({_PAPER_URL}) · [コード]({_REPO_URL})

> **非公式デモです。** 本デモは論文著者および
> `CyberAgentAILab/vae_speech_align` リポジトリとは無関係の、第三者が作成した
> ものです。公開されているコードと学習済みモデルを利用しています。

音声と、その読み上げテキストを入力すると、各音素の開始・終了時刻を表示します。
"""

FOOTER_EN = f"""
Pretrained models and sample audio come from
[CyberAgentAILab/vae_speech_align]({_REPO_URL}).
The samples are synthetic speech. See that repository for training data
and licensing.
"""

FOOTER_JA = f"""
学習済みモデルとサンプル音声は
[CyberAgentAILab/vae_speech_align]({_REPO_URL})
のものです。サンプルは合成音声です。学習データとライセンスについては
同リポジトリを参照してください。
"""

UI_TEXT: dict[str, dict] = {
    "en": {
        "intro": INTRO_EN,
        "footer": FOOTER_EN,
        "ui_language": "Interface",
        "language": "Language",
        "language_choices": {"ja": "Japanese", "en": "English"},
        "speech": "Speech",
        "text_label": {
            "ja": "Text (katakana, as pronounced)",
            "en": "Text",
        },
        "text_hint": {
            "ja": (
                "Write the reading in katakana: particles as they are "
                "pronounced (は→ワ, を→オ) and long vowels as ー."
            ),
            "en": "Plain English; g2p_en turns it into the VCTK ARPABET set.",
        },
        "phonemes": "Phonemes",
        "phonemes_info": (
            "Filled in by the text front-end — edit it if a word comes "
            "out wrong."
        ),
        "align": "Align",
        "samples": (
            "Samples (the Japanese and English sets say the same five "
            "sentences)"
        ),
        "table": "Phoneme table",
        "posterior": "Alignment posterior",
        "headers": ["#", "phoneme", "start (s)", "end (s)", "duration (ms)"],
        "ribbon_empty": (
            "Pick a language, add speech and its text, then press Align — "
            "the phoneme boundaries appear here."
        ),
        "ribbon_hint": (
            "Click a phoneme to hear it on its own · click the spectrogram "
            "to play from there · slow it down to hear the boundaries."
        ),
        "speed": "speed",
        "loop": "loop phoneme",
        "notice": (
            "⚠️ Ignored: {chars} — the kana front-end reads katakana only, "
            "so hiragana, kanji and half-width letters are dropped from "
            "the phonemes."
        ),
        "summary": (
            "**{tokens} phonemes** · {states} states ({per} per phoneme) · "
            "{duration:.2f} s of audio · {shift:.0f} ms frames  \n"
            "Viterbi log-likelihood **{likelihood:,.1f}** · aligned in "
            "{elapsed:.2f} s ({rtf:.2f}× real time) on `{device}` · "
            "acoustic features from `{ssl}`"
        ),
        "err_no_audio": (
            "Add a recording first (upload a file or use the microphone)."
        ),
        "err_no_text": "Add the text that is spoken in the recording.",
        "err_unknown": ("Not in the {language} model's phoneme set: {tokens}"),
        "err_too_short": "The recording is too short to align.",
        "err_too_long": (
            "The recording is {duration:.1f} s; this demo aligns up to "
            "{limit:.0f} s. Please trim it."
        ),
        "err_needs_audio": (
            "The text needs at least {needed:.1f} s of audio ({tokens} "
            "phonemes x {per} states) but the recording is "
            "{duration:.1f} s. Shorten the text or use a longer recording."
        ),
        "err_front_end": "Text front-end failed: {error}",
    },
    "ja": {
        "intro": INTRO_JA,
        "footer": FOOTER_JA,
        "ui_language": "表示言語",
        "language": "言語",
        "language_choices": {"ja": "日本語", "en": "英語"},
        "speech": "音声",
        "text_label": {
            "ja": "テキスト（発音どおりのカタカナ）",
            "en": "テキスト（英語）",
        },
        "text_hint": {
            "ja": (
                "読みをカタカナで入力してください。"
                "助詞は発音どおり（は→ワ、を→オ）、長音は「ー」で表記します。"
            ),
            "en": (
                "英語のテキストを入力してください。"
                "g2p_en が ARPABET の音素に変換します。"
            ),
        },
        "phonemes": "音素列",
        "phonemes_info": (
            "テキストから自動生成されます。誤りがあれば直接編集してください。"
        ),
        "align": "アライメント実行",
        "samples": "サンプル（日本語と英語は同じ内容の5文です）",
        "table": "音素の一覧",
        "posterior": "アライメント事後確率",
        "headers": ["#", "音素", "開始 (秒)", "終了 (秒)", "長さ (ms)"],
        "ribbon_empty": (
            "音声とテキストを入力して「アライメント実行」を押すと、"
            "ここに音素の境界が表示されます。"
        ),
        "ribbon_hint": (
            "音素をクリックするとその音だけ再生 · "
            "スペクトログラムをクリックするとその位置から再生 · "
            "速度を下げると境界が聞き取りやすくなります。"
        ),
        "speed": "速度",
        "loop": "音素をループ",
        "notice": (
            "⚠️ 無視されました: {chars} — カナのフロントエンドはカタカナのみを"
            "読みます。ひらがな・漢字・半角文字は音素に反映されません。"
        ),
        "summary": (
            "**{tokens} 音素** · {states} 状態（音素あたり {per}） · "
            "音声 {duration:.2f} 秒 · フレーム {shift:.0f} ms  \n"
            "Viterbi 対数尤度 **{likelihood:,.1f}** · 処理時間 "
            "{elapsed:.2f} 秒（実時間の {rtf:.2f} 倍） · `{device}` · "
            "音響特徴 `{ssl}`"
        ),
        "err_no_audio": (
            "先に音声を入力してください（ファイルのアップロードまたはマイク）。"
        ),
        "err_no_text": "音声で話されているテキストを入力してください。",
        "err_unknown": ("{language}モデルの音素セットにありません: {tokens}"),
        "err_too_short": "音声が短すぎてアライメントできません。",
        "err_too_long": (
            "音声が {duration:.1f} 秒です。このデモは {limit:.0f} 秒までです。"
            "短く切ってください。"
        ),
        "err_needs_audio": (
            "テキストには少なくとも {needed:.1f} 秒の音声が必要です"
            "（{tokens} 音素 × {per} 状態）が、音声は {duration:.1f} 秒です。"
            "テキストを短くするか、長い音声を使ってください。"
        ),
        "err_front_end": "テキスト変換に失敗しました: {error}",
    },
}


def ui_text(ui: str) -> dict:
    """Return the string table for the interface language."""
    return UI_TEXT.get(ui, UI_TEXT[DEFAULT_UI])


@dataclass
class Controls:
    """The input widgets, kept together so wiring stays readable."""

    ui: gr.Radio
    language: gr.Radio
    audio: gr.Audio
    text: gr.Textbox
    phonemes: gr.Textbox
    notice: gr.Markdown
    align: gr.Button


@dataclass
class Outputs:
    """The widgets the alignment writes to."""

    ribbon: gr.HTML
    info: gr.Markdown
    table: gr.Dataframe
    gamma: gr.Plot


@dataclass
class Chrome:
    """Page furniture that follows the interface language."""

    intro: gr.Markdown
    footer: gr.Markdown
    samples: gr.Dataset
    table_section: gr.Accordion
    posterior_section: gr.Accordion


UI_CHOICES = [("English", "en"), ("日本語", "ja")]


def _language_choices(ui: str) -> list[tuple[str, str]]:
    names = ui_text(ui)["language_choices"]
    return [(names[key], key) for key in LANGUAGES]


def _build_inputs(
    ui: str, ui_radio: gr.Radio, default_example: list
) -> Controls:
    """Build the input column, prefilled with the first sample."""
    strings = ui_text(ui)
    default_key, default_wav, default_text = default_example

    with gr.Row():
        with gr.Column(scale=2):
            language = gr.Radio(
                choices=_language_choices(ui),
                value=default_key,
                label=strings["language"],
            )
            audio = gr.Audio(
                sources=["upload", "microphone"],
                type="numpy",
                label=strings["speech"],
                value=default_wav,
            )
        with gr.Column(scale=3):
            text = gr.Textbox(
                label=strings["text_label"][default_key],
                placeholder=LANGUAGES[default_key].text_placeholder,
                info=strings["text_hint"][default_key],
                value=default_text,
                lines=2,
            )
            phonemes = gr.Textbox(
                label=strings["phonemes"],
                info=strings["phonemes_info"],
                value=text_to_phonemes(default_key, default_text),
                lines=2,
            )
            notice = gr.Markdown(
                front_end_notice(ui, default_key, default_text)
            )
            align = gr.Button(strings["align"], variant="primary")

    return Controls(ui_radio, language, audio, text, phonemes, notice, align)


def _build_outputs(ui: str) -> tuple[Outputs, gr.Accordion, gr.Accordion]:
    strings = ui_text(ui)

    ribbon = gr.HTML(
        value=ribbon_empty(ui),
        html_template=RIBBON_TEMPLATE,
        css_template=RIBBON_CSS,
        js_on_load=RIBBON_JS,
        apply_default_css=False,
        container=False,
        padding=False,
    )
    info = gr.Markdown("")

    with gr.Accordion(strings["table"], open=False) as table_section:
        table = gr.Dataframe(
            headers=strings["headers"],
            datatype=["number", "str", "number", "number", "number"],
            interactive=False,
            max_height=340,
        )

    with gr.Accordion(strings["posterior"], open=False) as posterior_section:
        gamma = gr.Plot(show_label=False)

    return (
        Outputs(ribbon, info, table, gamma),
        table_section,
        posterior_section,
    )


def on_ui_change(ui: str, lang_key: str, text: str, payload: str):
    """Re-render every piece of text in the chosen interface language."""
    strings = ui_text(ui)
    return (
        gr.update(value=strings["intro"]),
        gr.update(value=strings["footer"]),
        gr.update(label=strings["ui_language"]),
        gr.update(
            label=strings["language"],
            choices=_language_choices(ui),
            value=lang_key,
        ),
        gr.update(label=strings["speech"]),
        gr.update(
            label=strings["text_label"][lang_key],
            info=strings["text_hint"][lang_key],
        ),
        gr.update(label=strings["phonemes"], info=strings["phonemes_info"]),
        gr.update(value=strings["align"]),
        gr.update(label=strings["samples"]),
        gr.update(label=strings["table"]),
        gr.update(label=strings["posterior"]),
        gr.update(headers=strings["headers"]),
        relabel_payload(ui, payload),
        front_end_notice(ui, lang_key, (text or "").strip()),
        relabel_payload(ui, payload),
    )


def _wire_events(
    controls: Controls,
    outputs: Outputs,
    chrome: Chrome,
    source_state: gr.State,
    payload_state: gr.State,
) -> None:
    controls.ui.change(
        on_ui_change,
        inputs=[
            controls.ui,
            controls.language,
            controls.text,
            payload_state,
        ],
        outputs=[
            chrome.intro,
            chrome.footer,
            controls.ui,
            controls.language,
            controls.audio,
            controls.text,
            controls.phonemes,
            controls.align,
            chrome.samples,
            chrome.table_section,
            chrome.posterior_section,
            outputs.table,
            outputs.ribbon,
            controls.notice,
            payload_state,
        ],
    )

    controls.language.change(
        on_language_change,
        inputs=[controls.ui, controls.language],
        outputs=[
            controls.text,
            controls.phonemes,
            source_state,
            controls.notice,
        ],
    )

    # Follow every change to the text, not just blur/submit, so the
    # phoneme box can never sit there showing an older sentence.
    controls.text.change(
        fill_phonemes,
        inputs=[
            controls.ui,
            controls.language,
            controls.text,
            controls.phonemes,
            source_state,
        ],
        outputs=[controls.phonemes, source_state, controls.notice],
        trigger_mode="always_last",
        show_progress="hidden",
    )

    controls.align.click(
        run_alignment,
        inputs=[
            controls.ui,
            controls.language,
            controls.audio,
            controls.text,
            controls.phonemes,
            source_state,
        ],
        outputs=[
            outputs.ribbon,
            outputs.info,
            outputs.table,
            outputs.gamma,
            controls.phonemes,
            source_state,
            controls.notice,
            payload_state,
        ],
        concurrency_limit=1,
    )


def build_demo() -> gr.Blocks:
    examples = build_examples()
    default_example = examples[0]
    ui = DEFAULT_UI
    strings = ui_text(ui)

    with gr.Blocks(title="Phoneme Aligner") as demo:
        with gr.Row():
            with gr.Column(scale=5):
                intro = gr.Markdown(strings["intro"])
            with gr.Column(scale=1, min_width=150):
                ui_radio = gr.Radio(
                    choices=UI_CHOICES,
                    value=ui,
                    label=strings["ui_language"],
                )

        source_state = gr.State(default_example[2])
        payload_state = gr.State("")

        controls = _build_inputs(ui, ui_radio, default_example)
        outputs, table_section, posterior_section = _build_outputs(ui)

        samples = gr.Examples(
            examples=examples,
            inputs=[controls.language, controls.audio, controls.text],
            outputs=[controls.phonemes, source_state, controls.notice],
            fn=example_phonemes,
            run_on_click=True,
            cache_examples=False,
            label=strings["samples"],
            examples_per_page=10,
        )

        footer = gr.Markdown(strings["footer"])

        chrome = Chrome(
            intro=intro,
            footer=footer,
            samples=samples.dataset,
            table_section=table_section,
            posterior_section=posterior_section,
        )
        _wire_events(controls, outputs, chrome, source_state, payload_state)

    return demo


def warm_up() -> None:
    """Compile the numba DP kernels so the first alignment is not slow."""
    from vae_speech_align.forwardsum.numba_impl import (
        forwardsum_gamma_numba,
        forwardsum_viterbi_numba,
    )

    matrix = torch.full((1, 8, 5), -1.0)
    lengths = torch.LongTensor([6])
    states = torch.LongTensor([3])
    forwardsum_viterbi_numba(matrix, lengths, states)
    forwardsum_gamma_numba(matrix, lengths, states)
    _logger.info("Alignment kernels ready")


def preload_models() -> None:
    """Warm the caches in the background so the first click is fast."""
    for key in LANGUAGES:
        # A failed preload must not take the app down: the language is
        # simply loaded again, and reported, on first use.
        try:
            load_aligner(key)
        except Exception:
            _logger.exception(f"Could not preload the {key} model")
    try:
        _english_g2p()
    except Exception:
        _logger.exception("Could not preload the English text front-end")


if __name__ == "__main__":
    import threading

    warm_up()
    threading.Thread(target=preload_models, daemon=True).start()

    build_demo().queue(default_concurrency_limit=2).launch()