from abc import ABC, abstractmethod
from typing import Final, NamedTuple, Protocol, TypeAlias

import torch
from torch import nn
from xxhash import xxh3_64_intdigest

# Type Aliases
Trajectory: TypeAlias = list[tuple[float, str]]


# Protocols
class IndividualEffectsFn(Protocol):
    """The individual effects function protocol.

    Calls the individual effects function.

    It must be able to yield 2D or 3D tensors given inputs of `gamma` (population
    parameters), `x` (covariates matrix), and `b` (random effects). Note `b` is
    either 2D or 3D.

    Args:
        gamma (torch.Tensor | None): The population parameters.
        x (torch.Tensor | None): The fixed covariates matrix.
        b (torch.Tensor): The random effects.

    Returns:
        torch.Tensor: The individual parameters `psi`.

    Examples:
        >>> individual_effects_fn = lambda gamma, x, b: gamma + b
    """

    def __call__(
        self, gamma: torch.Tensor | None, x: torch.Tensor | None, b: torch.Tensor
    ) -> torch.Tensor: ...


class RegressionFn(Protocol):
    """The regression function protocol.

    It must be able to yield 3D and 4D tensors given 1D or 2D time inputs, as well
        as `psi` input of order 2 or 3. This is not very restrictive, but requires to be
        careful. The last dimension is the dimension of the response variable; second
        last is the repeated measurements; third last is individual based; possible
        fourth last is for parallelization of the MCMC sampler.

        It is identical to LinkFn.

    Args:
        t (torch.Tensor): The evaluation times.
        psi (torch.Tensor): The individual parameters.

    Returns:
        torch.Tensor: The response variable values.

    Examples:
        >>> def sigmoid(t: torch.Tensor, psi: torch.Tensor):
        ...     scale, offset, slope = psi.chunk(3, dim=-1)
        ...     # Fully broadcasted
        ...     return (scale * torch.sigmoid((t - offset) / slope)).unsqueeze(-1)
        >>> regression_fn = sigmoid
    """

    def __call__(self, t: torch.Tensor, psi: torch.Tensor) -> torch.Tensor: ...


class LinkFn(Protocol):
    """The link function protocol.

    It must be able to yield 3D and 4D tensors given 1D or 2D time inputs, as well
        as `psi` input of order 2 or 3. This is not very restrictive, but requires to be
        careful. The last dimension is the dimension of the response variable; second
        last is the repeated measurements; third last is individual based; possible
        fourth last is for parallelization of the MCMC sampler.

        It is identical to RegressionFn.

    Args:
        t (torch.Tensor): The evaluation times.
        psi (torch.Tensor): The individual parameters.

    Returns:
        torch.Tensor: The response variable values.

    Examples:
        >>> def sigmoid(t: torch.Tensor, psi: torch.Tensor):
        ...     scale, offset, slope = psi.chunk(3, dim=-1)
        ...     # Fully broadcasted
        ...     return (scale * torch.sigmoid((t - offset) / slope)).unsqueeze(-1)
        >>> link_fn = sigmoid
    """

    def __call__(self, t: torch.Tensor, psi: torch.Tensor) -> torch.Tensor: ...


class ClockMethod(Protocol):
    r"""The clock method protocol.

    This protocol is useful to differentiate between two natural mappings when dealing
    with base hazards:

    .. math::
        (t_0, t_1) \mapsto \begin{cases} t_1 - t_0 \text{(clock reset)}, \\ t_1
        \text{(clock forward)} \end{cases}.
    """

    def __call__(self, t0: torch.Tensor, t1: torch.Tensor) -> torch.Tensor: ...


# Abstract classes
class LogBaseHazardFn(nn.Module, ABC):
    """The log base hazard base class.

    This is not a protocol because caching is done, and therefore a key is required.
    Making this a `nn.Module`, one can check the value of the parameters and store their
    hashed values at the same time as the

    Note that the log base hazard function is in log scale, and expects a former
    transition time column vector `t0` as well as other times points at which the log
    base hazard is to be computed. `t1` is a matrix with the same number of rows as
    `t0`.

    Implement a `forward` and do not forget the `super().__init__()` when declaring your
    own class.

    Pass the parameters you want to optimize in the `ModelParams.extra` attribute as a
    list. If you do not want them to be optimized, then by default they do not require
    gradients for the given implementations.
    """

    @property
    def key(self) -> tuple[int, ...]:
        """Returns a hash containing the class type and parameters if there are any.

        Returns:
            tuple[int, ...]: A key used in caching operations.
        """
        ident = id(self)
        parameters = list(self.parameters())
        if parameters == []:
            return (ident,)

        params_flat = nn.utils.parameters_to_vector(self.parameters())
        return (ident, xxh3_64_intdigest(params_flat.detach().numpy()))  # type: ignore

    @abstractmethod
    def forward(self, t0: torch.Tensor, t1: torch.Tensor) -> torch.Tensor: ...


# Named tuples
class BucketData(NamedTuple):
    """A simple `NamedTuple` containing transition information.

    Attributes:
        idxs (torch.Tensor): The individual indices.
        t0 (torch.Tensor): A column vector of previous transition times.
        t1 (torch.Tensor): A column vecotr of next transition times.
    """

    idxs: torch.Tensor
    t0: torch.Tensor
    t1: torch.Tensor


class HazardInfo(NamedTuple):
    """A simple internal `NamedTuple` required for hazard computation.

    Attributes:
        t0 (torch.Tensor): A column vector of previous transition times.
        t1 (torch.Tensor): A matrix of next transition times.
        x (torch.Tensor | None): The covariates.
        psi (torch.Tensor): The individual parameters.
        alpha (torch.Tensor): The baseline hazard parameters.
        beta (torch.Tensor | None): The regression parameters.
        log_base_hazard_fn (LogBaseHazardFn): The log base hazard function.
        link_fn (LinkFn): The link function.
    """

    t0: torch.Tensor
    t1: torch.Tensor
    x: torch.Tensor | None
    psi: torch.Tensor
    alpha: torch.Tensor
    beta: torch.Tensor | None
    log_base_hazard_fn: LogBaseHazardFn
    link_fn: LinkFn


# Constants
LOG_TWO_PI: Final[torch.Tensor] = torch.log(torch.tensor(2.0 * torch.pi))
HAZARD_CACHE_KEYS: Final[tuple[str, ...]] = ("base", "half", "quad_c", "quad_lc")
SIGNIFICANCE_LEVELS: Final[tuple[float, ...]] = (
    0.001,
    0.01,
    0.05,
    0.1,
    float("inf"),
)
SIGNIFICANCE_CODES: Final[tuple[str, ...]] = (
    "[red3]***[/]",
    "[orange3]**[/]",
    "[yellow3]*[/]",
    ".",
    "",
)
