r"""Reverse diffusion samplers.

For a distribution :math:`p(X)` over :math:`\mathbb{R}^D`, the perturbation kernel (see
:mod:`azula.noise`)

.. math:: p(X_t \mid X) = \mathcal{N}(X_t \mid \alpha_t X, \sigma_t^2 I)

defines a series of marginal distributions

.. math:: p(X_t) = \int p(X_t \mid x) \, p(x) \, dx \, .

The goal of diffusion models is to generate samples from :math:`p(X_0)`. To this end,
reverse transition kernels :math:`q(X_s \mid X_t)` from :math:`t` to :math:`s < t` are
chosen. Then, starting from :math:`x_1 \sim p(X_1)`, :math:`T` transitions
:math:`x_{t_{i-1}} \sim q(X_{t_{i-1}} \mid x_{t_i})` are simulated from :math:`t_T = 1`
to :math:`t_0 = 0`. If the kernels are consistent with the marginals :math:`p(X_t)`,
that is if

.. math:: p(X_{t_{i-1}}) \approx
    \int q(X_{t_{i-1}} \mid x_{t_i}) \, p(x_{t_i}) \, dx_{t_i} \, ,

for all :math:`i = 1, \dots, T`, the tensors :math:`x_{t_i}` are distributed according
to :math:`p(X_{t_i})`, including the last one :math:`x_{t_0} = x_0`.
"""

__all__ = [
    "Sampler",
    "DDPMSampler",
    "DDIMSampler",
    "EulerSampler",
    "HeunSampler",
    "ItoSampler",
    "zABSampler",
    "vABSampler",
    "zEABSampler",
    "xEABSampler",
    "REABSampler",
    "PCSampler",
]

import abc
import math
import torch

from collections.abc import Iterable, Sequence
from functools import partial
from torch import Tensor
from tqdm import tqdm

from .denoise import Denoiser
from .nn.utils import promote_dtype


class Sampler(abc.ABC):
    r"""Abstract reverse diffusion sampler.

    Arguments:
        start: The starting time :math:`t_T`.
        stop: The stopping time :math:`t_0`.
        steps: The number of discretization steps :math:`T`. By default, the step size
            :math:`t_i - t_{i-1}` is constant.
        silent: Whether to hide the sampling progress bar or not.
        dtype: The time data type.
        device: The time device.
    """

    denoiser: Denoiser

    def __init__(
        self,
        start: float = 1.0,
        stop: float = 0.0,
        steps: int = 64,
        silent: bool = False,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
    ) -> None:
        self.start = start
        self.stop = stop
        self.steps = steps
        self.silent = silent

        self.dtype = dtype
        self.device = device

    @property
    def timesteps(self) -> Tensor:
        return torch.linspace(
            self.start,
            self.stop,
            self.steps + 1,
            dtype=self.dtype,
            device=self.device,
        )

    @torch.no_grad()
    def init(
        self,
        shape: Sequence[int],
        mean: float | Tensor = 0.0,
        var: float | Tensor = 1.0,
        **kwargs,
    ) -> Tensor:
        r"""Draws an initial noisy tensor :math:`x_{t_T}`.

        .. math:: x_{t_T} \sim \mathcal{N}(\alpha_{t_T} \mathbb{E}[X], \alpha_{t_T}^2 \mathbb{V}[X] + \sigma_{t_T}^2 I)

        Arguments:
            shape: The shape :math:`(*)` of the tensor.
            mean: The mean :math:`\mathbb{E}[X]` of :math:`p(X)`, with shape
                :math:`()` or :math:`(*)`.
            var: The variance :math:`\mathbb{V}[X]` of :math:`p(X)`, with shape
                :math:`()` or :math:`(*)`.
            kwargs: Keyword arguments passed to :func:`torch.Tensor.to`.

        Returns:
            A noisy tensor :math:`x_{t_T}`, with shape :math:`(*)`.
        """

        t_T = self.timesteps[0]

        alpha_T, sigma_T = self.denoiser.schedule(t_T)
        alpha_T, sigma_T = alpha_T.to(**kwargs), sigma_T.to(**kwargs)

        mean_T, std_T = alpha_T * mean, torch.sqrt(alpha_T**2 * var + sigma_T**2)
        mean_T, std_T = mean_T.expand(shape), std_T.expand(shape)

        return mean_T + std_T * torch.randn_like(mean_T)

    def progress_bar(self, it: Iterable) -> Iterable:
        if torch.is_tensor(it):
            it = it.unbind()

        if self.silent:
            return it
        else:
            return tqdm(it, miniters=1, unit="step", ncols=79, ascii=True)

    @torch.no_grad()
    def __call__(self, x: Tensor, **kwargs) -> Tensor:
        r"""Simulates the reverse process from :math:`t_T` to :math:`t_0`.

        Arguments:
            x: A noisy tensor :math:`x_{t_T}`, with shape :math:`(*)`.
            kwargs: Optional keyword arguments.

        Returns:
            The clean(er) tensor :math:`x_{t_0}`, with shape :math:`(*)`.
        """

        time_pairs = self.timesteps.unfold(0, 2, 1).to(device=x.device)

        x_t = x

        for t, s in self.progress_bar(time_pairs):
            x_s = self.step(x_t, t, s, **kwargs)
            x_t = x_s

        x = x_t

        return x

    def step(self, x_t: Tensor, t: Tensor, s: Tensor, **kwargs) -> Tensor:
        r"""Simulates the reverse process from :math:`t` to :math:`s`.

        Arguments:
            x_t: The current tensor :math:`x_t`, with shape :math:`(*)`.
            t: The current time :math:`t`, with shape :math:`()`.
            s: The target time :math:`s`, with shape :math:`()`.
            kwargs: Optional keyword arguments.

        Returns:
            The new tensor :math:`x_s \sim q(X_s \mid x_t)`, with shape :math:`(*)`.
        """

        raise NotImplementedError()


