from typing import Any

import numpy as np

from qxmt.datasets.generate.linear import (
    generate_linear_regression_data,
    generate_linear_separable_data,
)


class GeneratedDataLoader:
    """
    This class loads the generated data and label.
    The loaded data and label are returned as numpy arrays (X and y).
    The dataset is generated by the specified task and method.

    Supported task types:
    - classification
    - regression

    Supported generate methods:
    - linear

    Examples:
        >>> loader = GeneratedDataLoader(task_type="classification", generate_method="linear", random_seed=42)
        >>> X, y = loader.load()
        >>> loader = GeneratedDataLoader(task_type="regression", generate_method="linear", random_seed=42)
        >>> X, y = loader.load()
    """

    def __init__(
        self,
        task_type: str,
        generate_method: str,
        random_seed: int,
        params: dict[str, Any] = {},
    ) -> None:
        """Initialize the GeneratedDataLoader.

        Args:
            task_type (str): dataset task type
            generate_method (str): generate dataset method
            random_seed (int): random seed for reproducibility
            params (dict[str, Any]]): additional parameters for dataset generation
        """
        self.task_type = task_type
        self.generate_method = generate_method
        self.random_seed = random_seed
        self.params = params

    def load(self) -> tuple[np.ndarray, np.ndarray]:
        """Load dummy dataset. The dataset is generated by the specified task and method.

        Args:
            task_type (str): dataset task type. "classification" or "regression" is supported.
            generate_method (str): generate dataset method. "linear" is supported.
            random_seed (int): random seed for reproducibility.
            params (dict[str, Any]]): additional parameters for dataset generation.

        Raises:
            ValueError: not supported task_type or generate_method

        Returns:
            tuple[np.ndarray, np.ndarray]: generated data and label
        """
        match (self.task_type, self.generate_method):
            case ("classification", "linear"):
                X, y = generate_linear_separable_data(
                    n_samples=self.params.get("n_samples", 100),
                    n_features=self.params.get("n_features", 2),
                    n_classes=self.params.get("n_classes", 2),
                    noise=self.params.get("noise", 0.1),
                    scale=self.params.get("scale", 1.0),
                    random_seed=self.random_seed,
                )
            case ("regression", "linear"):
                X, y = generate_linear_regression_data(
                    n_samples=self.params.get("n_samples", 100),
                    n_features=self.params.get("n_features", 2),
                    noise=self.params.get("noise", 0.1),
                    scale=self.params.get("scale", 1.0),
                    random_seed=self.random_seed,
                )
            case _:
                raise ValueError(
                    f'Unsupported task_type="{self.task_type}" or generate_method="{self.generate_method}"'
                )

        return X, y
