"""Teacher model: ConvNeXt text encoder + DiT with AdaLN-Zero, RoPE, SwiGLU, QK-norm.

docs/04-architecture.md. Shapes (``B`` batch, ``T`` latent frames, ``C`` latent channels):

* ``x_t``: ``(B, T, C)`` noisy latent; ``t``: ``(B,)`` flow time in [0, 1] (flow.py).
* ``text_ids``: ``(B, T)`` token ids, F5-style padded with ``PAD_ID`` up to ``T``.
* ``cond_latent``: ``(B, T, C)`` speaker prompt — the clean latent on prompt frames, zeros
  on frames to generate.
* ``frame_mask``: ``(B, T)`` bool, **True = real frame** (False = batch padding; such keys
  are never attended to).
* Global conditions per item: ``emotion`` ``(B,)`` long in ``[0, n_emotions]`` where
  ``n_emotions`` is the **null** id; ``intensity`` ``(B,)`` float with
  ``intensity_known`` ``(B,)`` bool (False → null embedding); ``rate_bucket`` /
  ``loudness_bucket`` ``(B,)`` long with the last id as null.
* ``drop_text`` ``(B,)`` bool: replace text by all-PAD (classifier-free guidance).
* ``language`` ``(B,)`` long in ``[0, n_languages)`` (ADR-0007; ``nida.train.data.LANGUAGES``
  order, 0 = Turkish). With ``n_languages == 0`` (models before ADR-0007) every id must be 0
  and nothing is added. The language embedding is added to every text position and is
  **not** dropped with the text: the unconditional CFG branch keeps the language (accent).

``text_path="aligned"`` (E11, ADR-0008, :mod:`nida.model.align`): ``text_ids`` are unused
(all ``PAD_ID``); the characters come unpadded in ``char_ids`` ``(B, L)`` with
``char_mask``, ``char_unit`` ``(B, L)``, ``unit_first_char`` ``(B, U)`` and
``unit_bounds`` ``(B, U + 1)`` (frame boundaries of the units). Optional ``char_phase``
``(B, L)`` overrides the within-unit positions (inference: even spacing). The text
features reach the frames through :class:`~nida.model.align.GaussianAligner`; ``drop_text``
zeroes them. ``forward(..., return_aux=True)`` also returns the duration head's
log-durations ``(B, L)`` (``log(1 + frames)``).

``dur_cond="emotion"`` (E13, aligned path only): the duration head's input is modulated by
a FiLM of the emotion condition (:class:`~nida.model.align.EmotionFiLM`), computed from the
same ``emotion`` / ``intensity`` / ``intensity_known`` the flow model sees, so a dropped
emotion (classifier-free guidance) reaches the head as the null id. The FiLM is
zero-initialised: the model starts exactly equal to ``dur_cond="none"``, which has no extra
parameters (its ``state_dict`` is unchanged). The FiLM only feeds the returned/predicted
log-durations: the training within-unit phase is built from the unmodulated head output, so
the flow loss never reaches the FiLM and the emotion enters the aligner only through the
AdaLN condition, exactly as without ``dur_cond``.

Returns the velocity ``v̂`` ``(B, T, C)``. The output projection and every AdaLN gate are
zero-initialised, so a fresh model predicts exactly 0.
"""

from __future__ import annotations

import math
from collections.abc import Callable
from dataclasses import dataclass
from typing import Literal, overload

import torch
import torch.nn.functional as F
from torch import Tensor, nn

from nida.model.align import (
    DurationHead,
    EmotionFiLM,
    GaussianAligner,
    char_phase,
    frame_coordinates,
)
from nida.model.vocab import PAD_ID, VOCAB

__all__ = [
    "ALIGNED_FIELDS",
    "CONFIGS",
    "DUR_CONDS",
    "DUR_FILM_PREFIX",
    "TEXT_PATHS",
    "Conditions",
    "DiT",
    "ModelConfig",
]

#: How text conditions reach the latent frames (E11): F5 padding or explicit durations.
TEXT_PATHS = ("implicit", "aligned")

#: What the duration head is conditioned on (E13): the characters only, or also the emotion.
DUR_CONDS = ("none", "emotion")

#: State-dict prefix of the ``dur_cond="emotion"`` parameters (absent for ``"none"``).
DUR_FILM_PREFIX = "dur_film."