class DDPMSampler(Sampler):
    r"""Creates an DDPM sampler.

    .. math:: x_s \gets \alpha_s \mathbb{E}[X \mid x_t]
        + \sigma_s \, \sqrt{1 - \tau} \, \frac{x_t - \alpha_t \mathbb{E}[X \mid x_t]}{\sigma_t}
        + \sigma_s \, \sqrt{\tau} \, \varepsilon

    where :math:`\varepsilon \sim \mathcal{N}(0, I)` and

    .. math:: \tau = 1 - \frac{\alpha_t^2}{\alpha_s^2} \frac{\sigma_s^2}{\sigma_t^2} \, .

    References:
        | Denoising Diffusion Probabilistic Models (Ho et al., 2020)
        | https://arxiv.org/abs/2006.11239

    Arguments:
        denoiser: A denoiser :math:`q_\phi(X \mid X_t)`.
        kwargs: Keyword arguments passed to :class:`Sampler`.
    """

    def __init__(self, denoiser: Denoiser, **kwargs) -> None:
        super().__init__(**kwargs)

        self.denoiser = denoiser

    def step(self, x_t: Tensor, t: Tensor, s: Tensor, **kwargs) -> Tensor:
        alpha_s, sigma_s = self.denoiser.schedule(s)
        alpha_t, sigma_t = self.denoiser.schedule(t)

        tau = 1 - (alpha_t / alpha_s * sigma_s / sigma_t) ** 2

        q_t = self.denoiser(x_t, t, **kwargs)

        x_s = alpha_s * q_t.mean
        x_s = x_s + sigma_s * torch.sqrt(1 - tau) / sigma_t * (x_t - alpha_t * q_t.mean)
        x_s = x_s + sigma_s * torch.sqrt(tau) * torch.randn_like(x_t)

        return x_s


class DDIMSampler(Sampler):
    r"""Creates a DDIM sampler.

    .. math:: x_s \gets \alpha_s \mathbb{E}[X \mid x_t]
        + \sigma_s \, \sqrt{1 - \eta \, \tau} \, \frac{x_t - \alpha_t \mathbb{E}[X \mid x_t]}{\sigma_t}
        + \sigma_s \, \sqrt{\eta \, \tau} \, \varepsilon

    where :math:`\varepsilon \sim \mathcal{N}(0, I)` and

    .. math:: \tau = 1 - \frac{\alpha_t^2}{\alpha_s^2} \frac{\sigma_s^2}{\sigma_t^2} \, .

    References:
        | Denoising Diffusion Implicit Models (Song et al., 2021)
        | https://arxiv.org/abs/2010.02502

    Arguments:
        denoiser: A denoiser :math:`q_\phi(X \mid X_t)`.
        eta: The stochasticity parameter :math:`\eta \in \mathbb{R}_+`.
            If :math:`\eta = 1`, :class:`DDIMSampler` is equivalent to :class:`DDPMSampler`.
            If :math:`\eta = 0`, :class:`DDIMSampler` is equivalent to :class:`EulerSampler`.
        kwargs: Keyword arguments passed to :class:`Sampler`.
    """

    def __init__(self, denoiser: Denoiser, eta: float = 0.0, **kwargs) -> None:
        super().__init__(**kwargs)

        self.denoiser = denoiser
        self.eta = eta

    def step(self, x_t: Tensor, t: Tensor, s: Tensor, **kwargs) -> Tensor:
        alpha_s, sigma_s = self.denoiser.schedule(s)
        alpha_t, sigma_t = self.denoiser.schedule(t)

        tau = 1 - (alpha_t / alpha_s * sigma_s / sigma_t) ** 2
        tau = torch.clip(self.eta * tau, min=0, max=1)

        q_t = self.denoiser(x_t, t, **kwargs)

        x_s = alpha_s * q_t.mean
        x_s = x_s + sigma_s * torch.sqrt(1 - tau) / sigma_t * (x_t - alpha_t * q_t.mean)
        x_s = x_s + sigma_s * torch.sqrt(tau) * torch.randn_like(x_t)

        return x_s


