from __future__ import annotations

__all__ = ["PoissonTensorGenerator", "RandPoissonTensorGenerator"]


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

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


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

    Example usage:

    .. code-block:: 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: Generator | None = None) -> Tensor:
        return torch.poisson(self._rate.generate(size=size, rng=rng), generator=rng)


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

    Args:
    ----
        rate (float, optional): Specifies the rate of the Poisson
            distribution. This value has to be greater than 0.
            Default: ``1.0``

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

    Example usage:

    .. code-block:: 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:
            raise ValueError(f"rate has to be greater than 0 (received: {rate})")
        self._rate = float(rate)

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

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