@dataclass(frozen=True)
class ModelConfig:
    latent_dim: int = 64
    d_model: int = 512
    n_layers: int = 12
    n_heads: int = 8
    ffn_dim: int = 1365
    text_dim: int = 512
    text_layers: int = 4
    text_kernel: int = 7
    conv_pos_kernel: int = 31
    n_emotions: int = 7
    n_rate_buckets: int = 5
    n_loudness_buckets: int = 3
    n_languages: int = 0
    vocab_size: int = len(VOCAB)
    rope_base: float = 10000.0
    time_scale: float = 1000.0
    text_path: str = "implicit"
    align_heads: int = 4
    #: Weight of the duration loss next to the flow loss (aligned path; E11: 0.1).
    dur_loss_weight: float = 0.1
    #: ``"emotion"``: FiLM of the duration head's input on the emotion/intensity condition
    #: (E13; needs ``text_path="aligned"``). ``"none"``: characters only (E11/E12).
    dur_cond: str = "none"

    def __post_init__(self) -> None:
        sizes = ("latent_dim", "d_model", "n_layers", "n_heads", "ffn_dim", "text_dim")
        for name in (*sizes, "text_kernel", "conv_pos_kernel", "vocab_size"):
            value = getattr(self, name)
            if isinstance(value, bool) or not isinstance(value, int) or value < 1:
                raise ValueError(f"{name} must be a positive int, got {value!r}")
        for name in (
            "text_layers",
            "n_emotions",
            "n_rate_buckets",
            "n_loudness_buckets",
            "n_languages",
        ):
            value = getattr(self, name)
            if isinstance(value, bool) or not isinstance(value, int) or value < 0:
                raise ValueError(f"{name} must be a non-negative int, got {value!r}")
        if self.d_model % self.n_heads:
            raise ValueError("d_model must be divisible by n_heads")
        if (self.d_model // self.n_heads) % 2:
            raise ValueError("head dimension must be even (RoPE)")
        for name in ("text_kernel", "conv_pos_kernel"):
            if getattr(self, name) % 2 == 0:
                raise ValueError(f"{name} must be odd (same-length convolution)")
        if self.vocab_size != len(VOCAB):
            raise ValueError(f"vocab_size {self.vocab_size} != vocabulary size {len(VOCAB)}")
        if self.text_path not in TEXT_PATHS:
            raise ValueError(f"text_path must be one of {TEXT_PATHS}, got {self.text_path!r}")
        heads = self.align_heads
        if isinstance(heads, bool) or not isinstance(heads, int) or heads < 1:
            raise ValueError(f"align_heads must be a positive int, got {heads!r}")
        if self.text_path == "aligned" and (self.text_dim % heads or (self.text_dim // heads) % 2):
            raise ValueError("text_dim must split into align_heads heads of even size")
        if not (math.isfinite(self.dur_loss_weight) and self.dur_loss_weight >= 0):
            raise ValueError(f"dur_loss_weight must be >= 0, got {self.dur_loss_weight!r}")
        if self.dur_cond not in DUR_CONDS:
            raise ValueError(f"dur_cond must be one of {DUR_CONDS}, got {self.dur_cond!r}")
        if self.dur_cond != "none" and self.text_path != "aligned":
            raise ValueError(f"dur_cond={self.dur_cond!r} needs text_path='aligned'")


#: Starting sizes from docs/04 (parameter counts are measured, see tests).
CONFIGS: dict[str, ModelConfig] = {
    "tiny": ModelConfig(
        d_model=256, n_layers=4, n_heads=4, ffn_dim=683, text_dim=256, text_layers=2
    ),
    "small": ModelConfig(d_model=512, n_layers=12, n_heads=8, ffn_dim=1365, text_dim=512),
    "base": ModelConfig(d_model=768, n_layers=18, n_heads=12, ffn_dim=2048, text_dim=512),
    # ~315 M: the single "big" model sized for ~2,000 h (from-scratch decision, 2026-10-14);
    # ffn_dim is a multiple of 64 for the bf16 kernels.
    "large": ModelConfig(d_model=1024, n_layers=16, n_heads=16, ffn_dim=2752, text_dim=512),
}


@dataclass
class Conditions:
    """Per-item inputs besides ``x_t`` and ``t`` (see module docstring for shapes)."""

    text_ids: Tensor
    cond_latent: Tensor
    frame_mask: Tensor
    emotion: Tensor
    intensity: Tensor
    intensity_known: Tensor
    rate_bucket: Tensor
    loudness_bucket: Tensor
    drop_text: Tensor
    language: Tensor
    char_ids: Tensor | None = None
    char_mask: Tensor | None = None
    char_unit: Tensor | None = None
    unit_first_char: Tensor | None = None
    unit_bounds: Tensor | None = None
    char_phase: Tensor | None = None

    def map(self, fn: Callable[[Tensor], Tensor]) -> Conditions:
        """Apply ``fn`` to every tensor field (``None`` fields stay ``None``)."""
        values = {
            name: (None if getattr(self, name) is None else fn(getattr(self, name)))
            for name in Conditions.__dataclass_fields__
        }
        return Conditions(**values)  # type: ignore[arg-type]  # required fields are never None


#: Optional fields of :class:`Conditions` that ``text_path="aligned"`` requires.
ALIGNED_FIELDS = ("char_ids", "char_mask", "char_unit", "unit_first_char", "unit_bounds")


#: Standard deviation of the embedding tables at initialisation.
EMBED_INIT_STD = 0.02


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x: Tensor) -> Tensor:
        norm = x.float().pow(2).mean(dim=-1, keepdim=True).add(self.eps).rsqrt()
        return (x.float() * norm * self.weight.float()).to(x.dtype)


class ConvNeXtBlock(nn.Module):
    """Depthwise conv → LayerNorm → pointwise MLP (GELU) → residual; masked."""

    def __init__(self, dim: int, kernel: int) -> None:
        super().__init__()
        self.dw = nn.Conv1d(dim, dim, kernel, padding=kernel // 2, groups=dim)
        self.norm = nn.LayerNorm(dim)
        self.pw1 = nn.Linear(dim, dim * 2)
        self.pw2 = nn.Linear(dim * 2, dim)

    def forward(self, x: Tensor, mask: Tensor) -> Tensor:
        """x: (B, T, D); mask: (B, T) True = keep."""
        keep = mask[..., None]
        zero = torch.zeros((), dtype=x.dtype, device=x.device)
        y = self.dw(torch.where(keep, x, zero).transpose(1, 2)).transpose(1, 2)
        y = self.pw2(F.gelu(self.pw1(self.norm(y))))
        out: Tensor = torch.where(keep, x + y, zero)
        return out


def _rope(x: Tensor, base: float) -> Tensor:
    """Rotary position embedding on (B, H, T, Dh) with positions 0..T-1."""
    _, _, length, dim = x.shape
    half = dim // 2
    freqs = base ** (-torch.arange(half, device=x.device, dtype=torch.float32) / half)
    angles = torch.arange(length, device=x.device, dtype=torch.float32)[:, None] * freqs
    cos, sin = angles.cos().to(x.dtype), angles.sin().to(x.dtype)
    x1, x2 = x[..., :half], x[..., half:]
    return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)


class Attention(nn.Module):
    def __init__(self, cfg: ModelConfig) -> None:
        super().__init__()
        self.n_heads = cfg.n_heads
        self.head_dim = cfg.d_model // cfg.n_heads
        self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model)
        self.q_norm = RMSNorm(self.head_dim)
        self.k_norm = RMSNorm(self.head_dim)
        self.out = nn.Linear(cfg.d_model, cfg.d_model)
        self.rope_base = cfg.rope_base

    def forward(self, x: Tensor, key_mask: Tensor) -> Tensor:
        """x: (B, T, D); key_mask: (B, T) True = may be attended to."""
        batch, length, dim = x.shape
        qkv = self.qkv(x).view(batch, length, 3, self.n_heads, self.head_dim)
        q, k, v = qkv.permute(2, 0, 3, 1, 4)  # each (B, H, T, Dh)
        q = _rope(self.q_norm(q), self.rope_base)
        k = _rope(self.k_norm(k), self.rope_base)
        attn_mask = key_mask[:, None, None, :]  # (B, 1, 1, T) broadcast over heads/queries
        y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
        out: Tensor = self.out(y.transpose(1, 2).reshape(batch, length, dim))
        return out


class SwiGLU(nn.Module):
    def __init__(self, dim: int, hidden: int) -> None:
        super().__init__()
        self.w12 = nn.Linear(dim, 2 * hidden)
        self.w3 = nn.Linear(hidden, dim)

    def forward(self, x: Tensor) -> Tensor:
        a, b = self.w12(x).chunk(2, dim=-1)
        out: Tensor = self.w3(F.silu(a) * b)
        return out


class DiTBlock(nn.Module):
    """AdaLN-Zero block: 6 modulation vectors (shift/scale/gate × attention, MLP)."""

    def __init__(self, cfg: ModelConfig) -> None:
        super().__init__()
        self.norm1 = nn.LayerNorm(cfg.d_model, elementwise_affine=False, eps=1e-6)
        self.attn = Attention(cfg)
        self.norm2 = nn.LayerNorm(cfg.d_model, elementwise_affine=False, eps=1e-6)
        self.mlp = SwiGLU(cfg.d_model, cfg.ffn_dim)
        self.ada = nn.Linear(cfg.d_model, 6 * cfg.d_model)
        nn.init.zeros_(self.ada.weight)
        nn.init.zeros_(self.ada.bias)

    def forward(self, x: Tensor, c: Tensor, key_mask: Tensor) -> Tensor:
        """x: (B, T, D); c: (B, D) condition vector; key_mask: (B, T)."""
        mods = self.ada(F.silu(c))[:, None, :].chunk(6, dim=-1)
        shift1, scale1, gate1, shift2, scale2, gate2 = mods
        h = self.norm1(x) * (1 + scale1) + shift1
        x = x + gate1 * self.attn(h, key_mask)
        h = self.norm2(x) * (1 + scale2) + shift2
        out: Tensor = x + gate2 * self.mlp(h)
        return out


def timestep_embedding(t: Tensor, dim: int, scale: float) -> Tensor:
    """Sinusoidal embedding of ``t * scale``; t: (B,) → (B, dim)."""
    half = dim // 2
    freqs = torch.exp(-math.log(10000.0) * torch.arange(half, device=t.device) / half)
    args = (t.float() * scale)[:, None] * freqs[None]
    emb = torch.cat([args.cos(), args.sin()], dim=-1)
    return F.pad(emb, (0, dim % 2))


class DiT(nn.Module):
    def __init__(self, cfg: ModelConfig) -> None:
        super().__init__()
        self.cfg = cfg
        d = cfg.d_model
        self.text_embed = nn.Embedding(cfg.vocab_size, cfg.text_dim)
        self.text_blocks = nn.ModuleList(
            ConvNeXtBlock(cfg.text_dim, cfg.text_kernel) for _ in range(cfg.text_layers)
        )
        self.in_proj = nn.Linear(2 * cfg.latent_dim + cfg.text_dim, d)
        self.conv_pos = nn.Conv1d(
            d,
            d,
            cfg.conv_pos_kernel,
            padding=cfg.conv_pos_kernel // 2,
            groups=16 if d % 16 == 0 else 1,
        )
        self.time_mlp = nn.Sequential(nn.Linear(d, d), nn.SiLU(), nn.Linear(d, d))
        self.emotion_embed = nn.Embedding(cfg.n_emotions + 1, d)
        self.intensity_mlp = nn.Sequential(nn.Linear(1, d), nn.SiLU(), nn.Linear(d, d))
        self.intensity_null = nn.Parameter(torch.zeros(d))
        self.rate_embed = nn.Embedding(cfg.n_rate_buckets + 1, d)
        self.loudness_embed = nn.Embedding(cfg.n_loudness_buckets + 1, d)
        self.blocks = nn.ModuleList(DiTBlock(cfg) for _ in range(cfg.n_layers))
        self.final_norm = nn.LayerNorm(d, elementwise_affine=False, eps=1e-6)
        self.final_ada = nn.Linear(d, 2 * d)
        self.out_proj = nn.Linear(d, cfg.latent_dim)
        for layer in (self.final_ada, self.out_proj):
            nn.init.zeros_(layer.weight)
            nn.init.zeros_(layer.bias)
        # PyTorch's default N(0, 1) embeddings would dominate the summed condition vector
        # at init; use the usual transformer scale for the condition tables.
        # The text table keeps PyTorch's N(0, 1) init (as F5-TTS): it must stand next to the
        # unit-scale absolute position code added in forward().
        for embed in (self.emotion_embed, self.rate_embed, self.loudness_embed):
            nn.init.normal_(embed.weight, std=EMBED_INIT_STD)
        # ADR-0007: a learned language vector on the text stream (None = pre-ADR model, so
        # older checkpoints keep loading strictly). It keeps PyTorch's N(0, 1) init like the
        # text table it is added to: at 0.02 it would be ~2% of that stream and the model
        # could fit the accent from the shared Latin letters long before using it.
        self.lang_embed: nn.Embedding | None = None
        if cfg.n_languages > 0:
            self.lang_embed = nn.Embedding(cfg.n_languages, cfg.text_dim)
        # E11 aligned text path (None for implicit models, so their checkpoints load strictly).
        self.duration_head: DurationHead | None = None
        self.aligner: GaussianAligner | None = None
        if cfg.text_path == "aligned":
            self.duration_head = DurationHead(cfg.text_dim)
            self.aligner = GaussianAligner(cfg.text_dim, cfg.align_heads)
        # E13: created last, so that every other parameter draws the same random init as in
        # a dur_cond="none" model of the same seed. None for "none": strict checkpoint keys.
        self.dur_film: EmotionFiLM | None = None
        if cfg.dur_cond == "emotion":
            self.dur_film = EmotionFiLM(cfg.text_dim, cfg.n_emotions, embed_std=EMBED_INIT_STD)

    def _check(self, x_t: Tensor, t: Tensor, c: Conditions) -> None:
        batch, length, channels = x_t.shape
        if channels != self.cfg.latent_dim:
            raise ValueError(f"latent_dim {channels} != config {self.cfg.latent_dim}")
        expected = {
            "t": (t, (batch,)),
            "text_ids": (c.text_ids, (batch, length)),
            "cond_latent": (c.cond_latent, (batch, length, channels)),
            "frame_mask": (c.frame_mask, (batch, length)),
            "emotion": (c.emotion, (batch,)),
            "intensity": (c.intensity, (batch,)),
            "intensity_known": (c.intensity_known, (batch,)),
            "rate_bucket": (c.rate_bucket, (batch,)),
            "loudness_bucket": (c.loudness_bucket, (batch,)),
            "drop_text": (c.drop_text, (batch,)),
            "language": (c.language, (batch,)),
        }
        for name, (tensor, shape) in expected.items():
            if tuple(tensor.shape) != shape:
                raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}")
        for name in ("frame_mask", "intensity_known", "drop_text"):
            if getattr(c, name).dtype != torch.bool:
                raise ValueError(f"{name} must be bool")
        if not bool(c.frame_mask.any(dim=1).all()):
            raise ValueError("every item needs at least one real frame")
        if c.language.dtype != torch.long:
            raise ValueError("language must be long")
        upper = max(self.cfg.n_languages, 1)
        if bool(((c.language < 0) | (c.language >= upper)).any()):
            raise ValueError(
                f"language ids must be in [0, {upper}) for n_languages={self.cfg.n_languages}"
            )
        if self.cfg.text_path == "aligned":
            self._check_aligned(c, batch)

    @staticmethod
    def _check_aligned(c: Conditions, batch: int) -> None:
        missing = [name for name in ALIGNED_FIELDS if getattr(c, name) is None]
        if missing:
            raise ValueError(f"text_path='aligned' needs {missing}")
        char_ids, char_mask, char_unit = c.char_ids, c.char_mask, c.char_unit
        first, bounds = c.unit_first_char, c.unit_bounds
        assert char_ids is not None and char_mask is not None and char_unit is not None
        assert first is not None and bounds is not None
        length = char_ids.shape[1] if char_ids.ndim == 2 else -1
        units = first.shape[1] if first.ndim == 2 else -1
        shapes = {
            "char_ids": (char_ids, (batch, length)),
            "char_mask": (char_mask, (batch, length)),
            "char_unit": (char_unit, (batch, length)),
            "unit_first_char": (first, (batch, units)),
            "unit_bounds": (bounds, (batch, units + 1)),
        }
        if c.char_phase is not None:
            shapes["char_phase"] = (c.char_phase, (batch, length))
        for name, (tensor, shape) in shapes.items():
            if tuple(tensor.shape) != shape:
                raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}")
        if char_mask.dtype != torch.bool:
            raise ValueError("char_mask must be bool")
        if not bool(char_mask.any(dim=1).all()):
            raise ValueError("every item needs at least one character")
        if bool(((char_unit < 0) | (char_unit >= units)).any()):
            raise ValueError("char_unit outside [0, U)")
        if bool((bounds[:, 1:] < bounds[:, :-1]).any()):
            raise ValueError("unit_bounds must be non-decreasing")

    def condition_vector(self, t: Tensor, c: Conditions) -> Tensor:
        """Sum of time and global-condition embeddings: (B, D)."""
        d = self.cfg.d_model
        cond = self.time_mlp(
            timestep_embedding(t, d, self.cfg.time_scale).to(self.out_proj.weight.dtype)
        )
        cond = cond + self.emotion_embed(c.emotion)
        # Unknown intensities may be NaN placeholders: zero them BEFORE the MLP, otherwise
        # where() still back-propagates 0 * NaN into the MLP weights.
        intensity = torch.where(c.intensity_known, c.intensity, torch.zeros_like(c.intensity))
        known = self.intensity_mlp(intensity[:, None].to(cond.dtype))
        cond = cond + torch.where(c.intensity_known[:, None], known, self.intensity_null[None, :])
        cond = cond + self.rate_embed(c.rate_bucket) + self.loudness_embed(c.loudness_bucket)
        result: Tensor = cond
        return result

    def _encode_chars(self, c: Conditions) -> tuple[Tensor, Tensor, Tensor]:
        """Encoder output ``(B, L, text_dim)``, log-durations ``(B, L)`` and phase log-durations.

        The second value is the duration head's prediction (the one the duration loss and
        :meth:`predict_durations` use). With ``dur_cond="emotion"`` the head sees the
        FiLM-modulated features there; the third value is the head's output on the plain
        encoder output, which is what the training within-unit phase is built from, so
        the emotion reaches neither the aligner nor (through it) the FiLM via the flow loss.
        Without a FiLM the two are the same tensor.
        """
        assert self.duration_head is not None
        char_ids, char_mask = c.char_ids, c.char_mask
        assert char_ids is not None and char_mask is not None
        ids = torch.where(char_mask, char_ids, torch.full_like(char_ids, PAD_ID))
        h = self.text_embed(ids)
        if self.lang_embed is not None:
            h = h + self.lang_embed(c.language)[:, None, :]
        for block in self.text_blocks:
            h = block(h, char_mask)
        plain = self.duration_head(h)
        if self.dur_film is None:
            return h, plain, plain
        feats = self.dur_film(h, c.emotion, c.intensity, c.intensity_known)
        return h, self.duration_head(feats), plain

    @torch.no_grad()
    def log_durations(self, c: Conditions) -> Tensor:
        """The duration head's ``log(1 + frames)`` ``(B, L)`` (meaningful where ``char_mask``).

        Same inputs as :meth:`predict_durations`; used to compare the head under different
        emotion conditions (the validation tempo metric).
        """
        if self.cfg.text_path != "aligned":
            raise ValueError("log_durations needs text_path='aligned'")
        if c.char_ids is None or c.char_mask is None:
            raise ValueError("log_durations needs char_ids and char_mask")
        return self._encode_chars(c)[1]

    @torch.no_grad()
    def predict_durations(self, c: Conditions) -> Tensor:
        """Predicted character durations in frames ``(B, L)`` (>= 0; 0 on padding).

        Only the character fields of ``c`` are read (no frames needed), plus ``emotion``,
        ``intensity`` and ``intensity_known`` when ``dur_cond="emotion"`` (the durations are
        then those of that emotion; the null id gives the emotion-averaged rhythm).
        """
        if self.cfg.text_path != "aligned":
            raise ValueError("predict_durations needs text_path='aligned'")
        if c.char_ids is None or c.char_mask is None:
            raise ValueError("predict_durations needs char_ids and char_mask")
        log_dur = self.log_durations(c)
        zero = torch.zeros((), device=log_dur.device)
        return torch.where(c.char_mask, torch.expm1(log_dur), zero)

    def _aligned_text(self, c: Conditions, length: int) -> tuple[Tensor, Tensor]:
        """Aligned text features ``(B, T, text_dim)`` and log-durations ``(B, L)``."""
        assert self.aligner is not None
        char_mask, char_unit = c.char_mask, c.char_unit
        first, bounds = c.unit_first_char, c.unit_bounds
        assert char_mask is not None and char_unit is not None
        assert first is not None and bounds is not None
        h, log_dur, phase_log_dur = self._encode_chars(c)
        if c.char_phase is not None:
            phase = c.char_phase.float()
        else:
            phase = char_phase(torch.expm1(phase_log_dur), char_unit, first, char_mask)
        char_coord = char_unit.float() + phase
        frame_unit, frame_coord = frame_coordinates(bounds, length)
        z = self.aligner(h, char_mask, char_unit, char_coord, frame_unit, frame_coord)
        keep_text = ~c.drop_text[:, None, None]
        z = torch.where(keep_text, z, torch.zeros((), dtype=z.dtype, device=z.device))
        return z, log_dur

    @overload
    def forward(
        self, x_t: Tensor, t: Tensor, c: Conditions, *, return_aux: Literal[False] = ...
    ) -> Tensor: ...

    @overload
    def forward(
        self, x_t: Tensor, t: Tensor, c: Conditions, *, return_aux: Literal[True]
    ) -> tuple[Tensor, Tensor]: ...

    def forward(
        self, x_t: Tensor, t: Tensor, c: Conditions, *, return_aux: bool = False
    ) -> Tensor | tuple[Tensor, Tensor]:
        """Velocity ``(B, T, C)``; with ``return_aux`` (aligned path only) also the
        log-durations ``(B, L)``."""
        self._check(x_t, t, c)
        if self.cfg.text_path == "aligned":
            text, log_dur = self._aligned_text(c, x_t.shape[1])
            velocity = self._backbone(x_t, t, c, text)
            return (velocity, log_dur) if return_aux else velocity
        if return_aux:
            raise ValueError("return_aux needs text_path='aligned'")
        mask = c.frame_mask
        text_ids = torch.where(
            c.drop_text[:, None], torch.full_like(c.text_ids, PAD_ID), c.text_ids
        )
        # Absolute sinusoidal positions (F5-TTS TextEmbedding): the only absolute position
        # signal in the model. RoPE and the conv position code are relative, so without it
        # frame i cannot find "its" character near i * chars_per_frame (measured: a toy
        # model never used unaligned text without it, 2026-10-09).
        positions = torch.arange(text_ids.shape[1], device=text_ids.device)
        text = (
            self.text_embed(text_ids)
            + timestep_embedding(positions, self.cfg.text_dim, 1.0).to(
                self.text_embed.weight.dtype
            )[None]
        )
        if self.lang_embed is not None:
            text = text + self.lang_embed(c.language)[:, None, :]
        for block in self.text_blocks:
            text = block(text, mask)
        return self._backbone(x_t, t, c, text)

    def _backbone(self, x_t: Tensor, t: Tensor, c: Conditions, text: Tensor) -> Tensor:
        """DiT over per-frame ``[x_t, prompt, text]``; ``text`` ``(B, T, text_dim)``."""
        mask = c.frame_mask
        keep = mask[..., None]
        zero = torch.zeros((), dtype=x_t.dtype, device=x_t.device)
        x_in = torch.where(keep, x_t, zero)
        cond_in = torch.where(keep, c.cond_latent, zero)
        # The residual stream keeps the parameter dtype (float32 under bf16 autocast): block
        # outputs (bf16) are added to a float32 tensor, so small updates are not rounded away.
        h = self.in_proj(torch.cat([x_in, cond_in, text], dim=-1))
        h = h.to(self.in_proj.weight.dtype)
        h = h + self.conv_pos(torch.where(keep, h, zero).transpose(1, 2)).transpose(1, 2)
        cond = self.condition_vector(t, c)
        for dit_block in self.blocks:
            h = dit_block(h, cond, mask)
        shift, scale = self.final_ada(F.silu(cond))[:, None, :].chunk(2, dim=-1)
        out = self.out_proj(self.final_norm(h) * (1 + scale) + shift)
        result: Tensor = torch.where(keep, out, zero)
        return result