class EulerSampler(Sampler):
    r"""Creates a explicit Euler (1st order) sampler.

    Without loss of generality, let's assume :math:`\alpha_t = 1` and :math:`\sigma_t =
    t` such that

    .. math:: \frac{d x_t}{dt}
        & = \alpha_t' \, \mathbb{E}[X \mid x_t] + \sigma_t' \, \mathbb{E}[Z \mid x_t] \\
        & = \mathbb{E}[Z \mid x_t] = z(x_t)

    and

    .. math:: x_s = x_t + \int_t^s z(x_u) \, du \, .

    The explicit Euler step for this integral is

    .. math:: x_s \gets x_t + (s - t) \, z(x_t)

    Wikipedia:
        https://wikipedia.org/wiki/Euler_method

    Arguments:
        denoiser: A denoiser :math:`q_\phi(X \mid X_t)`.
        kwargs: Keyword arguments passed to :class:`Sampler`.
    """

    def __init__(self, denoiser: Denoiser, **kwargs) -> None:
        super().__init__(**kwargs)

        self.denoiser = denoiser

    def step(self, x_t: Tensor, t: Tensor, s: Tensor, **kwargs) -> Tensor:
        alpha_s, sigma_s = self.denoiser.schedule(s)
        alpha_t, sigma_t = self.denoiser.schedule(t)

        q_t = self.denoiser(x_t, t, **kwargs)
        z_t = (x_t - alpha_t * q_t.mean) / sigma_t
        x_s = alpha_s / alpha_t * x_t + alpha_s * (sigma_s / alpha_s - sigma_t / alpha_t) * z_t

        return x_s


class HeunSampler(Sampler):
    r"""Creates a explicit Heun (2nd order) sampler.

    Without loss of generality, let's assume :math:`\alpha_t = 1` and :math:`\sigma_t =
    t` such that

    .. math:: \frac{d x_t}{dt}
        & = \alpha_t' \, \mathbb{E}[X \mid x_t] + \sigma_t' \, \mathbb{E}[Z \mid x_t] \\
        & = \mathbb{E}[Z \mid x_t] = z(x_t)

    and

    .. math:: x_s = x_t + \int_t^s z(x_u) \, du \, .

    The explicit Heun step for this integral is

    .. math::
        x_s & \gets x_t + (s - t) \, z(x_t) \\
        x_s & \gets x_t + (s - t) \frac{z(x_t) + z(x_s)}{2}

    Wikipedia:
        https://wikipedia.org/wiki/Heun%27s_method

    Arguments:
        denoiser: A denoiser :math:`q_\phi(X \mid X_t)`.
        kwargs: Keyword arguments passed to :class:`Sampler`.
    """

    def __init__(self, denoiser: Denoiser, **kwargs) -> None:
        super().__init__(**kwargs)

        self.denoiser = denoiser

    def step(self, x_t: Tensor, t: Tensor, s: Tensor, **kwargs) -> Tensor:
        alpha_s, sigma_s = self.denoiser.schedule(s)
        alpha_t, sigma_t = self.denoiser.schedule(t)

        q_t = self.denoiser(x_t, t, **kwargs)
        z_t = (x_t - alpha_t * q_t.mean) / sigma_t
        x_s = alpha_s / alpha_t * x_t + alpha_s * (sigma_s / alpha_s - sigma_t / alpha_t) * z_t

        q_s = self.denoiser(x_s, s, **kwargs)
        z_s = (x_s - alpha_s * q_s.mean) / sigma_s
        z_t = (z_t + z_s) / 2
        x_s = alpha_s / alpha_t * x_t + alpha_s * (sigma_s / alpha_s - sigma_t / alpha_t) * z_t

        return x_s


