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

from __future__ import annotations

__all__ = ["PoissonTensorGenerator", "RandPoissonTensorGenerator"]


import torch
from coola.utils.format import str_indent, str_mapping

from startorch.random import rand_poisson
from startorch.tensor.base import BaseTensorGenerator, setup_tensor_generator


class PoissonTensorGenerator(BaseTensorGenerator):
    r"""Implement a class to generate tensors 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 tensor. 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.tensor import RandUniform, Poisson
    >>> generator = Poisson(rate=RandUniform(low=1.0, high=2.0))
    >>> generator
    PoissonTensorGenerator(
      (rate): RandUniformTensorGenerator(low=1.0, high=2.0)
    )
    >>> generator.generate((2, 6))
    tensor([[...]])

    ```
    """

    def __init__(self, rate: BaseTensorGenerator | dict) -> None:
        super().__init__()
        self._rate = setup_tensor_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, size: tuple[int, ...], rng: torch.Generator | None = None) -> torch.Tensor:
        return torch.poisson(self._rate.generate(size=size, rng=rng), generator=rng)


class RandPoissonTensorGenerator(BaseTensorGenerator):
    r"""Implement a class to generate tensors by sampling values from a
    Poisson distribution.

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

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

    Example usage:

    ```pycon
    >>> from startorch.tensor import RandPoisson
    >>> generator = RandPoisson(rate=1.0)
    >>> generator
    RandPoissonTensorGenerator(rate=1.0)
    >>> generator.generate((2, 6))
    tensor([[...]])

    ```
    """

    def __init__(self, rate: float = 1.0) -> None:
        super().__init__()
        if rate <= 0:
            msg = f"rate has to be greater than 0 (received: {rate})"
            raise ValueError(msg)
        self._rate = float(rate)

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

    def generate(self, size: tuple[int, ...], rng: torch.Generator | None = None) -> torch.Tensor:
        return rand_poisson(size=size, rate=self._rate, generator=rng)
