from typing import Any, Dict, Union

import pandas as pd

from dataset_hub._core.get_data import get_data

task_type = "classification"
# TODO think about **params in public functions. User dont know about this params.
# How it solve pandas?


def _get_data(dataset_name: str, **params) -> Union[Any, Dict[str, Any]]:
    """
    Load a dataset for the 'classification' task.

    This function provides a simple interface for users to load datasets
    relevant to the classification task.

    Args:
        dataset_name (str): Name of the dataset to load. Defaults to "titanic".
        **params: Additional parameters passed to the underlying loader (optional).

    Returns:
        Dict[str, Any]: A dictionary of tables (e.g., pandas DataFrames) corresponding
        to the dataset.
    """
    return get_data(dataset_name, task_type=task_type, **params)


def get_titanic(**params) -> pd.DataFrame:
    """
    Load and return the Titanic dataset (classification).

    A classic binary classification dataset containing information about
    passengers aboard the Titanic, including demographic and ticket-related features
    and survival outcome.

    Columns:

    - ``survived`` (int): target variable, 1 if survived, 0 otherwise
    - ``pclass`` (int): passenger class (1 = 1st, 2 = 2nd, 3 = 3rd)
    - ``name`` (str): full name of the passenger
    - ``sex`` (str): passenger gender
    - ``age`` (float): passenger age in years
    - ``fare`` (float): ticket fare
    - ``sibsp`` (int): number of siblings/spouses aboard
    - ``parch`` (int): number of parents/children aboard

    Args:
        **params: Additional parameters passed to the underlying data loader
            (see :ref:`get_data` for details).

    Returns:
        pandas.DataFrame: The Titanic dataset with all features including the target.

        - The DataFrame contains the columns listed above.
        - The target column is ``survived``.
        - All other columns can be used as features for classification tasks.

    Example::

        from dataset_hub.classification import get_titanic

        titanic = get_titanic()
        print(titanic.head())
        X = titanic.drop(columns=['survived'])
        y = titanic['survived']
    """
    return _get_data("titanic", **params)  # type: ignore[return-value]


def get_iris(**params) -> pd.DataFrame:
    """
    Load and return the Iris dataset (classification).

    A classic multiclass classification dataset containing measurements
    of iris flowers from three different species.

    Columns:

    - ``sepal_length`` (float): length of the sepal in cm
    - ``sepal_width`` (float): width of the sepal in cm
    - ``petal_length`` (float): length of the petal in cm
    - ``petal_width`` (float): width of the petal in cm
    - ``species`` (str): target variable, species name (setosa, versicolor, virginica)

    Args:
        **params: Additional parameters passed to the underlying data loader
            (see :ref:`get_data` for details).

    Returns:
        pandas.DataFrame: The Iris dataset with all features including the target.

        - The DataFrame contains the columns listed above.
        - The target column is ``species``.
        - All other columns can be used as features for classification tasks.

    Example::

        from dataset_hub.classification import get_iris

        iris = get_iris()
        print(iris.head())
        X = iris.drop(columns=['species'])
        y = iris['species']
    """

    return _get_data("iris", **params)  # type: ignore[return-value]