class ItoSampler(Sampler):
    r"""Creates an Itô SDE sampler.

    Let's consider the Itô SDE

    .. math:: dx_t = \left[ f_t \, x_t - \frac{1 + \eta^2}{2 \tau} g_t^2 \, \nabla_{x_t} \log p(x_t) \right] dt + \eta \, g_t \, dw_t

    where :math:`\eta \geq 0` controls stochasticity, :math:`\tau \geq 0` controls
    temperature, and

    .. math::
        f_t & = \partial_t \log \alpha_t \\
        g_t^2 & = \alpha_t^2 \, \partial_t \frac{\sigma_t^2}{\alpha_t^2} \, .

    The integral form of this semi-linear Itô SDE is

    .. math:: x_s = \Psi(t, s) \, x_t
        + \frac{1 + \eta^2}{2} \int_t^s \Psi(u, s) \, g_u^2 \, \nabla_{x_u} \log p(x_u) \, du
        + \eta \int_t^s \Psi(u, s) \, g_u \, dw_u

    where

    .. math:: \Psi(t, s) = \exp \left( \int_t^s f_u \, du \right) = \frac{\alpha_s}{\alpha_t} \, .

    By Tweedie's formula, we have

    .. math:: \int_t^s \Psi(u, s) \, g_u^2 \, \nabla_{x_u} \log p(x_u) \, du
        & = \int_t^s \Psi(u, s) \, g_u^2 \frac{\alpha_u \mathbb{E}[x_0 | x_u] - x_u}{\sigma_u^2} du \\
        & = 2 \alpha_s \int_t^s \partial_u \frac{\sigma_u}{\alpha_u} \frac{\alpha_u \mathbb{E}[x_0 | x_u] - x_u}{\sigma_u} du \\
        & \approx 2 \alpha_s \frac{\alpha_t \mathbb{E}[x_0 | x_t] - x_t}{\sigma_t} \int_t^s \partial_u \frac{\sigma_u}{\alpha_u} du \\
        & \approx 2 \left( \frac{\sigma_s}{\sigma_t} - \frac{\alpha_s}{\alpha_t} \right) (\alpha_t \mathbb{E}[x_0 | x_t] - x_t) \, .

    Finally, by Itô isometry, we have

    .. math:: \int_t^s \Psi(u, s) \, g_u \, dw_u
        & = \int_s^t \Psi(u, s) \, g_u \, dw_u \\
        & = \varepsilon \, \sqrt{\int_s^t \Psi(u, s)^2 \, g_u^2 \, du} \\
        & = \varepsilon \, \sqrt{\int_s^t \alpha_s^2 \, \partial_u \frac{\sigma_u^2}{\alpha_u^2} du} \\
        & = \varepsilon \, \alpha_s \sqrt{\frac{\sigma_t^2}{\alpha_t^2} - \frac{\sigma_s^2}{\alpha_s^2}}

    where :math:`\varepsilon \sim \mathcal{N}(0, 1)`.

    Arguments:
        denoiser: A denoiser :math:`q_\phi(X \mid X_t)`.
        eta: The stochasticity parameter :math:`\eta \geq 0`.
        temperature: The temperature parameter :math:`\tau \geq 0`.
        kwargs: Keyword arguments passed to :class:`Sampler`.
    """

    def __init__(
        self,
        denoiser: Denoiser,
        eta: float = 1.0,
        temperature: float = 1.0,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)

        self.denoiser = denoiser
        self.eta = eta
        self.temperature = temperature

    def step(self, x_t: Tensor, t: Tensor, s: Tensor, **kwargs) -> Tensor:
        alpha_s, sigma_s = self.denoiser.schedule(s)
        alpha_t, sigma_t = self.denoiser.schedule(t)

        q_t = self.denoiser(x_t, t, **kwargs)

        x_s = alpha_s / alpha_t * x_t
        x_s = x_s + (1 + self.eta**2) / self.temperature * (
            sigma_s / sigma_t - alpha_s / alpha_t
        ) * (x_t - alpha_t * q_t.mean)
        x_s = x_s + self.eta * alpha_s * torch.sqrt(
            torch.abs((sigma_t / alpha_t) ** 2 - (sigma_s / alpha_s) ** 2)
        ) * torch.randn_like(x_s)

        return x_s


class zABSampler(Sampler):
    r"""Creates an Adams-Bashforth (AB) multi-step sampler with noise (:math:`z`) prediction.

    Note:
        This sampler is equivalent to the :math:`\rho\mathrm{AB}` sampler from Zhang et al.
        (2023) and the linear multi-step (LMS) sampler from Katherine Crowson's
        `k-diffusion <https://github.com/crowsonkb/k-diffusion>`_.

    Without loss of generality, let's assume :math:`\alpha_t = 1` and :math:`\sigma_t =
    t` such that

    .. math:: \frac{d x_t}{dt}
        & = \alpha_t' \, \mathbb{E}[X \mid x_t] + \sigma_t' \, \mathbb{E}[Z \mid x_t] \\
        & = \mathbb{E}[Z \mid x_t] = z(x_t)

    and

    .. math:: x_s = x_t + \int_t^s z(x_u) \, du \, .

    The :math:`n`-th order Adams-Bashforth step for this integral is

    .. math:: x_s \gets x_t + \sum_{i=1}^{n} z(x_{t_i}) \int_t^s \ell_i(u) \, du

    where :math:`t_i` are previous time steps and the polynomials :math:`\ell_i(u) =
    \sum_{k=1}^{n} a_{ik} \, u^k` form their Lagrange basis. The coefficients
    :math:`a_{ik}` are determined by solving the system of linear equations
    :math:`\ell_i(t_j) = \delta_{ij}`. Then, the Adams-Bashforth coefficients are

    .. math:: \int_t^s \ell_i(u) \, du
        & = \sum_{k=1}^{n} a_{ik} \int_t^s u^k \, du \\
        & = \sum_{k=1}^{n} a_{ik} \left[ \frac{u^{k+1}}{k+1} \right]_t^s

    Wikipedia:
        https://wikipedia.org/wiki/Linear_multistep_method

    References:
        | Fast Sampling of Diffusion Models with Exponential Integrator (Zhang et al., 2023)
        | https://arxiv.org/abs/2204.13902

    Arguments:
        denoiser: A denoiser :math:`q_\phi(X \mid X_t)`.
        order: The order :math:`n` of the multi-step method.
        kwargs: Keyword arguments passed to :class:`Sampler`.
    """

    def __init__(self, denoiser: Denoiser, order: int = 2, **kwargs) -> None:
        super().__init__(**kwargs)

        self.denoiser = denoiser
        self.order = order

    @staticmethod
    @partial(promote_dtype, min_dtype=torch.float64)
    def _adams_bashforth(t: Tensor, /, i: int, n: int) -> Tensor:
        r"""Returns the coefficients of the :math:`n`-th order Adams-Bashforth method.

        Arguments:
            t: The integration variable, with shape :math:`(T)`.
            i: The current step :math:`i`.
            n: The method order :math:`n \leq i + 1`.

        Returns:
            The coefficients, with shape :math:`(n)`.
        """

        n = min(n, i + 1)
        k = torch.arange(n, device=t.device)

        # Vandermonde matrix t_i^k
        V = t[i + 1 - n : i + 1] ** k[:, None]

        # Integral of u^k from t_i to t_{i+1}
        b = t[i + 1] ** (k + 1) / (k + 1) - t[i] ** (k + 1) / (k + 1)

        return torch.linalg.solve(V, b)

    @torch.no_grad()
    def __call__(self, x: Tensor, **kwargs) -> Tensor:
        time = self.timesteps.to(device=x.device)
        alpha, sigma = self.denoiser.schedule(time)
        u = sigma / alpha

        x_t = x

        buffer = []

        for i, t in enumerate(self.progress_bar(time[:-1])):
            alpha_t, sigma_t = alpha[i], sigma[i]
            alpha_s = alpha[i + 1]

            q_t = self.denoiser(x_t, t, **kwargs)
            z_t = (x_t - alpha_t * q_t.mean) / sigma_t

            buffer.append(z_t)
            if len(buffer) > self.order:
                buffer.pop(0)

            coeffs = self._adams_bashforth(u, i=i, n=self.order)
            integral = sum(b * c for b, c in zip(buffer, coeffs, strict=True))

            x_s = alpha_s / alpha_t * x_t + alpha_s * integral
            x_t = x_s

        x = x_t

        return x


