from typing import Any

import numpy as np

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


def load_dummy_dataset(
    task_type: str, generate_method: str, random_seed: int, params: dict[str, Any] = {}
) -> 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 (task_type, generate_method):
        case ("classification", "linear"):
            X, y = generate_linear_separable_data(
                n_samples=params.get("n_samples", 100),
                n_features=params.get("n_features", 2),
                n_classes=params.get("n_classes", 2),
                noise=params.get("noise", 0.1),
                scale=params.get("scale", 1.0),
                random_seed=random_seed,
            )
        case ("regression", "linear"):
            X, y = generate_linear_regression_data(
                n_samples=params.get("n_samples", 100),
                n_features=params.get("n_features", 2),
                noise=params.get("noise", 0.1),
                scale=params.get("scale", 1.0),
                random_seed=random_seed,
            )
        case _:
            raise ValueError(f'Unsupported task_type="{task_type}" or generate_method="{generate_method}"')

    return X, y
