r"""Contain an example generator to generate data using the sparse
uncorrelated features."""

from __future__ import annotations

__all__ = [
    "LinearRegressionExampleGenerator",
    "make_linear_regression",
    "get_uniform_weights",
]

from typing import TYPE_CHECKING

import torch

from startorch import constants as ct
from startorch.example.base import BaseExampleGenerator
from startorch.random import rand_normal, rand_uniform
from startorch.utils.conversion import to_tensor
from startorch.utils.seed import get_torch_generator
from startorch.utils.validation import check_num_examples, check_std

if TYPE_CHECKING:
    from collections.abc import Sequence


class LinearRegressionExampleGenerator(BaseExampleGenerator):
    r"""Implement a regression example generator where the data are
    generated with an underlying linear model.

    The implementation is based on
    https://scikit-learn.org/stable/modules/generated/sklearn.datasets.make_regression.html

    Args:
        weights: Specifies the linear weights in the underlying linear
            model. It must be a float tensor of shape
            ``(feature_size,)``.
        bias: Specifies the bias term in the underlying linear model.
        noise_std: Specifies the standard deviation of the Gaussian
            noise.

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

    Example usage:

    ```pycon
    >>> from startorch.example import LinearRegression
    >>> generator = LinearRegression.create_uniform_weights()
    >>> generator
    LinearRegressionExampleGenerator(feature_size=100, bias=0.0, noise_std=0.0)
    >>> batch = generator.generate(batch_size=10)
    >>> batch
    {'target': tensor([...]), 'feature': tensor([[...]])}

    ```
    """

    def __init__(
        self,
        weights: torch.Tensor | Sequence[float],
        bias: float = 0.0,
        noise_std: float = 0.0,
    ) -> None:
        self._weights = to_tensor(weights)
        self._bias = bias

        check_std(noise_std, "noise_std")
        self._noise_std = float(noise_std)

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

    @property
    def bias(self) -> float:
        r"""The bias of the underlying linear model."""
        return self._bias

    @property
    def feature_size(self) -> int:
        r"""The feature size."""
        return self._weights.shape[0]

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

    @property
    def weights(self) -> torch.Tensor:
        r"""``torch.Tensor``: The weights of the underlying linear
        model."""
        return self._weights

    def generate(
        self, batch_size: int = 1, rng: torch.Generator | None = None
    ) -> dict[str, torch.Tensor]:
        return make_linear_regression(
            num_examples=batch_size,
            weights=self._weights,
            bias=self._bias,
            noise_std=self._noise_std,
            generator=rng,
        )

    @classmethod
    def create_uniform_weights(
        cls,
        feature_size: int = 100,
        informative_feature_size: int = 10,
        bias: float = 0.0,
        noise_std: float = 0.0,
        random_seed: int = 17532042831661189422,
    ) -> LinearRegressionExampleGenerator:
        # TODO(thibaut): add missing documentation # noqa: TD003
        return cls(
            weights=get_uniform_weights(
                feature_size=feature_size,
                informative_feature_size=informative_feature_size,
                generator=get_torch_generator(random_seed),
            ),
            bias=bias,
            noise_std=noise_std,
        )


def make_linear_regression(
    weights: torch.Tensor,
    bias: float = 0.0,
    num_examples: int = 100,
    noise_std: float = 0.0,
    generator: torch.Generator | None = None,
) -> dict[str, torch.Tensor]:
    r"""Generate a regression dataset where the data are generated with
    an underlying linear model.

    The features are sampled from a Normal distribution.
    Then, the targets are generated by applying a random linear
    regression model.

    Args:
        weights: Specifies the linear weights in the underlying linear
            model. It must be a float tensor of shape
            ``(feature_size,)``.
        bias: Specifies the bias term in the underlying linear model.
        num_examples: Specifies the number of examples to generate.
        noise_std: Specifies the standard deviation of the Gaussian
            noise.
        generator: Specifies an optional random generator.

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

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

    Example usage:

    ```pycon
    >>> from startorch.example import make_linear_regression
    >>> batch = make_linear_regression(weights=torch.rand(10), num_examples=10)
    >>> batch
    {'target': tensor([...]), 'feature': tensor([[...]])}

    ```
    """
    check_num_examples(num_examples)
    check_std(noise_std, "noise_std")
    feature_size = weights.shape[0]
    features = rand_normal(size=(num_examples, feature_size), generator=generator)
    targets = torch.mm(features, weights.view(feature_size, 1)) + bias
    if noise_std > 0.0:
        features += rand_normal(
            size=(num_examples, feature_size), std=noise_std, generator=generator
        )
    return {ct.TARGET: targets.flatten(), ct.FEATURE: features}


def get_uniform_weights(
    feature_size: int,
    informative_feature_size: int,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    """Generate the weights of the linear combination used to generate
    the targets from the features.

    The weights of the informative features are sampled from a uniform
    distribution. The other weights are set to 0.
    This function was designed to be used with
    ``make_normal_regression``.

    Args:
        feature_size: Specifies the feature size i.e. the number of
            features.
        informative_feature_size: Specifies the number of informative
            features.
        generator: Specifies an optional random generator.

    Returns:
        The generated weights as a float tensor of shape
            ``(feature_size,)``.

    Example usage:

    ```pycon
    >>> from startorch.example.regression import get_uniform_weights
    >>> weights = get_uniform_weights(feature_size=10, informative_feature_size=5)
    >>> weights
    tensor([...])

    ```
    """
    informative_feature_size = min(feature_size, informative_feature_size)
    weights = torch.zeros(feature_size)
    weights[:informative_feature_size] = 100 * rand_uniform(
        size=(informative_feature_size,), generator=generator
    )
    permutation = torch.randperm(feature_size, generator=generator)
    return weights[permutation]