class vABSampler(zABSampler):
    r"""Creates an Adams-Bashforth (AB) multi-step sampler with velocity (:math:`v`) prediction.

    Without loss of generality, let's assume :math:`\alpha_t = 1 - t` and
    :math:`\sigma_t = t` such that

    .. math:: \frac{d x_t}{dt}
        & = \alpha_t' \, \mathbb{E}[X \mid x_t] + \sigma_t' \, \mathbb{E}[Z \mid x_t] \\
        & = -\mathbb{E}[X \mid x_t] + \mathbb{E}[Z \mid x_t] = v(x_t)

    and

    .. math:: x_s = x_t + \int_t^s v(x_u) \, du \, .

    The :math:`n`-th order Adams-Bashforth step for this integral is

    .. math:: x_s \gets x_t + \sum_{i=1}^{n} v(x_{t_i}) \int_t^s \ell_i(u) \, du

    where :math:`t_i` are previous time steps and the polynomials :math:`\ell_i(u)` form their Lagrange basis.

    See also:
        :class:`zABSampler`

    Arguments:
        denoiser: A denoiser :math:`q_\phi(X \mid X_t)`.
        order: The order :math:`n` of the multi-step method.
        kwargs: Keyword arguments passed to :class:`Sampler`.
    """

    @torch.no_grad()
    def __call__(self, x: Tensor, **kwargs) -> Tensor:
        time = self.timesteps.to(device=x.device)
        alpha, sigma = self.denoiser.schedule(time)
        u = sigma / (alpha + sigma)

        x_t = x

        buffer = []

        for i, t in enumerate(self.progress_bar(time[:-1])):
            alpha_t, sigma_t = alpha[i], sigma[i]
            alpha_s, sigma_s = alpha[i + 1], sigma[i + 1]

            q_t = self.denoiser(x_t, t, **kwargs)
            v_t = 1 / sigma_t * x_t - (1 + alpha_t / sigma_t) * q_t.mean

            buffer.append(v_t)
            if len(buffer) > self.order:
                buffer.pop(0)

            coeffs = self._adams_bashforth(u, i=i, n=self.order)
            integral = sum(b * c for b, c in zip(buffer, coeffs, strict=True))

            x_s = (alpha_s + sigma_s) / (alpha_t + sigma_t) * x_t + (alpha_s + sigma_s) * integral
            x_t = x_s

        x = x_t

        return x


