r"""Contain the implementation of sequence generators where the values
are sampled from a Wiener process."""

from __future__ import annotations

__all__ = ["RandWienerProcessSequenceGenerator", "wiener_process"]

import math

import torch

from startorch.sequence.base import BaseSequenceGenerator


class RandWienerProcessSequenceGenerator(BaseSequenceGenerator):
    r"""Implement a sequence generator to generate sequences where the
    values are generated by a Wiener process.

    Useful link: https://en.wikipedia.org/wiki/Wiener_process

    Args:
        step_size: Specifies the time step size.

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

    Example usage:

    ```pycon
    >>> from startorch.sequence import RandWienerProcess
    >>> generator = RandWienerProcess()
    >>> generator
    RandWienerProcessSequenceGenerator(step_size=1.0)
    >>> generator.generate(seq_len=6, batch_size=2)
    tensor([[...]])

    ```
    """

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

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

    def generate(
        self, seq_len: int, batch_size: int = 1, rng: torch.Generator | None = None
    ) -> torch.Tensor:
        return wiener_process(
            seq_len=seq_len,
            batch_size=batch_size,
            step_size=self._step_size,
            generator=rng,
        ).unsqueeze(dim=2)


def wiener_process(
    seq_len: int,
    step_size: float = 1.0,
    batch_size: int = 1,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    r"""Create a batch of sequences generated by a Wiener process.

    Useful link: https://en.wikipedia.org/wiki/Wiener_process

    Args:
        seq_len: Specifies the sequence length.
        step_size: Specifies the time step size.
        batch_size: Specifies the batch size.
        generator: Specifies an optional random generator.

    Returns:
        A batch of sequences generated with a Wiener process.
            It is a tensor of type float and shape
            ``(batch size, sequence length)``.

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

    Example usage:

    ```pycon
    >>> from startorch.sequence.wiener import wiener_process
    >>> wiener_process(seq_len=12, batch_size=4)
    tensor([[...]])

    ```
    """
    if step_size < 0:
        msg = f"step_size has to be greater than 0 (received: {step_size})"
        raise ValueError(msg)
    increments = torch.randn((batch_size, seq_len), generator=generator).mul(math.sqrt(step_size))
    increments[:, 0] = 0
    return torch.cumsum(increments, dim=1)
