from __future__ import annotations

__all__ = ["RandWienerProcessSequenceGenerator", "wiener_process"]

import math

import torch
from redcat import BatchedTensorSeq
from torch import Generator, Tensor

from startorch.sequence.base import BaseSequenceGenerator


class RandWienerProcessSequenceGenerator(BaseSequenceGenerator):
    r"""Implements 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 (float, optional): Specifies the time
            step size. Default: ``1.0``

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

    Example usage:

    .. code-block:: pycon

        >>> from startorch.sequence import RandWienerProcess
        >>> generator = RandWienerProcess()
        >>> generator
        RandWienerProcessSequenceGenerator(time_step_size=1.0)
        >>> generator.generate(seq_len=6, batch_size=2)
        tensor([[...]], batch_dim=0, seq_dim=1)
    """

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

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

    def generate(
        self, seq_len: int, batch_size: int = 1, rng: Generator | None = None
    ) -> BatchedTensorSeq:
        return BatchedTensorSeq(
            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: Generator | None = None,
) -> Tensor:
    r"""Creates a batch of sequences generated by a Wiener process.

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

    Args:
    ----
        seq_len (int): Specifies the sequence length.
        step_size (float, optional): Specifies the time
            step size. Default: ``1.0``
        batch_size (int, optional): Specifies the batch size.
            Default: ``1``
        generator (``torch.Generator`` or None): Specifies an optional
            random generator. Default: ``None``

    Returns:
    -------
        ``torch.Tensor`` of shape ``(batch size, sequence length)``
            and type float: A batch of Wiener processes.

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

    Example usage:

    .. code-block:: pycon

        >>> from startorch.sequence.wiener import wiener_process
        >>> wiener_process(seq_len=12, batch_size=4)
        tensor([[...]])
    """
    if step_size < 0:
        raise ValueError(f"step_size has to be greater than 0 (received: {step_size})")
    increments = torch.randn((batch_size, seq_len), generator=generator).mul(math.sqrt(step_size))
    increments[:, 0] = 0
    return torch.cumsum(increments, dim=1)
