from __future__ import annotations

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


from coola.utils.format import str_indent, str_mapping
from redcat import BatchedTensorSeq
from torch import Generator

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


class ExponentialSequenceGenerator(BaseSequenceGenerator):
    r"""Implements 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 (``BaseSequenceGenerator`` or dict):
            Specifies the rate generator or its configuration.
            The rate generator should return valid rate values.

    Example usage:

    .. code-block:: 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([[...]], batch_dim=0, seq_dim=1)
    """

    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: Generator | None = None
    ) -> BatchedTensorSeq:
        return BatchedTensorSeq(
            exponential(
                self._rate.generate(seq_len=seq_len, batch_size=batch_size, rng=rng).data,
                generator=rng,
            )
        )

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

        Args:
        ----
            rate (float, optional): Specifies the rate of the
                Exponential distribution. Default: ``1.0``
            feature_size (tuple or list or int, optional): Specifies the
                feature size. Default: ``1``

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

        Example usage:

        .. code-block:: 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([[...]], batch_dim=0, seq_dim=1)
        """
        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"""Implements 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 (float, optional): Specifies the minimum rate
                value. Default: ``0.01``
            max_rate (float, optional): Specifies the maximum rate
                value. Default: ``1.0``
            feature_size (tuple or list or int, optional): Specifies the
                feature size. Default: ``1``

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

        Example usage:

        .. code-block:: 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([[...]], batch_dim=0, seq_dim=1)
        """
        # 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"""Implements a class to generate sequences by sampling values from
    an Exponential distribution.

    Args:
    ----
        rate (float, optional): Specifies the rate of the Exponential
            distribution. Default: ``1.0``
        feature_size (tuple or list or int, optional): Specifies the
            feature size. Default: ``1``

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

    Example usage:

    .. code-block:: 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([[...]], batch_dim=0, seq_dim=1)
    """

    def __init__(
        self,
        rate: float = 1.0,
        feature_size: tuple[int, ...] | list[int] | int = 1,
    ) -> None:
        super().__init__()
        if rate <= 0:
            raise ValueError(f"rate has to be greater than 0 (received: {rate})")
        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: Generator | None = None
    ) -> BatchedTensorSeq:
        return BatchedTensorSeq(
            rand_exponential(
                size=(batch_size, seq_len) + self._feature_size,
                rate=self._rate,
                generator=rng,
            )
        )


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

    Args:
    ----
        rate (float, optional): Specifies the rate of the Exponential
            distribution. Default: ``1.0``
        max_value (float, optional): Specifies the maximum value.
            Default: ``5.0``
        feature_size (tuple or list or int, optional): Specifies the
            feature size. Default: ``1``

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

    Example usage:

    .. code-block:: 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([[...]], batch_dim=0, seq_dim=1)
    """

    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:
            raise ValueError(f"rate has to be greater than 0 (received: {rate})")
        self._rate = float(rate)
        if max_value <= 0:
            raise ValueError(f"max_value has to be greater than 0 (received: {max_value})")
        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: Generator | None = None
    ) -> BatchedTensorSeq:
        return BatchedTensorSeq(
            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"""Implements a class to generate sequence by sampling values from
    an Exponential distribution.

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

    Example usage:

    .. code-block:: 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([[...]], batch_dim=0, seq_dim=1)
    """

    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: Generator | None = None
    ) -> BatchedTensorSeq:
        return BatchedTensorSeq(
            trunc_exponential(
                rate=self._rate.generate(seq_len=seq_len, batch_size=batch_size, rng=rng).data,
                max_value=self._max_value.generate(
                    seq_len=seq_len, batch_size=batch_size, rng=rng
                ).data,
                generator=rng,
            )
        )
