r"""Contain functions to sample values from continuous uni-variate
distribution supported on a bounded interval."""

from __future__ import annotations

__all__ = [
    "asinh_uniform",
    "log_uniform",
    "rand_asinh_uniform",
    "rand_log_uniform",
    "rand_trunc_cauchy",
    "rand_trunc_exponential",
    "rand_trunc_half_cauchy",
    "rand_trunc_half_normal",
    "rand_trunc_log_normal",
    "rand_trunc_normal",
    "rand_uniform",
    "trunc_cauchy",
    "trunc_exponential",
    "trunc_half_cauchy",
    "trunc_half_normal",
    "trunc_log_normal",
    "trunc_normal",
    "uniform",
]

import math

import torch
from torch.distributions import (
    Cauchy,
    Exponential,
    HalfCauchy,
    HalfNormal,
    LogNormal,
    Normal,
)

from startorch.utils.tensor import shapes_are_equal


def rand_trunc_cauchy(
    size: list[int] | tuple[int, ...],
    loc: float = 0.0,
    scale: float = 1.0,
    min_value: float = -2.0,
    max_value: float = 2.0,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    r"""Create a tensor filled with values sampled from a truncated
    Cauchy distribution.

    Args:
        size: Specifies the tensor shape.
        loc: Specifies the location of the Cauchy distribution.
        scale: Specifies the scale of the Cauchy distribution.
        min_value: Specifies the minimum value.
        max_value: Specifies the maximum value.
        generator: Specifies an optional random generator.

    Returns:
        A tensor filled with values sampled from a truncated Cauchy
            distribution

    Raises:
        ValueError: if the ``max_value`` and ``min_value`` parameters
            are not valid.

    Example usage:

    ```pycon
    >>> import torch
    >>> from startorch.random import rand_trunc_cauchy
    >>> rand_trunc_cauchy((2, 3), loc=1.0, scale=2.0, min_value=-3.0, max_value=3.0)
    tensor([[...]])

    ```
    """
    if max_value < min_value:
        msg = f"`max_value` ({max_value}) has to be greater or equal to `min_value` ({min_value})"
        raise ValueError(msg)
    distribution = Cauchy(loc=loc, scale=scale)
    # Values are generated by using a truncated uniform distribution and
    # then using the inverse CDF for the Cauchy distribution.
    # Get upper and lower CDF values
    lo, up = distribution.cdf(torch.tensor([min_value, max_value]))
    return distribution.icdf(rand_uniform(size, low=lo, high=up, generator=generator)).clamp(
        min=min_value, max=max_value
    )


def trunc_cauchy(
    loc: torch.Tensor,
    scale: torch.Tensor,
    min_value: torch.Tensor,
    max_value: torch.Tensor,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    r"""Create a tensor filled with values sampled from a truncated
    Cauchy distribution.

    Args:
        loc: Specifies the location/median of the Cauchy distribution.
            It must be a float tensor of shape ``(d0, d1, ..., dn)``.
        scale: Specifies the scale of the Cauchy distribution.
            It must be a float tensor of shape ``(d0, d1, ..., dn)``.
        min_value: Specifies the minimum value. It must be a float
            tensor of shape ``(d0, d1, ..., dn)``.
        max_value: Specifies the maximum value. It must be a float
            tensor of shape ``(d0, d1, ..., dn)``.
        generator: Specifies an optional random generator.

    Returns:
        A tensor of shape ``(d0, d1, ..., dn)`` filled with values
            sampled from a truncated Cauchy distribution

    Raises:
        ValueError: if the ``loc``, ``scale``, ``max_value`` and
            ``min_value`` parameters are not valid.

    Example usage:

    ```pycon
    >>> import torch
    >>> from startorch.random import trunc_cauchy
    >>> trunc_cauchy(
    ...     loc=torch.tensor([1.0, 0.0, -1.0]),
    ...     scale=torch.tensor([1.0, 3.0, 5.0]),
    ...     min_value=torch.tensor([-5.0, -10.0, -15.0]),
    ...     max_value=torch.tensor([5.0, 10.0, 15.0]),
    ... )
    tensor([...])

    ```
    """
    if not shapes_are_equal((loc, scale, min_value, max_value)):
        msg = (
            f"Incorrect shapes. The shapes of all the input tensors must be equal: loc={loc.shape}"
            f"  scale={scale.shape}  min_value={min_value.shape}  max_value={max_value.shape}"
        )
        raise ValueError(msg)
    if torch.any(scale.le(0.0)):
        msg = "All the `scale` values have to be greater than 0"
        raise ValueError(msg)
    if torch.any(max_value.lt(min_value)):
        msg = (
            "Found at least one value in `min_value` that is higher than its associated "
            "`max_value`"
        )
        raise ValueError(msg)
    distribution = Cauchy(loc=loc, scale=scale)
    # Values are generated by using a truncated uniform distribution and
    # then using the inverse CDF for the Cauchy distribution.
    # Get upper and lower CDF values
    lo, up = distribution.cdf(torch.stack([min_value, max_value]))
    return distribution.icdf(uniform(low=lo, high=up, generator=generator)).clamp(
        min=min_value, max=max_value
    )


def rand_trunc_exponential(
    size: list[int] | tuple[int, ...],
    rate: float = 1.0,
    max_value: float = 5.0,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    r"""Create a tensor filled with values sampled from a truncated
    Exponential distribution.

    Args:
        size: Specifies the tensor shape.
        rate: Specifies the rate of the Exponential distribution.
        max_value: Specifies the maximum value.
        generator: Specifies an optional random generator.

    Returns:
        A tensor filled with values sampled from a truncated
            Exponential distribution

    Raises:
        ValueError: if the ``max_value`` parameter is not valid.

    Example usage:

    ```pycon
    >>> import torch
    >>> from startorch.random import rand_trunc_exponential
    >>> rand_trunc_exponential((2, 3), rate=1.0, max_value=3.0)
    tensor([[...]])

    ```
    """
    if max_value <= 0:
        msg = f"`max_value` has to be greater than 0 (received: {max_value})"
        raise ValueError(msg)
    distribution = Exponential(rate=rate)
    # Values are generated by using a truncated uniform distribution and
    # then using the inverse CDF for the Exponential distribution.
    # Get upper and lower CDF values
    up = distribution.cdf(torch.tensor(max_value)).item()
    return distribution.icdf(rand_uniform(size, low=0.0, high=up, generator=generator)).clamp(
        max=max_value
    )


def trunc_exponential(
    rate: torch.Tensor,
    max_value: torch.Tensor,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    r"""Create a tensor filled with values sampled from a truncated
    Exponential distribution.

    Args:
        rate: Specifies the rate of the Exponential distribution.
            It must be a float tensor of shape ``(d0, d1, ..., dn)``.
        max_value: Specifies the maximum value. It must be a float
            tensor of shape ``(d0, d1, ..., dn)``.
        generator: Specifies an optional random generator.

    Returns:
        A tensor of shape ``(d0, d1, ..., dn)`` filled with values
            sampled from a truncated Exponential distribution.

    Raises:
        ValueError: if the ``rate`` and ``max_value`` parameter are not
            valid.

    Example usage:

    ```pycon
    >>> import torch
    >>> from startorch.random import trunc_exponential
    >>> trunc_exponential(
    ...     rate=torch.tensor([1.0, 3.0, 5.0]),
    ...     max_value=torch.tensor([5.0, 10.0, 15.0]),
    ... )
    tensor([...])

    ```
    """
    if rate.shape != max_value.shape:
        msg = (
            "Incorrect shapes. The shapes of all the input tensors must be equal:  "
            f"rate={rate.shape}  max_value={max_value.shape})"
        )
        raise ValueError(msg)
    if torch.any(rate.le(0.0)):
        msg = "All the `rate` values have to be greater than 0"
        raise ValueError(msg)
    if torch.any(max_value.le(0.0)):
        msg = "Found at least one value in `max_value` that is lower or equal to 0"
        raise ValueError(msg)
    distribution = Exponential(rate=rate)
    # Values are generated by using a truncated uniform distribution and
    # then using the inverse CDF for the Exponential distribution.
    # Get upper and lower CDF values
    u = distribution.cdf(max_value)
    return distribution.icdf(torch.rand(rate.shape, generator=generator).mul(u)).clamp(
        max=max_value
    )


def rand_trunc_half_cauchy(
    size: list[int] | tuple[int, ...],
    scale: float = 1.0,
    max_value: float = 4.0,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    r"""Create a tensor filled with values sampled from a truncated half-
    Cauchy distribution.

    Args:
        size: Specifies the tensor shape.
        scale: Specifies the scale of the half-Cauchy distribution.
        max_value: Specifies the maximum value.
        generator: Specifies an optional random generator.

    Returns:
        A tensor filled with values sampled from a truncated
            half-Cauchy distribution.

    Raises:
        ValueError: if the ``max_value`` parameter is not valid.

    Example usage:

    ```pycon
    >>> import torch
    >>> from startorch.random import rand_trunc_half_cauchy
    >>> rand_trunc_half_cauchy((2, 3), scale=1.0, max_value=3.0)
    tensor([[...]])

    ```
    """
    if max_value <= 0:
        msg = f"`max_value` has to be greater than 0 (received: {max_value})"
        raise ValueError(msg)
    distribution = HalfCauchy(scale=scale)
    # Values are generated by using a truncated uniform distribution and
    # then using the inverse CDF for the Cauchy distribution.
    # Get upper CDF value
    u = distribution.cdf(torch.tensor(max_value))
    return distribution.icdf(rand_uniform(size, low=0.0, high=u, generator=generator)).clamp(
        max=max_value
    )


def trunc_half_cauchy(
    scale: torch.Tensor,
    max_value: torch.Tensor,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    r"""Create a tensor filled with values sampled from a truncated half-
    Cauchy distribution.

    Args:
        scale: Specifies the scale of the half-Cauchy distribution.
            It must be a float tensor of shape ``(d0, d1, ..., dn)``.
        max_value: Specifies the maximum value. It must be a float
            tensor of shape ``(d0, d1, ..., dn)``.
        generator: Specifies an optional random generator.

    Returns:
        A tensor of shape ``(d0, d1, ..., dn)`` filled with values
            sampled from a truncated Cauchy distribution.

    Raises:
        ValueError: if the ``scale`` and ``max_value`` parameter are
            not valid.

    Example usage:

    ```pycon
    >>> import torch
    >>> from startorch.random import trunc_half_cauchy
    >>> trunc_half_cauchy(
    ...     scale=torch.tensor([1.0, 3.0, 5.0]),
    ...     max_value=torch.tensor([5.0, 10.0, 15.0]),
    ... )
    tensor([...])

    ```
    """
    if scale.shape != max_value.shape:
        msg = (
            "Incorrect shapes. The shapes of all the input tensors must be equal:  "
            f"scale={scale.shape}  max_value={max_value.shape}"
        )
        raise ValueError(msg)
    if torch.any(scale.le(0.0)):
        msg = "All the `scale` values have to be greater than 0"
        raise ValueError(msg)
    if torch.any(max_value.le(0.0)):
        msg = "Found at least one value in `max_value` that is lower or equal to 0"
        raise ValueError(msg)
    distribution = HalfCauchy(scale=scale)
    # Values are generated by using a truncated uniform distribution and
    # then using the inverse CDF for the Cauchy distribution.
    # Get upper CDF value
    u = distribution.cdf(max_value)
    return distribution.icdf(torch.rand(scale.shape, generator=generator).mul(u)).clamp(
        max=max_value
    )


def rand_trunc_half_normal(
    size: list[int] | tuple[int, ...],
    std: float = 1.0,
    max_value: float = 5.0,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    r"""Create a tensor filled with values sampled from a truncated half-
    Normal distribution.

    Args:
        size: Specifies the tensor shape.
        std: Specifies the standard deviation of the half-Normal
            distribution.
        max_value: Specifies the maximum value.
        generator: Specifies an optional random generator.

    Returns:
        A tensor filled with values sampled from a truncated
            half-Normal distribution.

    Raises:
        ValueError: if the ``max_value`` parameter is not valid.

    Example usage:

    ```pycon
    >>> import torch
    >>> from startorch.random import rand_trunc_half_normal
    >>> rand_trunc_half_normal((2, 3), std=1.0, max_value=3.0)
    tensor([[...]])

    ```
    """
    if max_value <= 0:
        msg = f"`max_value` has to be greater than 0 (received: {max_value})"
        raise ValueError(msg)
    distribution = HalfNormal(scale=std)
    # Values are generated by using a truncated uniform distribution and
    # then using the inverse CDF for the Normal distribution.
    # Get upper cdf value
    u = distribution.cdf(torch.tensor(max_value))
    return distribution.icdf(rand_uniform(size, low=0.0, high=u, generator=generator)).clamp(
        max=max_value
    )


def trunc_half_normal(
    std: torch.Tensor,
    max_value: torch.Tensor,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    r"""Create a tensor filled with values sampled from a truncated half-
    Normal distribution.

    Args:
        std: Specifies the standard deviation of the half-Normal
            distribution. It must be a float tensor of shape
            ``(d0, d1, ..., dn)``.
        max_value: Specifies the maximum value. It must be a float
            tensor of shape ``(d0, d1, ..., dn)``.
        generator: Specifies an optional random generator.

    Returns:
        A tensor of shape ``(d0, d1, ..., dn)`` filled with values
            sampled from a truncated half-Normal distribution.

    Raises:
        ValueError: if the ``std`` and ``max_value`` parameters are not
            valid.

    Example usage:

    ```pycon
    >>> import torch
    >>> from startorch.random import trunc_half_normal
    >>> trunc_half_normal(
    ...     std=torch.tensor([1.0, 3.0, 5.0]), max_value=torch.tensor([5.0, 10.0, 15.0])
    ... )
    tensor([...])

    ```
    """
    if std.shape != max_value.shape:
        msg = (
            "Incorrect shapes. The shapes of all the input tensors must be equal: "
            f"std={std.shape}  max_value={max_value.shape}"
        )
        raise ValueError(msg)
    if torch.any(std.le(0.0)):
        msg = "All the `std` values have to be greater than 0"
        raise ValueError(msg)
    if torch.any(max_value.le(0.0)):
        msg = "All the `max_value` values must be greater than 0"
        raise ValueError(msg)
    distribution = HalfNormal(scale=std)
    # Values are generated by using a uniform distribution and
    # then using the inverse CDF for the Normal distribution.
    # Get upper and lower CDF values
    u = distribution.cdf(max_value)
    return distribution.icdf(torch.rand(std.shape, generator=generator).mul(u)).clamp(max=max_value)


def rand_trunc_log_normal(
    size: list[int] | tuple[int, ...],
    mean: float = 0.0,
    std: float = 1.0,
    min_value: float = 0.0,
    max_value: float = 5.0,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    r"""Create a tensor filled with values sampled from a truncated log-
    Normal distribution.

    Args:
        size: Specifies the tensor shape.
        mean: Specifies the mean of the underlying Normal distribution.
        std: Specifies the standard deviation of the underlying Normal
            distribution.
        min_value: Specifies the minimum value.
        max_value: Specifies the maximum value.
        generator: Specifies an optional random generator.

    Returns:
        A tensor filled with values sampled from a truncated log-Normal
            distribution.

    Raises:
        ValueError: if the ``min_value`` and ``max_value`` parameters
            are not valid.

    Example usage:

    ```pycon
    >>> import torch
    >>> from startorch.random import rand_trunc_log_normal
    >>> rand_trunc_log_normal((2, 3), mean=0.0, std=1.0, min_value=1.0, max_value=4.0)
    tensor([[...]])

    ```
    """
    if max_value < min_value:
        mag = f"`max_value` ({max_value}) has to be greater or equal to `min_value` ({min_value})"
        raise ValueError(mag)
    distribution = LogNormal(loc=mean, scale=std)
    # Values are generated by using a truncated uniform distribution and
    # then using the inverse CDF for the log-Normal distribution.
    # Get upper and lower CDF values
    lo, up = distribution.cdf(torch.tensor([min_value, max_value]))
    return distribution.icdf(rand_uniform(size, low=lo, high=up, generator=generator)).clamp(
        min=min_value, max=max_value
    )


def trunc_log_normal(
    mean: torch.Tensor,
    std: torch.Tensor,
    min_value: torch.Tensor,
    max_value: torch.Tensor,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    r"""Create a tensor filled with values sampled from a truncated log-
    Normal distribution.

    Args:
        mean: Specifies the mean of the underlying Normal distribution.
            It must be a float tensor of shape ``(d0, d1, ..., dn)``.
        std: Specifies the standard deviation of the underlying Normal
            distribution. It must be a float tensor of shape
            ``(d0, d1, ..., dn)``.
        min_value: Specifies the minimum value. It must be a float
            tensor of shape ``(d0, d1, ..., dn)``.
        max_value: Specifies the maximum value. It must be a float
            tensor of shape ``(d0, d1, ..., dn)``.
        generator: Specifies an optional random generator.

    Returns:
        A tensor of shape ``(d0, d1, ..., dn)`` filled with values
            sampled from a truncated log-Normal distribution.

    Raises:
        ValueError: if the ``mean``, ``std``, ``min_value`` and
            ``max_value`` parameters are not valid.

    Example usage:

    ```pycon

    >>> import torch
    >>> from startorch.random import trunc_log_normal
    >>> trunc_log_normal(
    ...     mean=torch.tensor([-1.0, 0.0, 1.0]),
    ...     std=torch.tensor([1.0, 3.0, 5.0]),
    ...     min_value=torch.tensor([0.0, 1.0, 2.0]),
    ...     max_value=torch.tensor([5.0, 10.0, 15.0]),
    ... )
    tensor([...])

    ```
    """
    if not shapes_are_equal((mean, std, min_value, max_value)):
        msg = (
            "Incorrect shapes. The shapes of all the input tensors must be equal: "
            f"mean={mean.shape}  std={std.shape}  min_value={min_value.shape}  "
            f"max_value={max_value.shape}"
        )
        raise ValueError(msg)
    if torch.any(std.le(0.0)):
        msg = "All the `std` values have to be greater than 0"
        raise ValueError(msg)
    if torch.any(max_value.lt(min_value)):
        msg = (
            "Found at least one value in `min_value` that is higher than its associated "
            "`max_value`"
        )
        raise ValueError(msg)
    distribution = LogNormal(loc=mean, scale=std)
    # Values are generated by using a uniform distribution and
    # then using the inverse CDF for the log-Normal distribution.
    # Get upper and lower CDF values
    lo, up = distribution.cdf(torch.stack([min_value, max_value]))
    return distribution.icdf(uniform(low=lo, high=up, generator=generator)).clamp(
        min=min_value, max=max_value
    )


def rand_trunc_normal(
    size: list[int] | tuple[int, ...],
    mean: float = 0.0,
    std: float = 1.0,
    min_value: float = -3.0,
    max_value: float = 3.0,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    r"""Create a tensor filled with values sampled from a truncated
    Normal distribution.

    Args:
        size: Specifies the tensor shape.
        mean: Specifies the mean of the Normal distribution.
        std: Specifies the standard deviation of the Normal
            distribution.
        min_value: Specifies the minimum value.
        max_value: Specifies the maximum value.
        generator: Specifies an optional random generator.

    Returns:
        A tensor filled with values sampled from a truncated Normal
            distribution.

    Raises:
        ValueError: if the ``min_value`` and  ``max_value`` parameters
            are not valid.

    Example usage:

    ```pycon
    >>> import torch
    >>> from startorch.random import rand_trunc_normal
    >>> rand_trunc_normal((2, 3), mean=1.0, std=2.0, min_value=-5.0, max_value=5.0)
    tensor([[...]])

    ```
    """
    if max_value < min_value:
        msg = f"`max_value` ({max_value}) has to be greater or equal to `min_value` ({min_value})"
        raise ValueError(msg)
    distribution = Normal(loc=mean, scale=std)
    # Values are generated by using a truncated uniform distribution and
    # then using the inverse CDF for the Normal distribution.
    # Get upper and lower cdf values
    lo, up = distribution.cdf(torch.tensor([min_value, max_value]))
    return distribution.icdf(rand_uniform(size, low=lo, high=up, generator=generator)).clamp(
        min=min_value, max=max_value
    )


def trunc_normal(
    mean: torch.Tensor,
    std: torch.Tensor,
    min_value: torch.Tensor,
    max_value: torch.Tensor,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    r"""Create a tensor filled with values sampled from a truncated
    Normal distribution.

    Args:
        mean: Specifies the mean of the Normal distribution.
            It must be a float tensor of shape ``(d0, d1, ..., dn)``.
        std: Specifies the standard deviation of the Normal
            distribution. It must be a float tensor of shape
            ``(d0, d1, ..., dn)``.
        min_value: Specifies the minimum value. It must be a float
            tensor of shape ``(d0, d1, ..., dn)``.
        max_value: Specifies the maximum value. It must be a float
            tensor of shape ``(d0, d1, ..., dn)``.
        generator: Specifies an optional random generator.

    Returns:
        A tensor of shape ``(d0, d1, ..., dn)`` filled with values
            sampled from a truncated Normal distribution.

    Raises:
        ValueError: if the ``min_value`` and  ``max_value`` parameters
            are not valid.

    Example usage:

    ```pycon
    >>> import torch
    >>> from startorch.random import trunc_normal
    >>> trunc_normal(
    ...     mean=torch.tensor([1.0, 0.0, -1.0]),
    ...     std=torch.tensor([1.0, 3.0, 5.0]),
    ...     min_value=torch.tensor([-5.0, -10.0, -15.0]),
    ...     max_value=torch.tensor([5.0, 10.0, 15.0]),
    ... )
    tensor([...])

    ```
    """
    if not shapes_are_equal((mean, std, min_value, max_value)):
        msg = (
            "Incorrect shapes. The shapes of all the input tensors must be equal: "
            f"mean={mean.shape}  std={std.shape}  min_value={min_value.shape}  "
            f"max_value={max_value.shape}"
        )
        raise ValueError(msg)
    if torch.any(std.le(0.0)):
        msg = "All the `std` values have to be greater than 0"
        raise ValueError(msg)
    if torch.any(max_value.lt(min_value)):
        msg = (
            "Found at least one value in `min_value` that is higher than its associated "
            "`max_value`"
        )
        raise ValueError(msg)
    distribution = Normal(loc=mean, scale=std)
    # Values are generated by using a uniform distribution and
    # then using the inverse CDF for the Normal distribution.
    # Get upper and lower CDF values
    lo, up = distribution.cdf(torch.stack([min_value, max_value]))
    return distribution.icdf(uniform(low=lo, high=up, generator=generator)).clamp(
        min=min_value, max=max_value
    )


def rand_uniform(
    size: list[int] | tuple[int, ...],
    low: float = 0.0,
    high: float = 1.0,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    r"""Create a tensor filled with values sampled from a uniform
    distribution.

    Args:
        size: Specifies the tensor shape.
        low: Specifies the minimum value (inclusive).
        high: Specifies the maximum value (exclusive).
        generator: Specifies an optional random generator.

    Returns:
        A tensor filled with values sampled from a uniform
            distribution.

    Raises:
        ValueError: if the ``low`` and  ``high`` parameters  are not
            valid.

    Example usage:

    ```pycon
    >>> import torch
    >>> from startorch.random import rand_uniform
    >>> rand_uniform((2, 3), low=-1.0, high=2.0)
    tensor([[...]])

    ```
    """
    if high < low:
        msg = f"`high` ({high}) has to be greater or equal to `low` ({low})"
        raise ValueError(msg)
    return (
        torch.rand(size=size, generator=generator).mul(high - low).add(low).clamp(min=low, max=high)
    )


def uniform(
    low: torch.Tensor, high: torch.Tensor, generator: torch.Generator | None = None
) -> torch.Tensor:
    r"""Create a tensor filled with values sampled from a uniform
    distribution.

    Unlike ``rand_uniform``, this function allows to sample values
    from different uniform distributions at the same time.
    The shape of the ``low`` and ``high`` tensors are used to infer
    the output size.

    Args:
        low: Specifies the minimum values (inclusive). It must be a
            float tensor of shape ``(d0, d1, ..., dn)``.
        high: Specifies the maximum values (exclusive). It must be a
            float tensor of shape ``(d0, d1, ..., dn)``.
        generator: Specifies an optional random generator.

    Returns:
        A tensor of shape ``(d0, d1, ..., dn)`` filled with values
            sampled from a uniform distribution where the minimum and
            maximum values are given as input.

    Raises:
        ValueError: if the input tensor shapes do not match or if at
            least one value in ``low`` tensor is higher than its
            associated high value in ``high``

    Example usage:

    ```pycon
    >>> import torch
    >>> from startorch.random import uniform
    >>> uniform(low=torch.tensor([-1.0, 0.0, 1.0]), high=torch.tensor([1.0, 3.0, 5.0]))
    tensor([...])

    ```
    """
    if low.shape != high.shape:
        msg = (
            f"Incorrect shapes. The shapes of all the input tensors must be equal: "
            f"low={low.shape}  high={high.shape})"
        )
        raise ValueError(msg)
    if torch.any(high.lt(low)):
        msg = (
            "Found at least one value in `low` that is higher than its associated "
            "value in `high`.`uniform` expects to return a [low, high) range"
        )
        raise ValueError(msg)
    return (
        torch.rand(size=high.shape, generator=generator)
        .mul(high - low)
        .add(low)
        .clamp(min=low, max=high)
    )


def rand_log_uniform(
    size: list[int] | tuple[int, ...],
    low: float,
    high: float,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    r"""Create a tensor filled with values sampled from a uniform
    distribution in the log space.

    Args:
        size: Specifies the tensor shape.
        low: Specifies the minimum value (inclusive).
            This value needs to be positive.
        high: Specifies the maximum value (exclusive).
        generator: Specifies an optional random generator.

    Returns:
        A tensor filled with values sampled from a uniform distribution
            in the log space.

    Raises:
        ValueError: if the ``low`` and  ``high`` parameters  are not
            valid.

    Example usage:

    ```pycon
    >>> from startorch.random import rand_log_uniform
    >>> rand_log_uniform((2, 3), low=0.1, high=1000.0)
    tensor([[...]])

    ```
    """
    if high < low:
        msg = f"`high` ({high}) has to be greater or equal to `low` ({low})"
        raise ValueError(msg)
    log_low = math.log(low)
    return (
        torch.rand(size=size, generator=generator)
        .mul(math.log(high) - log_low)
        .add(log_low)
        .exp()
        .clamp(min=low, max=high)
    )


def log_uniform(
    low: torch.Tensor, high: torch.Tensor, generator: torch.Generator | None = None
) -> torch.Tensor:
    r"""Create a tensor filled with values sampled from a uniform
    distribution in the log space.

    Args:
        low: Specifies the minimum values (inclusive). It must be a
            float tensor of shape ``(d0, d1, ..., dn)``.
        high: Specifies the maximum values (exclusive). It must be a
            float tensor of shape ``(d0, d1, ..., dn)``.
        generator: Specifies an optional random generator.

    Returns:
        A tensor of shape ``(d0, d1, ..., dn)`` filled with values
            sampled from a uniform distribution in the log space where
            the minimum and maximum values are given as input.

    Raises:
        ValueError: if the ``low`` and  ``high`` parameters  are not
            valid.

    Example usage:

    ```pycon
    >>> from startorch.random import log_uniform
    >>> log_uniform(low=torch.tensor([0.01, 0.1, 1.0]), high=torch.tensor([1.0, 10.0, 100.0]))
    tensor([...])

    ```
    """
    if low.shape != high.shape:
        msg = (
            "Incorrect shapes. The shapes of all the input tensors must be equal: "
            f"low={low.shape}  high={high.shape}"
        )
        raise ValueError(msg)
    if torch.any(high.lt(low)):
        msg = (
            "Found at least one value in `low` that is higher than its associated "
            "value in `high`. `log_uniform` expects to return a [low, high) range"
        )
        raise ValueError(msg)
    log_low = low.log()
    return (
        torch.rand(size=high.shape, generator=generator)
        .mul(high.log() - log_low)
        .add(log_low)
        .exp()
        .clamp(min=low, max=high)
    )


def rand_asinh_uniform(
    size: list[int] | tuple[int, ...],
    low: float,
    high: float,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    r"""Create a tensor filled with values sampled from a uniform
    distribution in the inverse hyperbolic sine space.

    Args:
        size: Specifies the tensor shape.
        low: Specifies the minimum value (inclusive).
            This value needs to be positive.
        high: Specifies the maximum value (exclusive).
        generator: Specifies an optional random generator.

    Returns:
        A tensor filled with values sampled from a uniform distribution
            in the inverse hyperbolic sine space.

    Raises:
        ValueError: if the ``low`` and  ``high`` parameters  are not
            valid.

    Example usage:

    ```pycon
    >>> from startorch.random import rand_asinh_uniform
    >>> rand_asinh_uniform((2, 3), low=-1000.0, high=1000.0)
    tensor([[...]])

    ```
    """
    if high < low:
        msg = f"`high` ({high}) has to be greater or equal to `low` ({low})"
        raise ValueError(msg)
    log_low = math.asinh(low)
    return (
        torch.rand(size=size, generator=generator)
        .mul(math.asinh(high) - log_low)
        .add(log_low)
        .sinh()
        .clamp(min=low, max=high)
    )


def asinh_uniform(
    low: torch.Tensor, high: torch.Tensor, generator: torch.Generator | None = None
) -> torch.Tensor:
    r"""Create a tensor filled with values sampled from a uniform
    distribution in the inverse hyperbolic sine space.

    Args:
        low: Specifies the minimum values (inclusive). It must be a
            float tensor of shape ``(d0, d1, ..., dn)``.
        high: Specifies the maximum values (exclusive). It must be a
            float tensor of shape ``(d0, d1, ..., dn)``.
        generator: Specifies an optional random generator.

    Returns:
        A tensor of shape ``(d0, d1, ..., dn)`` filled with values
            sampled from a uniform distribution in the inverse
            hyperbolic sine space where the minimum and maximum values
            are given as input.

    Raises:
        ValueError: if the ``low`` and  ``high`` parameters  are not
            valid.

    Example usage:

    ```pycon
    >>> from startorch.random import asinh_uniform
    >>> asinh_uniform(
    ...     low=torch.tensor([-10.0, 0.0, 1.0]),
    ...     high=torch.tensor([1.0, 10.0, 100.0]),
    ... )
    tensor([...])

    ```
    """
    if low.shape != high.shape:
        msg = (
            "Incorrect shapes. The shapes of all the input tensors must be equal: "
            f"low={low.shape}  high={high.shape}"
        )
        raise ValueError(msg)
    if torch.any(high.lt(low)):
        msg = (
            "Found at least one value in `low` that is higher than its associated "
            "value in `high`. `asinh_uniform` expects to return a [low, high) range"
        )
        raise ValueError(msg)
    log_low = low.asinh()
    return (
        torch.rand(size=high.shape, generator=generator)
        .mul(high.asinh() - log_low)
        .add(log_low)
        .sinh()
        .clamp(min=low, max=high)
    )
