from __future__ import annotations

__all__ = ["SineWaveSequenceGenerator"]

import math

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

from startorch.sequence.base import BaseSequenceGenerator, setup_sequence_generator


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

    Example usage:

    .. code-block:: pycon

        >>> import torch
        >>> from startorch.sequence import Arange, SineWave, RandUniform, Constant, RandLogUniform
        >>> generator = SineWave(
        ...     value=Arange(),
        ...     frequency=Constant(RandLogUniform(low=0.01, high=0.1)),
        ...     phase=Constant(RandUniform(low=-1.0, high=1.0)),
        ...     amplitude=Constant(RandLogUniform(low=0.1, high=1.0)),
        ... )
        >>> generator
        SineWaveSequenceGenerator(
          (value): ArangeSequenceGenerator(feature_size=(1,))
          (frequency): ConstantSequenceGenerator(
              (sequence): RandLogUniformSequenceGenerator(low=0.01, high=0.1, feature_size=(1,))
            )
          (phase): ConstantSequenceGenerator(
              (sequence): RandUniformSequenceGenerator(low=-1.0, high=1.0, feature_size=(1,))
            )
          (amplitude): ConstantSequenceGenerator(
              (sequence): RandLogUniformSequenceGenerator(low=0.1, high=1.0, feature_size=(1,))
            )
        )
        >>> generator.generate(seq_len=12, batch_size=4)
        tensor([[...]], batch_dim=0, seq_dim=1)
    """

    def __init__(
        self,
        value: BaseSequenceGenerator | dict,
        frequency: BaseSequenceGenerator | dict,
        phase: BaseSequenceGenerator | dict,
        amplitude: BaseSequenceGenerator | dict,
    ) -> None:
        super().__init__()
        self._value = setup_sequence_generator(value)
        self._frequency = setup_sequence_generator(frequency)
        self._phase = setup_sequence_generator(phase)
        self._amplitude = setup_sequence_generator(amplitude)

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

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