r"""Contain a periodic sequence generator that generates periodic
sequence by sampling values with a sine wave pattern."""

from __future__ import annotations

__all__ = ["SineWavePeriodicSequenceGenerator"]

import math
from typing import TYPE_CHECKING

from coola.utils.format import str_indent, str_mapping

from startorch.periodic.sequence import (
    BasePeriodicSequenceGenerator,
    setup_periodic_sequence_generator,
)

if TYPE_CHECKING:
    import torch


class SineWavePeriodicSequenceGenerator(BasePeriodicSequenceGenerator):
    r"""Implement a periodic sequence generator that generates periodic
    sequence by sampling values with a sine wave pattern.

    The sequences are generated by using the following formula:

    ``output = amplitude * sin(2 * pi * frequency * value + phase)``

    Args:
        value: Specifies a sequence generator (or its configuration)
            to generate the sequence values.
        phase: Specifies a sequence generator (or its configuration)
            to generate the phase values.
        amplitude: Specifies a sequence generator (or its
            configuration) to generate the amplitude values.

    Example usage:

    ```pycon
    >>> import torch
    >>> from startorch.periodic.sequence import Repeat, SineWave
    >>> from startorch.sequence import RandUniform
    >>> generator = SineWave(
    ...     value=Repeat(RandUniform(low=-1.0, high=1.0)),
    ...     phase=Repeat(RandUniform(low=-1.0, high=1.0)),
    ...     amplitude=Repeat(RandUniform(low=-1.0, high=1.0)),
    ... )
    >>> generator
    SineWavePeriodicSequenceGenerator(
      (value): RepeatPeriodicSequenceGenerator(
          (generator): RandUniformSequenceGenerator(low=-1.0, high=1.0, feature_size=(1,))
        )
      (phase): RepeatPeriodicSequenceGenerator(
          (generator): RandUniformSequenceGenerator(low=-1.0, high=1.0, feature_size=(1,))
        )
      (amplitude): RepeatPeriodicSequenceGenerator(
          (generator): RandUniformSequenceGenerator(low=-1.0, high=1.0, feature_size=(1,))
        )
    )
    >>> generator.generate(seq_len=12, period=4, batch_size=4)
    tensor([[...]])

    ```
    """

    def __init__(
        self,
        value: BasePeriodicSequenceGenerator | dict,
        phase: BasePeriodicSequenceGenerator | dict,
        amplitude: BasePeriodicSequenceGenerator | dict,
    ) -> None:
        super().__init__()
        self._value = setup_periodic_sequence_generator(value)
        self._phase = setup_periodic_sequence_generator(phase)
        self._amplitude = setup_periodic_sequence_generator(amplitude)

    def __repr__(self) -> str:
        args = str_indent(
            str_mapping({"value": self._value, "phase": self._phase, "amplitude": self._amplitude})
        )
        return f"{self.__class__.__qualname__}(\n  {args}\n)"

    def generate(
        self, seq_len: int, period: int, batch_size: int = 1, rng: torch.Generator | None = None
    ) -> torch.Tensor:
        return (
            self._value.generate(seq_len=seq_len, period=period, batch_size=batch_size, rng=rng)
            .mul(2.0 * math.pi / period)
            .add(
                self._phase.generate(seq_len=seq_len, period=period, batch_size=batch_size, rng=rng)
            )
            .sin()
            .mul(
                self._amplitude.generate(
                    seq_len=seq_len, period=period, batch_size=batch_size, rng=rng
                )
            )
        )
