"""Explicit-duration text path (E11, ADR-0008): duration head + windowed Gaussian alignment.

Instead of padding the characters to the latent length (F5 style, ``text_path="implicit"``),
every latent frame attends softly to the characters of its own unit and the two
neighbouring units (a unit is a word or an event tag, see ``nida.data.timings``), under a
learned Gaussian prior on the distance between continuous positions (EMA-TTS report §3,
docs/experiments/E11-aligned.md). Notation (``B`` items, ``L`` characters, ``U`` units,
``T`` frames, ``H`` alignment heads):

* ``char_unit`` ``(B, L)`` long — unit index ``w(i)`` of every character;
* ``unit_first_char`` ``(B, U)`` long — index of the first character of every unit;
* ``unit_bounds`` ``(B, U + 1)`` float — frame boundaries: unit ``u`` covers frames
  ``[unit_bounds[u], unit_bounds[u + 1])``; padded units have zero length at the item end;
* masks are bool with **True = real** (character / unit / frame).

Coordinates are "unit index + fraction inside the unit", so the last character of one word
and the first of the next are close (they must interact for coarticulation):

* frame ``t``: ``g^f_t = w(t) + (t + 0.5 − s_{w(t)}) / D_{w(t)}``;
* character ``i``: ``g^c_i = w(i) + φ_i`` with
  ``φ_i = (Σ_{j<i, w(j)=w(i)} d_j + d_i / 2) / Σ_{w(j)=w(i)} d_j`` (durations ``d`` from
  the duration head during training, so the flow loss reaches the within-unit division).
"""

from __future__ import annotations

import itertools
import math
from collections.abc import Sequence

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

__all__ = [
    "DurationHead",
    "EmotionFiLM",
    "GaussianAligner",
    "char_phase",
    "even_char_phase",
    "frame_coordinates",
    "inference_bounds",
]

#: Floor added to unit duration sums so an all-zero unit never divides by zero.
_DUR_EPS = 1e-4


class DurationHead(nn.Module):
    """``l_i = log(1 + d_i) = softplus(MLP(h_i)) >= 0`` per character.

    Shapes: ``h`` ``(B, L, D)`` → ``(B, L)`` log-durations (frames). ``d = expm1(l) >= 0``.
    The softplus keeps the duration non-negative, so positions built from it are ordered.
    """

    def __init__(self, dim: int) -> None:
        super().__init__()
        self.net = nn.Sequential(nn.Linear(dim, dim), nn.SiLU(), nn.Linear(dim, 1))

    def forward(self, h: Tensor) -> Tensor:
        out: Tensor = F.softplus(self.net(h).squeeze(-1).float())
        return out


class EmotionFiLM(nn.Module):
    """FiLM of the duration-head input on the emotion condition (E13, ``dur_cond="emotion"``).

    ``h' = h * (1 + gamma(e)) + beta(e)`` with ``e = E[emotion] + g(intensity)`` and
    ``(gamma, beta) = W e`` (one ``Linear(dim, 2 * dim)``, **zero-initialised**, so a fresh
    module is the identity and the model starts exactly equal to ``dur_cond="none"``).
    ``E`` has ``n_emotions + 1`` rows, the last one being the null id (emotion dropped for
    classifier-free guidance); ``g`` is the same small MLP as the AdaLN intensity embedding
    but without a learned null: an unknown intensity contributes nothing.

    Shapes: ``h`` ``(B, L, D)``; ``emotion`` ``(B,)`` long in ``[0, n_emotions]``;
    ``intensity`` ``(B,)`` float; ``intensity_known`` ``(B,)`` bool → ``(B, L, D)``.
    """

    def __init__(self, dim: int, n_emotions: int, *, embed_std: float = 0.02) -> None:
        super().__init__()
        self.emotion_embed = nn.Embedding(n_emotions + 1, dim)
        nn.init.normal_(self.emotion_embed.weight, std=embed_std)
        self.intensity_mlp = nn.Sequential(nn.Linear(1, dim), nn.SiLU(), nn.Linear(dim, dim))
        self.proj = nn.Linear(dim, 2 * dim)
        nn.init.zeros_(self.proj.weight)
        nn.init.zeros_(self.proj.bias)

    def forward(
        self, h: Tensor, emotion: Tensor, intensity: Tensor, intensity_known: Tensor
    ) -> Tensor:
        e = self.emotion_embed(emotion)
        # Unknown intensities may be NaN placeholders: zero them BEFORE the MLP, otherwise
        # where() still back-propagates 0 * NaN into the MLP weights.
        level = torch.where(intensity_known, intensity, torch.zeros_like(intensity))
        known = self.intensity_mlp(level[:, None].to(e.dtype))
        e = e + torch.where(intensity_known[:, None], known, torch.zeros_like(known))
        gamma, beta = self.proj(e)[:, None, :].chunk(2, dim=-1)
        out: Tensor = h * (1 + gamma) + beta
        return out


