from __future__ import annotations

__all__ = ["SineWavePeriodicSequenceGenerator"]

import math

from coola.utils.format import str_indent, str_mapping
from redcat import BatchedTensorSeq
from torch import Generator

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


class SineWavePeriodicSequenceGenerator(BasePeriodicSequenceGenerator):
    r"""Implements a class to generate 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 (``BasePeriodicSequenceGenerator`` or dict): Specifies a
            sequence generator (or its configuration) to generate the
            sequence values.
        phase (``BasePeriodicSequenceGenerator`` or dict): Specifies a
            sequence generator (or its configuration) to generate the
            phase values.
        amplitude (``BasePeriodicSequenceGenerator`` or dict): Specifies a
            sequence generator (or its configuration) to generate the
            amplitude values.

    Example usage:

    .. code-block:: 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([[...]], batch_dim=0, seq_dim=1)
    """

    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: Generator | None = None
    ) -> BatchedTensorSeq:
        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
                )
            )
        )
