r"""Contain the implementation of sequence generators where the values
are sampled from an Exponential distribution."""

from __future__ import annotations

__all__ = [
    "ExponentialSequenceGenerator",
    "RandExponentialSequenceGenerator",
    "RandTruncExponentialSequenceGenerator",
    "TruncExponentialSequenceGenerator",
]

from typing import TYPE_CHECKING

from coola.utils.format import str_indent, str_mapping

from startorch.random import (
    exponential,
    rand_exponential,
    rand_trunc_exponential,
    trunc_exponential,
)
from startorch.sequence.base import BaseSequenceGenerator, setup_sequence_generator
from startorch.sequence.constant import ConstantSequenceGenerator, FullSequenceGenerator
from startorch.utils.conversion import to_tuple

if TYPE_CHECKING:
    import torch


class ExponentialSequenceGenerator(BaseSequenceGenerator):
    r"""Implement a class to generate sequence by sampling values from an
    Exponential distribution.

    The rates of the Exponential distribution are generated by the
    rate generator. The rate generator should return the rate for each
    value in the sequence.

    Args:
        rate: Specifies the rate generator or its configuration.
            The rate generator should return valid rate values.

    Example usage:

    ```pycon
    >>> from startorch.sequence import Exponential, RandUniform
    >>> generator = Exponential(rate=RandUniform(low=1.0, high=10.0))
    >>> generator
    ExponentialSequenceGenerator(
      (rate): RandUniformSequenceGenerator(low=1.0, high=10.0, feature_size=(1,))
    )
    >>> generator.generate(seq_len=6, batch_size=2)
    tensor([[...]])

    ```
    """

    def __init__(self, rate: BaseSequenceGenerator | dict) -> None:
        super().__init__()
        self._rate = setup_sequence_generator(rate)

    def __repr__(self) -> str:
        args = str_indent(str_mapping({"rate": self._rate}))
        return f"{self.__class__.__qualname__}(\n  {args}\n)"

    def generate(
        self, seq_len: int, batch_size: int = 1, rng: torch.Generator | None = None
    ) -> torch.Tensor:
        return exponential(
            self._rate.generate(seq_len=seq_len, batch_size=batch_size, rng=rng),
            generator=rng,
        )

    @classmethod
    def create_fixed_rate(
        cls, rate: float = 1.0, feature_size: tuple[int, ...] | list[int] | int = 1
    ) -> ExponentialSequenceGenerator:
        r"""Implement a sequence generator where the values are sampled
        from an Exponential distribution with a fixed rate.

        Args:
            rate: Specifies the rate of the Exponential distribution.
            feature_size: Specifies the feature size.

        Returns:
            A sequence generator where the rates of the Exponential
                distribution are a fixed given value.

        Example usage:

        ```pycon
        >>> from startorch.sequence import Exponential, RandUniform
        >>> generator = Exponential.create_fixed_rate(rate=1.0)
        >>> generator
        ExponentialSequenceGenerator(
          (rate): FullSequenceGenerator(value=1.0, feature_size=(1,))
        )
        >>> generator.generate(seq_len=6, batch_size=2)
        tensor([[...]])

        ```
        """
        return cls(FullSequenceGenerator(value=rate, feature_size=feature_size))

    @classmethod
    def create_uniform_rate(
        cls,
        min_rate: float = 0.01,
        max_rate: float = 1.0,
        feature_size: tuple[int, ...] | list[int] | int = 1,
    ) -> ExponentialSequenceGenerator:
        r"""Implement a sequence generator where the rates of the
        Exponential distribution are sampled from a uniform
        distribution.

        One rate is sampled per sequence.

        Args:
            min_rate: Specifies the minimum rate value.
            max_rate: Specifies the maximum rate value.
            feature_size: Specifies the feature size.

        Returns:
            A sequence generator where the rates for each sequence are
                sampled from a uniform distribution.

        Example usage:

        ```pycon
        >>> from startorch.sequence import Exponential, RandUniform
        >>> generator = Exponential.create_uniform_rate(min_rate=0.1, max_rate=1.0)
        >>> generator
        ExponentialSequenceGenerator(
          (rate): ConstantSequenceGenerator(
              (sequence): RandUniformSequenceGenerator(low=0.1, high=1.0, feature_size=(1,))
            )
        )
        >>> generator.generate(seq_len=6, batch_size=2)
        tensor([[...]])

        ```
        """
        # The import is here to do not generate circular dependencies
        from startorch.sequence.uniform import RandUniformSequenceGenerator

        return cls(
            ConstantSequenceGenerator(
                RandUniformSequenceGenerator(
                    low=min_rate,
                    high=max_rate,
                    feature_size=feature_size,
                )
            ),
        )


