"""Flow-matching convention — the ONLY place it is defined (docs/04-architecture.md).

* ``x0 ~ N(0, I)`` is noise, ``x1`` is the data latent, ``t ∈ [0, 1]``;
  **t = 0 is noise, t = 1 is data**.
* ``x_t = (1 − t) · x0 + t · x1``; the target velocity is ``v = x1 − x0``.
* Loss: masked MSE ``‖v̂(x_t, t, c) − v‖²`` over the frames to generate (``loss_mask``).
* Training ``t``: logit-normal by default (``sigmoid(N(mean, std))``), or uniform.
* Sampling integrates ``dx/dt = v̂`` from ``t = 0`` to ``t = 1`` (Euler or midpoint), on a
  time grid optionally warped by "sway" (``t + s · (cos(π t / 2) − 1 + t)``, F5-TTS).

Shapes: latents are ``[batch, frames, channels]``; ``t`` is ``[batch]``; masks are
``[batch, frames]`` booleans (True = counted).
"""

from __future__ import annotations

import math
from collections.abc import Callable
from typing import Literal

import torch
from torch import Tensor

__all__ = [
    "VelocityFn",
    "interpolate",
    "masked_mse",
    "sample_ode",
    "sample_t",
    "target_velocity",
    "time_grid",
]

VelocityFn = Callable[[Tensor, Tensor], Tensor]  # (x_t [B,T,C], t [B]) -> v̂ [B,T,C]


def _check_t(t: Tensor, batch: int) -> None:
    if t.ndim != 1 or t.shape[0] != batch:
        raise ValueError(f"t must have shape [{batch}], got {tuple(t.shape)}")


def interpolate(x0: Tensor, x1: Tensor, t: Tensor) -> Tensor:
    """``x_t = (1 − t) · x0 + t · x1`` with ``t`` broadcast over frames and channels."""
    if x0.shape != x1.shape or x0.ndim != 3:
        raise ValueError(f"x0 and x1 must be equal [B,T,C], got {x0.shape} and {x1.shape}")
    _check_t(t, x0.shape[0])
    tt = t.to(x0.dtype)[:, None, None]
    return (1.0 - tt) * x0 + tt * x1


def target_velocity(x0: Tensor, x1: Tensor) -> Tensor:
    """``v = x1 − x0`` (constant along the straight path)."""
    if x0.shape != x1.shape:
        raise ValueError(f"shape mismatch {x0.shape} vs {x1.shape}")
    return x1 - x0


def masked_mse(pred: Tensor, target: Tensor, loss_mask: Tensor) -> Tensor:
    """Mean squared error over the masked frames (and all channels).

    Raises:
        ValueError: shape mismatch, or a mask that selects no frame at all (the loss would
            be 0/0 — a silent NaN or a silent zero are both wrong).
    """
    if pred.shape != target.shape or pred.ndim != 3:
        raise ValueError(f"pred/target must be equal [B,T,C], got {pred.shape}, {target.shape}")
    if loss_mask.shape != pred.shape[:2] or loss_mask.dtype != torch.bool:
        raise ValueError(f"loss_mask must be bool [B,T], got {loss_mask.dtype} {loss_mask.shape}")
    count = loss_mask.sum()
    if int(count) == 0:
        raise ValueError("loss_mask selects no frame")
    # Select BEFORE squaring: NaN/inf on masked frames must reach neither the loss nor its
    # gradient (where() after the square would still back-propagate 2 * NaN * 0 = NaN).
    diff = torch.where(
        loss_mask[..., None], pred.float() - target.float(), torch.zeros((), device=pred.device)
    )
    return diff.pow(2).mean(dim=-1).sum() / count


def sample_t(
    batch: int,
    *,
    kind: Literal["logit_normal", "uniform"],
    mean: float = 0.0,
    std: float = 1.0,
    generator: torch.Generator | None = None,
    device: torch.device | None = None,
) -> Tensor:
    """Training time steps in ``(0, 1)``, shape ``[batch]``."""
    if batch < 1:
        raise ValueError("batch must be >= 1")
    if kind == "logit_normal":
        if not (math.isfinite(std) and std > 0 and math.isfinite(mean)):
            raise ValueError(f"invalid logit-normal parameters mean={mean} std={std}")
        normal = torch.randn(batch, generator=generator, device=device) * std + mean
        return torch.sigmoid(normal)
    if kind == "uniform":
        return torch.rand(batch, generator=generator, device=device)
    raise ValueError(f"unknown t sampling {kind!r}")


def time_grid(steps: int, *, sway: float = 0.0, device: torch.device | None = None) -> Tensor:
    """``steps + 1`` increasing times from 0 to 1, optionally sway-warped.

    Sway (F5-TTS): ``t' = t + s · (cos(π t / 2) − 1 + t)``; ``s < 0`` puts more steps near
    ``t = 0`` (noise end). Endpoints stay exactly 0 and 1; the grid must stay increasing.

    Raises:
        ValueError: ``steps < 1`` or a sway that makes the grid non-monotonic.
    """
    if steps < 1:
        raise ValueError("steps must be >= 1")
    t = torch.linspace(0.0, 1.0, steps + 1, dtype=torch.float64, device=device)
    if sway != 0.0:
        t = t + sway * (torch.cos(math.pi / 2 * t) - 1 + t)
        t[0], t[-1] = 0.0, 1.0
    if not bool(torch.all(t[1:] > t[:-1])):
        raise ValueError(f"sway {sway} gives a non-increasing time grid")
    return t


@torch.no_grad()
def sample_ode(
    velocity: VelocityFn,
    x0: Tensor,
    *,
    steps: int,
    method: Literal["euler", "midpoint"] = "euler",
    sway: float = 0.0,
) -> Tensor:
    """Integrate from noise ``x0`` (t = 0) to data (t = 1).

    With the oracle ``v̂ = x1 − x0`` both methods return ``x1`` exactly (up to float error):
    the path is a straight line (unit test of the convention).
    """
    if x0.ndim != 3:
        raise ValueError(f"x0 must be [B,T,C], got {tuple(x0.shape)}")
    if method not in ("euler", "midpoint"):
        raise ValueError(f"unknown ODE method {method!r}")
    grid = time_grid(steps, sway=sway, device=x0.device)
    x = x0
    batch = x0.shape[0]
    for i in range(steps):
        t0, t1 = float(grid[i]), float(grid[i + 1])
        dt = t1 - t0
        t_now = torch.full((batch,), t0, dtype=x0.dtype, device=x0.device)
        v = velocity(x, t_now)
        if method == "euler":
            x = x + dt * v
        elif method == "midpoint":
            t_mid = torch.full((batch,), t0 + dt / 2, dtype=x0.dtype, device=x0.device)
            x = x + dt * velocity(x + (dt / 2) * v, t_mid)
        else:
            raise ValueError(f"unknown ODE method {method!r}")
    return x