"""
Discretization methods for continuous-time systems.
"""

from typing import Callable, Optional, Tuple, Union

import torch
from torch import Tensor


def zoh(
    A: Tensor, delta: Tensor, integration_timesteps: Optional[Tensor] = None
) -> tuple[Tensor, Tensor]:
    """
    Zero-Order Hold (ZOH) discretization method, used across most models,
    $$
        A_bar = exp{delta * A}
        gamma_bar = A^{-1} (A_bar - I)
    $$
    Reference: https://hazyresearch.stanford.edu/blog/2022-01-14-s4-3

    Args
    ----
        A (Tensor): Continuous-time system matrix, shape: (N,), i.e., only diagonal elements.
        delta (Tensor): Time step for discretization.
        integration_timesteps (Tensor, optional): Not used in ZOH discretization.

    Returns
    -------
        tuple[Tensor, Tensor]: Discretized system, A_bar and input_normalizer, gamma_bar.
    """

    Identity = torch.ones(A.shape[0], device=A.device)
    A_bar = torch.exp(delta * A)
    gamma_bar = (1 / A) * (A_bar - Identity)
    return A_bar, gamma_bar


def bilinear(
    A: Tensor, delta: Tensor, integration_timesteps: Optional[Tensor] = None
) -> tuple[Tensor, Tensor]:
    """
    Bilinear method first used in S4,
    $$
        A_bar = (I + 0.5 * delta * A)^{-1} (I - 0.5 * delta * A)
        gamma_bar = (I + 0.5 * delta * A)^{-1} delta
    $$
    Reference: https://hazyresearch.stanford.edu/blog/2022-01-14-s4-3

    Args
    ----
        A (Tensor): Continuous-time system matrix, shape: (N,), i.e., only diagonal elements.
        delta (Tensor): Time step for discretization.
        integration_timesteps (Tensor, optional): Not used in bilinear discretization.

    Returns
    -------
        tuple[Tensor, Tensor]: Discretized system, A_bar and input_normalizer, gamma_bar.
    """

    Identity = torch.ones(A.shape[0], device=A.device)
    A_bar = (1 / (Identity + 0.5 * delta * A)) * (Identity - 0.5 * delta * A)
    gamma_bar = (1 / (Identity + 0.5 * delta * A)) * delta
    return A_bar, gamma_bar


def dirac(
    A: Tensor, delta: Tensor, integration_timesteps: Optional[Tensor] = None
) -> tuple[Tensor, float]:
    """
    Dirac discretization method,
    $$
        A_bar = exp{delta * A}
        gamma_bar = 1.0
    $$
    Reference: https://github.com/Efficient-Scalable-Machine-Learning/event-ssm/blob/main/event_ssm/ssm.py

    Args
    ----
        A (Tensor): Continuous-time system matrix.
        delta (Tensor): Time step for discretization.
        integration_timesteps (Tensor, optional): Not used in dirac discretization.

    Returns
    -------
        tuple[Tensor, Tensor]: Discretized system, A_bar and input_normalizer, gamma_bar.
    """

    A_bar = torch.exp(delta * A)
    gamma_bar = 1.0
    return A_bar, gamma_bar


def async_(
    A: Tensor, delta: Tensor, integration_timesteps: Optional[Tensor] = None
) -> tuple[Tensor, Tensor]:
    """
    Asynchronous discretization method, introduced in https://arxiv.org/abs/2404.18508,
    this helps provide a strong inductive bias for asynchronous event-streams,
    $$
        A_bar = exp{delta * integration_timesteps * A}
        gamma_bar = A^{-1} (exp{delta * A} - I)
    $$
    This method is only for legacy reasons, it is not possible to use this method (or any other
    discretization with async timesteps) with LTI models.

    Args
    ----
        A (Tensor): Continuous-time system matrix.
        delta (Tensor): Time step for discretization.
        integration_timesteps (Tensor): Timesteps for async discretization, ideally of shape (B, L), i.e., difference in timesteps of events.

    Returns
    -------
        tuple[Tensor, Tensor]: Discretized system, A_bar and input_normalizer, gamma_bar.
    """
    assert (
        integration_timesteps is not None
    ), "Integration timesteps must be provided for async discretization."
    Identity = torch.ones(A.shape[0], device=A.device)
    A_bar = torch.exp(delta * integration_timesteps * A)
    gamma_bar = (1 / A) * (A_bar - Identity)

    return A_bar, gamma_bar


def no_discretization(
    A: Tensor, delta: Tensor, integration_timesteps: Optional[Tensor] = None
) -> tuple[Tensor, float]:
    """
    No discretization method, identity operation.
    $$
        A_bar = A
        gamma_bar = 1.0
    $$

    Args
    ----
        A (Tensor): Continuous-time system matrix.
        delta (Tensor): Time step for discretization (unused).
        integration_timesteps (Tensor, optional): Not used in no_discretization.

    Returns
    -------
        tuple[Tensor, float]:
                - A_bar (Tensor): Same as A.
                - gamma_bar (float): 1.0, as B_bar = B.
    """

    return A, 1.0


DISCRETIZE_FNS: dict[
    str,
    Callable[
        [Tensor, Tensor, Optional[Tensor]], Tuple[Tensor, Union[Tensor, float]]
    ],
] = {
    "zoh": zoh,
    "bilinear": bilinear,
    "dirac": dirac,
    "async": async_,
    "no_discretization": no_discretization,
}