class zEABSampler(Sampler):
    r"""Creates an exponential Adams-Bashforth (EAB) multi-step sampler with noise (:math:`z`) prediction.

    Note:
        This sampler is a multi-step generalization of the DPM-Solver sampler from Lu et
        al. (2022).

    Without loss of generality, let's assume :math:`\alpha_t = 1` and :math:`\sigma_t =
    e^t` such that

    .. math:: \frac{d x_t}{dt}
        & = \alpha_t' \, \mathbb{E}[X \mid x_t] + \sigma_t' \, \mathbb{E}[Z \mid x_t] \\
        & = e^t \, \mathbb{E}[Z \mid x_t] = e^t \, z(x_t)

    and

    .. math:: x_s = x_t + \int_t^s e^u \, z(x_u) \, du \, .

    The :math:`n`-th order exponential Adams-Bashforth step for this integral is

    .. math:: x_s \gets x_t + \sum_{i=1}^{n} z(x_{t_i}) \int_t^s e^u \, \ell_i(u) \, du

    where :math:`t_i` are previous time steps and the polynomials :math:`\ell_i(u) =
    \sum_{k=1}^{n} a_{ik} \, u^k` form their Lagrange basis. The coefficients
    :math:`a_{ik}` are determined by solving the system of linear equations
    :math:`\ell_i(t_j) = \delta_{ij}`. Then, the exponential Adams-Bashforth
    coefficients are

    .. math:: \int_t^s e^u \, \ell_i(u) \, du
        & = \sum_{k=1}^{n} a_{ik} \int_t^s e^u \, u^k \, du \\
        & = \sum_{k=1}^{n} a_{ik} \left[ (-1)^k \, k! \, e^u \sum_{j=0}^k \frac{(-u)^j}{j!} \right]_t^s

    References:
        | DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps (Lu et al., 2022)
        | https://arxiv.org/abs/2206.00927

        | Exponential Adams Bashforth ODE solver for stiff problems (Coudière et al., 2018)
        | https://arxiv.org/abs/1804.09927

    Arguments:
        denoiser: A denoiser :math:`q_\phi(X \mid X_t)`.
        order: The order :math:`n` of the multi-step method.
        kwargs: Keyword arguments passed to :class:`Sampler`.
    """

    def __init__(self, denoiser: Denoiser, order: int = 2, **kwargs) -> None:
        super().__init__(**kwargs)

        self.denoiser = denoiser
        self.order = order

    @staticmethod
    @partial(promote_dtype, min_dtype=torch.float64)
    def _exponential_adams_bashforth(t: Tensor, /, i: int, n: int) -> Tensor:
        r"""Returns the coefficients of the :math:`n`-th order exponential Adams-Bashforth method.

        Arguments:
            t: The integration variable, with shape :math:`(T)`.
            i: The current step :math:`i`.
            n: The method order :math:`n \leq i + 1`.

        Returns:
            The coefficients, with shape :math:`(n)`.
        """

        n = min(n, i + 1)
        k = torch.arange(n, device=t.device)
        k_fact = torch.cumprod(torch.clip(k, min=1), dim=0)

        # Vandermonde matrix t_i^k
        V = t[i + 1 - n : i + 1] ** k[:, None]

        # Integral of exp(u) u^k from t_i to t_{i+1}
        b = (
            (-1) ** k
            * k_fact
            * (
                torch.exp(t[i + 1]) * torch.cumsum((-t[i + 1]) ** k / k_fact, dim=0)
                - torch.exp(t[i]) * torch.cumsum((-t[i]) ** k / k_fact, dim=0)
            )
        )

        return torch.linalg.solve(V, b)

    @torch.no_grad()
    def __call__(self, x: Tensor, **kwargs) -> Tensor:
        time = self.timesteps.to(device=x.device)
        alpha, sigma = self.denoiser.schedule(time)
        u = sigma.log() - alpha.log()

        x_t = x

        buffer = []

        for i, t in enumerate(self.progress_bar(time[:-1])):
            alpha_t, sigma_t = alpha[i], sigma[i]
            alpha_s = alpha[i + 1]

            q_t = self.denoiser(x_t, t, **kwargs)
            z_t = (x_t - alpha_t * q_t.mean) / sigma_t

            buffer.append(z_t)
            if len(buffer) > self.order:
                buffer.pop(0)

            coeffs = self._exponential_adams_bashforth(u, i=i, n=self.order)
            integral = sum(b * c for b, c in zip(buffer, coeffs, strict=True))

            x_s = alpha_s / alpha_t * x_t + alpha_s * integral
            x_t = x_s

        x = x_t

        return x


