from __future__ import annotations

__all__ = ["HypercubeExampleGenerator", "create_hypercube_examples"]

import torch
from redcat import BatchDict, BatchedTensor

from startorch import constants as ct
from startorch.example.base import BaseExampleGenerator


class HypercubeExampleGenerator(BaseExampleGenerator[BatchedTensor]):
    r"""Implements a classification example generator.

    The data are generated by using a hypercube. The targets are some
    vertices of the hypercube. Each input feature is a 1-hot
    representation of the target plus a Gaussian noise. These data can
    be used for a multi-class classification task.

    Args:
    ----
        num_classes (int, optional): Specifies the number of classes.
            Default: 50
        feature_size (int, optional): Specifies the feature size.
            The feature size has to be greater than the number of
            classes. Default: ``64``
        noise_std (float, optional): Specifies the standard deviation
            of the Gaussian noise. Default: ``0.2``

    Raises:
    ------
        ValueError if one of the parameters is not valid.


    Example usage:

    .. code-block:: pycon

        >>> from startorch.example import Hypercube
        >>> generator = Hypercube(num_classes=5, feature_size=6)
        >>> generator
        HypercubeExampleGenerator(num_classes=5, feature_size=6, noise_std=0.2)
        >>> batch = generator.generate(batch_size=10)
        >>> batch
        BatchDict(
          (target): tensor([...], batch_dim=0)
          (feature): tensor([[...]], batch_dim=0)
        )
    """

    def __init__(
        self,
        num_classes: int = 50,
        feature_size: int = 64,
        noise_std: float = 0.2,
    ) -> None:
        if num_classes < 1:
            raise ValueError(f"The number of classes ({num_classes}) has to be greater than 0")
        self._num_classes = int(num_classes)

        if feature_size < num_classes:
            raise ValueError(
                f"The feature dimension ({feature_size:,}) has to be greater or equal to the "
                f"number of classes ({num_classes:,})"
            )
        self._feature_size = int(feature_size)

        if noise_std < 0:
            raise ValueError(
                f"The standard deviation of the Gaussian noise ({noise_std}) has to be "
                "greater or equal than 0"
            )
        self._noise_std = float(noise_std)

    def __repr__(self) -> str:
        return (
            f"{self.__class__.__qualname__}("
            f"num_classes={self._num_classes:,}, "
            f"feature_size={self._feature_size:,}, "
            f"noise_std={self._noise_std:,})"
        )

    @property
    def num_classes(self) -> int:
        r"""``int``: The number of classes when the data are created."""
        return self._num_classes

    @property
    def feature_size(self) -> int:
        r"""``int``: The feature size when the data are created."""
        return self._feature_size

    @property
    def noise_std(self) -> float:
        r"""``float``: The standard deviation of the Gaussian noise."""
        return self._noise_std

    def generate(
        self, batch_size: int = 1, rng: torch.Generator | None = None
    ) -> BatchDict[BatchedTensor]:
        return create_hypercube_examples(
            num_examples=batch_size,
            num_classes=self._num_classes,
            feature_size=self._feature_size,
            noise_std=self._noise_std,
            generator=rng,
        )


def create_hypercube_examples(
    num_examples: int = 1000,
    num_classes: int = 50,
    feature_size: int = 64,
    noise_std: float = 0.2,
    generator: torch.Generator | None = None,
) -> BatchDict[BatchedTensor]:
    r"""Create a toy classification dataset based on hypercube vertex
    structure.

    The data are generated by using a hypercube. The targets are some
    vertices of the hypercube. Each input feature is a 1-hot
    representation of the target plus a Gaussian noise. These data can
    be used for a multi-class classification task.

    Args:
    ----
        num_examples (int, optional): Specifies the number of examples.
            Default: ``1000``
        num_classes (int, optional): Specifies the number of classes.
            Default: 50
        feature_size (int, optional): Specifies the feature size.
            The feature size has to be greater than the number of
            classes. Default: ``64``
        noise_std (float, optional): Specifies the standard deviation
            of the Gaussian noise. Default: ``0.2``
        generator (``torch.Generator`` or ``None``, optional):
            Specifies an optional random generator. Default: ``None``

    Returns:
    -------
        dict: A dictionary with two keys:
            - ``'input'``: a ``torch.Tensor`` of type float and
                shape ``(num_examples, feature_size)``. This
                tensor represents the input features.
            - ``'target'``: a ``torch.Tensor`` of type long and
                shape ``(num_examples,)``. This tensor represents
                the targets.

    Raises:
    ------
        RuntimeError if one of the parameters is not valid.

    Example usage:

    .. code-block:: pycon

        >>> from startorch.example.hypercube import create_hypercube_examples
        >>> batch = create_hypercube_examples(num_examples=10, num_classes=5, feature_size=10)
        >>> batch
        BatchDict(
          (target): tensor([...], batch_dim=0)
          (feature): tensor([[...]], batch_dim=0)
        )
    """

    if num_examples < 1:
        raise RuntimeError(f"The number of examples ({num_examples}) has to be greater than 0")
    if num_classes < 1:
        raise RuntimeError(f"The number of classes ({num_classes}) has to be greater than 0")
    if feature_size < num_classes:
        raise RuntimeError(
            f"The feature dimension ({feature_size:,}) has to be greater or equal to the "
            f"number of classes ({num_classes:,})"
        )
    if noise_std < 0:
        raise RuntimeError(
            f"The standard deviation of the Gaussian noise ({noise_std}) has to be "
            "greater or equal than 0"
        )
    # Generate the target of each example.
    targets = torch.randint(0, num_classes, (num_examples,), generator=generator)
    # Generate the features. Each class should be a vertex of the hyper-cube
    # plus Gaussian noise.
    features = torch.randn(num_examples, feature_size, generator=generator).mul(noise_std)
    features.scatter_add_(1, targets.view(num_examples, 1), torch.ones(num_examples, 1))
    return BatchDict({ct.TARGET: BatchedTensor(targets), ct.FEATURE: BatchedTensor(features)})
