r"""Contain the implementation of a sequence generator where the values
are sampled using a Markov chain/process."""

from __future__ import annotations

__all__ = ["MarkovChainSequenceGenerator", "make_markov_chain"]


import torch

from startorch.sequence.base import BaseSequenceGenerator
from startorch.tensor import BaseTensorGenerator, setup_tensor_generator
from startorch.transition import BaseTransitionGenerator, setup_transition_generator


class MarkovChainSequenceGenerator(BaseSequenceGenerator):
    r"""Implement a sequence generator to generate a batch of sequences
    where the data are generated by using a Markov chain/process.

    Args:
        transition: The transition matrix of the Markov chain.
        init: The initial probabilities. If ``None``, the initial
            values are sampled with a uniform probability.

    Example usage:

    ```pycon

    >>> import torch
    >>> from startorch.sequence import MarkovChain
    >>> generator = MarkovChain(transition=torch.rand(5, 5))
    >>> generator
    MarkovChainSequenceGenerator(transition=(5, 5), init=None)
    >>> generator.generate(seq_len=6, batch_size=2)
    tensor([[...]])

    ```
    """

    def __init__(
        self,
        transition: torch.Tensor,
        init: torch.Tensor | None = None,
    ) -> None:
        self._transition = transition
        self._init = init

    def __repr__(self) -> str:
        init = self._init if self._init is None else tuple(self._init.shape)
        return f"{self.__class__.__qualname__}(transition={tuple(self._transition.shape)}, init={init})"

    def generate(
        self, seq_len: int, batch_size: int = 1, rng: torch.Generator | None = None
    ) -> torch.Tensor:
        return make_markov_chain(
            batch_size=batch_size,
            seq_len=seq_len,
            transition=self._transition,
            init=self._init,
            rng=rng,
        )

    @classmethod
    def create_from_generators(
        cls,
        n: int,
        transition: BaseTransitionGenerator | dict,
        init: BaseTensorGenerator | dict | None = None,
        rng: torch.Generator | None = None,
    ) -> MarkovChainSequenceGenerator:
        r"""Create the Markov chain sequence generator where the
        transition matrix and initial probabilities are generated from
        generators.

        Args:
            n: The number of states in the Markov chain i.e. the size
                of the transition matrix.
            transition: The transition matrix generator or its
                configuration.
            init: The initial probabilities generator or its
                configuration.
            rng: An optional random number generator.

        Returns:
            The instantiated Markov chain sequence generator.

        Example usage:

        ```pycon

        >>> import torch
        >>> from startorch.sequence import MarkovChain
        >>> generator = MarkovChain.create_from_generators(
        ...     n=6,
        ...     transition={'_target_': 'startorch.transition.Diagonal'},
        ...     init={'_target_': 'startorch.tensor.RandUniform'},
        ... )
        >>> generator
        MarkovChainSequenceGenerator(transition=(6, 6), init=(6,))

        ```
        """
        transition = setup_transition_generator(transition)
        if init is not None:
            init = setup_tensor_generator(init).generate(size=(n,), rng=rng)
        return cls(transition=transition.generate(n=n, rng=rng), init=init)


def make_markov_chain(
    batch_size: int,
    seq_len: int,
    transition: torch.Tensor,
    init: torch.Tensor | None = None,
    rng: torch.Generator | None = None,
) -> torch.Tensor:
    r"""Generate a batch of sequences where the data are generated by
    using a Markov chain/process.

    Args:
        batch_size: The batch size.
        seq_len: The sequence length.
        transition: The transition matrix of the Markov chain.
        init: The initial probabilities. If ``None``, the initial
            values are sampled with a uniform probability.
        rng: An optional random number generator.

    Returns:
        The generated batch of sequences.

    Example usage:

    ```pycon

    >>> import torch
    >>> from startorch.sequence.markov import make_markov_chain
    >>> out = make_markov_chain(batch_size=2, seq_len=6, transition=torch.rand(5, 5))
    >>> out
    tensor([[...]])

    ```
    """
    out = torch.zeros(batch_size, seq_len + 1, dtype=torch.long)
    if init is None:
        init = torch.ones(batch_size, transition.shape[0])
    for i in range(seq_len):
        weights = init if i == 0 else transition[out[:, i]]
        out[:, i + 1] = torch.multinomial(
            weights, num_samples=1, replacement=True, generator=rng
        ).flatten()
    return out[:, 1:]