class xEABSampler(Sampler):
    r"""Creates an exponential Adams-Bashforth (EAB) multi-step sampler with data (:math:`x`) prediction.

    Note:
        This sampler is a multi-step generalization of the DPM-Solver++ sampler from Lu et al. (2022).

    Without loss of generality, let's assume :math:`\alpha_t = 1` and :math:`\sigma_t =
    e^t` such that

    .. math:: \frac{d x_t}{dt}
        & = \alpha_t' \, \mathbb{E}[X \mid x_t] + \sigma_t' \, \mathbb{E}[Z \mid x_t] \\
        & = e^t \, \mathbb{E}[Z \mid x_t] = x_t - \mathbb{E}[X \mid x_t] = x_t - x(x_t)

    and

    .. math:: x_s = e^{s-t} \, x_t - \int_t^s e^{s-u} \, x(x_u) \, du \, .

    The :math:`n`-th order exponential Adams-Bashforth step for this integral is

    .. math:: x_s \gets e^{s-t} \, x_t
        - e^s \sum_{i=1}^{n} x(x_{t_i}) \int_t^s e^{-u} \, \ell_i(u) \, du

    where :math:`t_i` are previous time steps and the polynomials :math:`\ell_i(u) =
    \sum_{k=1}^{n} a_{ik} \, u^k` form their Lagrange basis. The coefficients
    :math:`a_{ik}` are determined by solving the system of linear equations
    :math:`\ell_i(t_j) = \delta_{ij}`. Then, the exponential Adams-Bashforth
    coefficients are

    .. math:: \int_t^s e^{-u} \, \ell_i(u) \, du
        & = \sum_{k=1}^{n} a_{ik} \int_t^s e^{-u} \, u^k \, du \\
        & = \sum_{k=1}^{n} a_{ik} \left[ -k! \, e^{-u} \sum_{j=0}^k \frac{u^j}{j!} \right]_t^s

    References:
        | DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic Models (Lu et al., 2022)
        | https://arxiv.org/abs/2211.01095

    Arguments:
        denoiser: A denoiser :math:`q_\phi(X \mid X_t)`.
        order: The order :math:`n` of the multi-step method.
        kwargs: Keyword arguments passed to :class:`Sampler`.
    """

    def __init__(self, denoiser: Denoiser, order: int = 2, **kwargs) -> None:
        super().__init__(**kwargs)

        self.denoiser = denoiser
        self.order = order

    @staticmethod
    @partial(promote_dtype, min_dtype=torch.float64)
    def _exponential_adams_bashforth(t: Tensor, /, i: int, n: int) -> Tensor:
        r"""Returns the coefficients of the :math:`n`-th order exponential Adams-Bashforth method.

        Arguments:
            t: The integration variable, with shape :math:`(T)`.
            i: The current step :math:`i`.
            n: The method order :math:`n \leq i + 1`.

        Returns:
            The coefficients, with shape :math:`(n)`.
        """

        n = min(n, i + 1)
        k = torch.arange(n, device=t.device)
        k_fact = torch.cumprod(torch.clip(k, min=1), dim=0)

        # Vandermonde matrix t_i^k
        V = t[i + 1 - n : i + 1] ** k[:, None]

        # Integral of exp(-u) u^k from t_i to t_{i+1}
        b = -k_fact * (
            torch.exp(-t[i + 1]) * torch.cumsum(t[i + 1] ** k / k_fact, dim=0)
            - torch.exp(-t[i]) * torch.cumsum(t[i] ** k / k_fact, dim=0)
        )

        return torch.linalg.solve(V, b)

    @torch.no_grad()
    def __call__(self, x: Tensor, **kwargs) -> Tensor:
        time = self.timesteps.to(device=x.device)
        alpha, sigma = self.denoiser.schedule(time)
        u = sigma.log() - alpha.log()

        x_t = x

        buffer = []

        for i, t in enumerate(self.progress_bar(time[:-1])):
            sigma_t = sigma[i]
            sigma_s = sigma[i + 1]

            q_t = self.denoiser(x_t, t, **kwargs)

            buffer.append(q_t.mean)
            if len(buffer) > self.order:
                buffer.pop(0)

            coeffs = self._exponential_adams_bashforth(u, i=i, n=self.order)
            integral = sum(b * c for b, c in zip(buffer, coeffs, strict=True))

            x_s = sigma_s / sigma_t * x_t - sigma_s * integral
            x_t = x_s

        x = x_t

        return x