def _segment_prefix(values: Tensor, char_unit: Tensor, unit_first_char: Tensor) -> Tensor:
    """Exclusive prefix sum of ``values`` ``(B, L)`` restarted at every unit's first char."""
    cumsum = torch.cumsum(values, dim=1)
    exclusive = cumsum - values
    # exclusive sum at the first character of each unit, gathered back to its characters
    start = torch.gather(exclusive, 1, unit_first_char.clamp(min=0))  # (B, U)
    return exclusive - torch.gather(start, 1, char_unit.clamp(min=0))


def char_phase(
    durations: Tensor, char_unit: Tensor, unit_first_char: Tensor, char_mask: Tensor
) -> Tensor:
    """Within-unit position ``φ`` ``(B, L)`` in ``[0, 1]`` of each character's midpoint.

    Args:
        durations: ``(B, L)`` non-negative character durations (frames).
        char_unit, unit_first_char: see module docstring (contiguous units).
        char_mask: ``(B, L)`` bool, True = real character (padding gets ``φ = 0``).
    """
    d = torch.where(char_mask, durations.float(), torch.zeros((), device=durations.device))
    prefix = _segment_prefix(d, char_unit, unit_first_char)
    total = torch.zeros(
        d.shape[0], unit_first_char.shape[1], device=d.device, dtype=d.dtype
    ).scatter_add(1, char_unit.clamp(min=0), d)
    unit_total = torch.gather(total, 1, char_unit.clamp(min=0)) + _DUR_EPS
    phase = (prefix + 0.5 * d) / unit_total
    return torch.where(char_mask, phase.clamp(0.0, 1.0), torch.zeros((), device=d.device))


def even_char_phase(char_unit: Tensor, unit_first_char: Tensor, char_mask: Tensor) -> Tensor:
    """``φ_i = (k + 0.5) / n_{w(i)}`` (even spacing inside a unit; inference, EMA §3.6)."""
    ones = char_mask.float()
    return char_phase(ones, char_unit, unit_first_char, char_mask)


def frame_coordinates(unit_bounds: Tensor, num_frames: int) -> tuple[Tensor, Tensor]:
    """Unit index ``w(t)`` ``(B, T)`` long and coordinate ``g^f_t`` ``(B, T)`` float.

    Frames at or beyond the last boundary (batch padding) get the last unit and a
    coordinate clamped inside it; they are masked by the caller.
    """
    batch, n_bounds = unit_bounds.shape
    centres = torch.arange(num_frames, device=unit_bounds.device, dtype=torch.float32) + 0.5
    centres = centres[None, :].expand(batch, num_frames).contiguous()
    bounds = unit_bounds.float().contiguous()
    unit = torch.searchsorted(bounds, centres, right=True) - 1
    unit = unit.clamp(0, n_bounds - 2)
    start = torch.gather(bounds, 1, unit)
    end = torch.gather(bounds, 1, unit + 1)
    frac = ((centres - start) / (end - start).clamp(min=1e-6)).clamp(0.0, 1.0)
    return unit, unit.float() + frac


def inference_bounds(
    known_bounds: Sequence[float], unit_frames: Sequence[float], keep_units: int, total: int
) -> list[int]:
    """Unit boundaries for generation: the first ``keep_units`` units keep their known
    boundaries (the speaker prompt), the rest share ``total − known_bounds[keep_units]``
    frames in proportion to ``unit_frames`` (predicted durations), rounded so that every
    boundary is an int, non-decreasing, and the last one is exactly ``total``.

    Args:
        known_bounds: ``U + 1`` reference boundaries (only ``0..keep_units`` are used).
        unit_frames: ``U`` non-negative predicted unit durations.
        keep_units: number of leading units to keep (``0..U``).
        total: frame count of the item.

    Raises:
        ValueError: inconsistent lengths, negative durations, or a kept boundary outside
            ``[0, total]``.
    """
    units = len(unit_frames)
    if len(known_bounds) != units + 1 or not 0 <= keep_units <= units:
        raise ValueError("need U + 1 known bounds and 0 <= keep_units <= U")
    if any(d < 0 or not math.isfinite(d) for d in unit_frames):
        raise ValueError("unit durations must be finite and >= 0")
    kept = [round(b) for b in known_bounds[: keep_units + 1]]
    if any(b < 0 or b > total for b in kept) or any(b < a for a, b in itertools.pairwise(kept)):
        raise ValueError("kept boundaries must be non-decreasing within [0, total]")
    start = kept[-1]
    rest = list(unit_frames[keep_units:])
    room = total - start
    weight = sum(rest)
    shares = (
        [room / len(rest)] * len(rest)
        if rest and weight <= 0
        else [room * d / weight for d in rest]
    )
    out = kept[:]
    acc = float(start)
    for share in shares:
        acc += share
        out.append(round(acc))
    out[-1] = total
    return out


