from __future__ import annotations

__all__ = ["LinearSequenceGenerator"]


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

from startorch.sequence.base import BaseSequenceGenerator, setup_sequence_generator


class LinearSequenceGenerator(BaseSequenceGenerator):
    r"""Implements a class to generate sequence by sampling values with a
    linear pattern.

    The sequences are generated by using the following formula:

    ``output = slope * value + intercept``

    Args:
    ----
        value (``BaseSequenceGenerator`` or dict): Specifies a
            sequence generator (or its configuration) to generate the
            sequence values. The linear transformation is applied
            on these values.
        slope (``BaseSequenceGenerator`` or dict): Specifies a sequence
            generator (or its configuration) to generate the slope
            values.
        intercept (``BaseSequenceGenerator`` or dict): Specifies a
            sequence generator (or its configuration) to generate the
            intercept values.

    Example usage:

    .. code-block:: pycon

        >>> from startorch.sequence import Linear, RandUniform
        >>> generator = Linear(
        ...     value=RandUniform(low=-1.0, high=1.0),
        ...     slope=RandUniform(low=1.0, high=2.0),
        ...     intercept=RandUniform(low=-10.0, high=-5.0),
        ... )
        >>> generator
        LinearSequenceGenerator(
          (value): RandUniformSequenceGenerator(low=-1.0, high=1.0, feature_size=(1,))
          (slope): RandUniformSequenceGenerator(low=1.0, high=2.0, feature_size=(1,))
          (intercept): RandUniformSequenceGenerator(low=-10.0, high=-5.0, feature_size=(1,))
        )
        >>> generator.generate(seq_len=6, batch_size=2)
        tensor([[...]], batch_dim=0, seq_dim=1)
    """

    def __init__(
        self,
        value: BaseSequenceGenerator | dict,
        slope: BaseSequenceGenerator | dict,
        intercept: BaseSequenceGenerator | dict,
    ) -> None:
        super().__init__()
        self._value = setup_sequence_generator(value)
        self._slope = setup_sequence_generator(slope)
        self._intercept = setup_sequence_generator(intercept)

    def __repr__(self) -> str:
        args = str_indent(
            str_mapping({"value": self._value, "slope": self._slope, "intercept": self._intercept})
        )
        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._slope.generate(seq_len=seq_len, batch_size=batch_size, rng=rng))
            .add(self._intercept.generate(seq_len=seq_len, batch_size=batch_size, rng=rng))
        )