class REABSampler(Sampler):
    r"""Creates a Rosenbrock-type exponential Adams-Bashforth (REAB) multi-step sampler.

    Note:
        This sampler is a multi-step generalization of the DPM-Solver-v3 sampler from Zheng et al. (2023).

    Without loss of generality, let's assume :math:`\alpha_t = 1` and :math:`\sigma_t =
    e^t` such that

    .. math:: \frac{d x_t}{dt}
        & = \alpha_t' \, \mathbb{E}[X \mid x_t] + \sigma_t' \, \mathbb{E}[Z \mid x_t] \\
        & = e^t \, \mathbb{E}[Z \mid x_t] = x_t - \mathbb{E}[X \mid x_t] \\
        & = a_t \, x_t + b_t \frac{(1 - a_t) \, x_t - \mathbb{E}[X \mid x_t]}{b_t}
            = a_t \, x_t + b_t \, f(x_t)

    and

    .. math:: x_s = \Psi(t, s) \, x_t + \int_t^s \Psi(u, s) \, b_u \, f(x_u) \, du

    where :math:`a_t = \frac{e^{2t}}{1 + e^{2t}}`, :math:`b_t = \sqrt{a_t}` and

    .. math:: \Psi(t, s)
        = \exp \left( \int_t^s a_u \, du \right)
        = \sqrt{\frac{1 + e^{2s}}{1 + e^{2t}}} \, .

    The :math:`n`-th order exponential Adams-Bashforth step for this integral is

    .. math:: x_s \gets \sqrt{\frac{1 + e^{2s}}{1 + e^{2t}}} \, x_t
        + \sqrt{1 + e^{2s}} \sum_{i=1}^{n} f(x_{t_i}) \int_t^s \frac{e^u}{1 + e^{2u}} \ell_i(u) \, du

    where :math:`t_i` are previous time steps and the polynomials :math:`\ell_i(u) =
    \sum_{k=1}^{n} a_{ik} \, u^k` form their Lagrange basis. The coefficients
    :math:`a_{ik}` are determined by solving the system of linear equations
    :math:`\ell_i(t_j) = \delta_{ij}`. Then, the exponential Adams-Bashforth
    coefficients are

    .. math:: \int_t^s \frac{e^u}{1 + e^{2u}} \ell_i(u) \, du
        = \sum_{k=1}^{n} a_{ik} \int_t^s \frac{e^u \, u^k}{1 + e^{2u}} du

    where the last integral is estimated by numerical integration.

    References:
        | DPM-Solver-v3: Improved Diffusion ODE Solver with Empirical Model Statistics (Zheng et al., 2023)
        | https://arxiv.org/abs/2310.13268

        | Exponential Rosenbrock-Type Methods (Hochbruck et al., 2009)
        | https://epubs.siam.org/doi/10.1137/080717717

    Arguments:
        denoiser: A denoiser :math:`q_\phi(X \mid X_t)`.
        order: The order :math:`n` of the multi-step method.
        kwargs: Keyword arguments passed to :class:`Sampler`.
    """

    def __init__(self, denoiser: Denoiser, order: int = 2, **kwargs) -> None:
        super().__init__(**kwargs)

        self.denoiser = denoiser
        self.order = order

    @staticmethod
    @partial(promote_dtype, min_dtype=torch.float64)
    def _exponential_adams_bashforth(t: Tensor, /, i: int, n: int) -> Tensor:
        r"""Returns the coefficients of the :math:`n`-th order exponential Adams-Bashforth method.

        Arguments:
            t: The integration variable, with shape :math:`(T)`.
            i: The current step :math:`i`.
            n: The method order :math:`n \leq i + 1`.

        Returns:
            The coefficients, with shape :math:`(n)`.
        """

        n = min(n, i + 1)
        k = torch.arange(n, device=t.device)

        # Vandermonde matrix t_i^k
        V = t[i + 1 - n : i + 1] ** k[:, None]

        # Integral of exp(u) / (1 + exp(2u)) u^k from t_i to t_{i+1}
        u = torch.linspace(t[i], t[i + 1], steps=256 + 1, dtype=t.dtype, device=t.device)
        y = torch.exp(u) / (1 + torch.exp(2 * u)) * (u ** k[:, None])
        b = torch.trapezoid(y, u, dim=-1)

        return torch.linalg.solve(V, b)

    @torch.no_grad()
    def __call__(self, x: Tensor, **kwargs) -> Tensor:
        time = self.timesteps.to(device=x.device)
        alpha, sigma = self.denoiser.schedule(time)
        u = sigma.log() - alpha.log()

        x_t = x

        buffer = []

        for i, t in enumerate(self.progress_bar(time[:-1])):
            alpha_t, sigma_t = alpha[i], sigma[i]
            alpha_s, sigma_s = alpha[i + 1], sigma[i + 1]

            q_t = self.denoiser(x_t, t, **kwargs)

            a_t = sigma_t**2 / (alpha_t**2 + sigma_t**2)
            b_t = sigma_t * torch.rsqrt(alpha_t**2 + sigma_t**2)

            f_t = (1 - a_t) / b_t / alpha_t * x_t - 1 / b_t * q_t.mean

            buffer.append(f_t)
            if len(buffer) > self.order:
                buffer.pop(0)

            coeffs = self._exponential_adams_bashforth(u, i=i, n=self.order)
            integral = sum(b * c for b, c in zip(buffer, coeffs, strict=True))

            x_s = (
                torch.sqrt((alpha_s**2 + sigma_s**2) / (alpha_t**2 + sigma_t**2)) * x_t
                + torch.sqrt(alpha_s**2 + sigma_t**2) * integral
            )
            x_t = x_s

        x = x_t

        return x


class PCSampler(Sampler):
    r"""Creates a predictor-corrector (PC) sampler.

    Arguments:
        denoiser: A denoiser :math:`q_\phi(X \mid X_t)`.
        corrections: The number of corrector steps for each predictor step.
        delta: The amplitude of corrector steps :math:`\delta \in [0,1]`.
        kwargs: Keyword arguments passed to :class:`Sampler`.
    """

    def __init__(
        self,
        denoiser: Denoiser,
        corrections: int = 1,
        delta: float = 0.01,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)

        self.denoiser = denoiser
        self.corrections = corrections
        self.delta = delta

    def step(self, x_t: Tensor, t: Tensor, s: Tensor, **kwargs) -> Tensor:
        alpha_s, sigma_s = self.denoiser.schedule(s)
        alpha_t, sigma_t = self.denoiser.schedule(t)

        # Corrector
        for _ in range(self.corrections):
            q_t = self.denoiser(x_t, t, **kwargs)
            x_t = (
                alpha_t * q_t.mean
                + math.sqrt(1 - self.delta) * (x_t - alpha_t * q_t.mean)
                + math.sqrt(self.delta) * sigma_t * torch.randn_like(x_t)
            )

        # Predictor
        q_t = self.denoiser(x_t, t, **kwargs)
        x_s = alpha_s * q_t.mean + sigma_s / sigma_t * (x_t - alpha_t * q_t.mean)

        return x_s