def _rotate(x: Tensor, positions: Tensor, base: float) -> Tensor:
    """RoPE at continuous ``positions``: ``x`` ``(..., N, H, Dh)``, ``positions`` ``(..., N)``."""
    dim = x.shape[-1]
    half = dim // 2
    freqs = base ** (-torch.arange(half, device=x.device, dtype=torch.float32) / half)
    angles = positions.float()[..., None, None] * freqs  # (..., N, 1, half)
    cos, sin = angles.cos(), angles.sin()
    x1, x2 = x[..., :half].float(), x[..., half:].float()
    return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)


class GaussianAligner(nn.Module):
    """Windowed Gaussian cross-attention from frames to characters (E11 §Değişiklik 4).

    ``A^{(m)}_{t,i} = softmax_i(τ ⟨q_t, k_i⟩ / √d_h − β_m (g^f_t − g^c_i)² / (2 sigma_m^2)
    + M_{t,i})`` with ``q_t = R(gamma g^f_t) q`` (one learned query, no frame content),
    ``k_i = R(gamma g^c_i)
    W_k h_i``, ``M = 0`` iff ``|w(i) − w(t)| <= window`` (and ``i`` is a real character),
    else ``−inf``. Output ``z_t = W_o concat_m Σ_i A^{(m)}_{t,i} W_v h_i`` ``(B, T, D)``.
    """

    def __init__(
        self,
        dim: int,
        heads: int,
        *,
        sigma_init: tuple[float, float] = (0.18, 1.2),
        gamma: float = 8.0,
        rope_base: float = 100.0,
        window: int = 1,
    ) -> None:
        super().__init__()
        if dim % heads or (dim // heads) % 2:
            raise ValueError("dim must split into heads of even size")
        if not 0.0 < sigma_init[0] <= sigma_init[1]:
            raise ValueError("sigma_init must be 0 < low <= high")
        self.heads = heads
        self.head_dim = dim // heads
        self.gamma = gamma
        self.rope_base = rope_base
        self.window = window
        self.query = nn.Parameter(torch.randn(heads, self.head_dim) / math.sqrt(self.head_dim))
        self.k_proj = nn.Linear(dim, dim)
        self.v_proj = nn.Linear(dim, dim)
        self.out = nn.Linear(dim, dim)
        sig = torch.linspace(math.log(sigma_init[0]), math.log(sigma_init[1]), heads)
        self.log_sigma = nn.Parameter(sig)
        self.log_beta = nn.Parameter(torch.zeros(heads))
        self.log_tau = nn.Parameter(torch.zeros(()))

    def forward(
        self,
        h: Tensor,
        char_mask: Tensor,
        char_unit: Tensor,
        char_coord: Tensor,
        frame_unit: Tensor,
        frame_coord: Tensor,
    ) -> Tensor:
        """``h`` ``(B, L, D)``; char tensors ``(B, L)``; frame tensors ``(B, T)`` → ``(B, T, D)``.

        Every frame keeps at least one admissible character when its unit has one; a
        frame whose window holds no real character (an empty item, or batch padding) gets
        a zero output instead of a NaN.
        """
        batch, length, dim = h.shape
        hd = self.head_dim
        k = self.k_proj(h).view(batch, length, self.heads, hd)
        v = self.v_proj(h).view(batch, length, self.heads, hd)
        k = _rotate(k, self.gamma * char_coord, self.rope_base)  # (B, L, H, hd) float
        frames = frame_coord.shape[1]
        q = self.query[None, None].expand(batch, frames, self.heads, hd)
        q = _rotate(q, self.gamma * frame_coord, self.rope_base)  # (B, T, H, hd)
        content = torch.einsum("bthd,blhd->bhtl", q, k) / math.sqrt(hd)
        dist2 = (frame_coord[:, None, :, None].float() - char_coord[:, None, None, :].float()) ** 2
        sigma2 = torch.exp(2 * self.log_sigma)[None, :, None, None]
        beta = torch.exp(self.log_beta)[None, :, None, None]
        logits = torch.exp(self.log_tau) * content - beta * dist2 / (2 * sigma2)
        near = (frame_unit[:, :, None] - char_unit[:, None, :]).abs() <= self.window
        allowed = (near & char_mask[:, None, :])[:, None]  # (B, 1, T, L)
        logits = logits.masked_fill(~allowed, float("-inf"))
        has_any = allowed.any(dim=-1, keepdim=True)  # (B, 1, T, 1)
        logits = torch.where(has_any, logits, torch.zeros((), device=h.device))
        attn = torch.softmax(logits, dim=-1) * has_any
        z = torch.einsum("bhtl,blhd->bthd", attn.to(v.dtype), v).reshape(batch, frames, dim)
        out: Tensor = self.out(z)
        return out