class RandExponentialSequenceGenerator(BaseSequenceGenerator):
    r"""Implement a class to generate sequences by sampling values from
    an Exponential distribution.

    Args:
        rate: Specifies the rate of the Exponential distribution.
        feature_size: Specifies the feature size.

    Raises:
        ValueError: if ``rate`` is not a positive number.

    Example usage:

    ```pycon
    >>> from startorch.sequence import RandExponential
    >>> generator = RandExponential(rate=1.0)
    >>> generator
    RandExponentialSequenceGenerator(rate=1.0, feature_size=(1,))
    >>> generator.generate(seq_len=6, batch_size=2)
    tensor([[...]])

    ```
    """

    def __init__(
        self,
        rate: float = 1.0,
        feature_size: tuple[int, ...] | list[int] | int = 1,
    ) -> None:
        super().__init__()
        if rate <= 0:
            msg = f"rate has to be greater than 0 (received: {rate})"
            raise ValueError(msg)
        self._rate = float(rate)
        self._feature_size = to_tuple(feature_size)

    def __repr__(self) -> str:
        return (
            f"{self.__class__.__qualname__}(rate={self._rate}, "
            f"feature_size={self._feature_size})"
        )

    def generate(
        self, seq_len: int, batch_size: int = 1, rng: torch.Generator | None = None
    ) -> torch.Tensor:
        return rand_exponential(
            size=(batch_size, seq_len) + self._feature_size,
            rate=self._rate,
            generator=rng,
        )


class RandTruncExponentialSequenceGenerator(BaseSequenceGenerator):
    r"""Implement a class to generate sequences by sampling values from a
    truncated Exponential distribution.

    Args:
        rate: Specifies the rate of the Exponential distribution.
        max_value: Specifies the maximum value.
        feature_size: Specifies the feature size.

    Raises:
        ValueError: if ``rate`` is not a positive number.
        ValueError: if ``max_value`` is not a positive number.

    Example usage:

    ```pycon
    >>> from startorch.sequence import RandTruncExponential
    >>> generator = RandTruncExponential(rate=1.0, max_value=3.0)
    >>> generator
    RandTruncExponentialSequenceGenerator(rate=1.0, max_value=3.0, feature_size=(1,))
    >>> generator.generate(seq_len=6, batch_size=2)
    tensor([[...]])

    ```
    """

    def __init__(
        self,
        rate: float = 1.0,
        max_value: float = 5.0,
        feature_size: tuple[int, ...] | list[int] | int = 1,
    ) -> None:
        super().__init__()
        if rate <= 0:
            msg = f"rate has to be greater than 0 (received: {rate})"
            raise ValueError(msg)
        self._rate = float(rate)
        if max_value <= 0:
            msg = f"max_value has to be greater than 0 (received: {max_value})"
            raise ValueError(msg)
        self._max_value = float(max_value)
        self._feature_size = to_tuple(feature_size)

    def __repr__(self) -> str:
        return (
            f"{self.__class__.__qualname__}(rate={self._rate}, max_value={self._max_value}, "
            f"feature_size={self._feature_size})"
        )

    def generate(
        self, seq_len: int, batch_size: int = 1, rng: torch.Generator | None = None
    ) -> torch.Tensor:
        return rand_trunc_exponential(
            size=(batch_size, seq_len) + self._feature_size,
            rate=self._rate,
            max_value=self._max_value,
            generator=rng,
        )


class TruncExponentialSequenceGenerator(BaseSequenceGenerator):
    r"""Implement a class to generate sequence by sampling values from an
    Exponential distribution.

    Args:
        rate: Specifies a sequence generator (or its configuration)
            to generate the rate.
        max_value: Specifies a sequence generator (or its
            configuration) to generate the maximum value (excluded).

    Example usage:

    ```pycon
    >>> from startorch.sequence import RandUniform, TruncExponential
    >>> generator = TruncExponential(
    ...     rate=RandUniform(low=1.0, high=10.0),
    ...     max_value=RandUniform(low=1.0, high=100.0),
    ... )
    >>> generator
    TruncExponentialSequenceGenerator(
      (rate): RandUniformSequenceGenerator(low=1.0, high=10.0, feature_size=(1,))
      (max_value): RandUniformSequenceGenerator(low=1.0, high=100.0, feature_size=(1,))
    )
    >>> generator.generate(seq_len=6, batch_size=2)
    tensor([[...]])

    ```
    """

    def __init__(
        self,
        rate: BaseSequenceGenerator | dict,
        max_value: BaseSequenceGenerator | dict,
    ) -> None:
        super().__init__()
        self._rate = setup_sequence_generator(rate)
        self._max_value = setup_sequence_generator(max_value)

    def __repr__(self) -> str:
        args = str_indent(str_mapping({"rate": self._rate, "max_value": self._max_value}))
        return f"{self.__class__.__qualname__}(\n  {args}\n)"

    def generate(
        self, seq_len: int, batch_size: int = 1, rng: torch.Generator | None = None
    ) -> torch.Tensor:
        return trunc_exponential(
            rate=self._rate.generate(seq_len=seq_len, batch_size=batch_size, rng=rng),
            max_value=self._max_value.generate(seq_len=seq_len, batch_size=batch_size, rng=rng),
            generator=rng,
        )
