r"""Contain the implementation of sequence generators where the values
are sampled from a Poisson distribution."""

from __future__ import annotations

__all__ = ["PoissonSequenceGenerator", "RandPoissonSequenceGenerator"]


import torch
from coola.utils import str_indent, str_mapping

from startorch.random import rand_poisson
from startorch.sequence.base import BaseSequenceGenerator, setup_sequence_generator
from startorch.sequence.constant import ConstantSequenceGenerator
from startorch.utils.conversion import to_tuple


class PoissonSequenceGenerator(BaseSequenceGenerator):
    r"""Implement a class to generate sequence by sampling values from a
    Poisson distribution.

    The rates of the Poisson distribution are generated by the rate
    generator. The rate generator should return the rate for each value
    in the sequence. The rate values should be greater than 0.

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

    Example usage:

    ```pycon
    >>> from startorch.sequence import RandUniform, Poisson
    >>> generator = Poisson(rate=RandUniform(low=1.0, high=2.0))
    >>> generator
    PoissonSequenceGenerator(
      (rate): RandUniformSequenceGenerator(low=1.0, high=2.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 torch.poisson(
            self._rate.generate(seq_len=seq_len, batch_size=batch_size, rng=rng),
            generator=rng,
        )

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

        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 of the Poisson
                distribution are sampled from a uniform distribution
                (``UniformConstantSequenceGenerator``).
        """
        # The import is here to do not generate circular dependencies
        from startorch.sequence.uniform import RandUniformSequenceGenerator

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


class RandPoissonSequenceGenerator(BaseSequenceGenerator):
    r"""Implement a class to generate sequence by sampling values from a
    Poisson distribution.

    Args:
        rate: Specifies the rate of the Poisson distribution.
            This value has to be greater than 0.
        feature_size: Specifies the feature size.

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

    Example usage:

    ```pycon
    >>> from startorch.sequence import RandPoisson
    >>> generator = RandPoisson(rate=1.0)
    >>> generator
    RandPoissonSequenceGenerator(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}, feature_size={self._feature_size})"
        )

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