"""Module de prétraitement des données pour les modèles de machine learning.
Ce module contient des fonctions pour identifier les types de features (numériques et catégorielles)
et pour construire des préprocesseurs adaptés aux différents types de modèles (linéaires, basés sur les arbres, boosting)."""
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import OneHotEncoder, StandardScaler


def get_feature_types(X: pd.DataFrame):
    """
    Identifie les colonnes numériques et catégorielles.
    args:
        X : pd.DataFrame
    returns:
        numeric_features : list
        categorical_features : list
    """
    numeric_features = X.select_dtypes(include=["int64", "float64"]).columns.tolist()
    categorical_features = X.select_dtypes(include=["object", "category"]).columns.tolist()

    return numeric_features, categorical_features


def build_linear_preprocessor(X: pd.DataFrame):
    """
    Preprocessing pour modèles linéaires.
    Numériques :
        StandardScaler
    Catégorielles :
        OneHotEncoder
    args:
        X : pd.DataFrame
    returns:
        preprocessor : ColumnTransformer
    """
    numeric_features, categorical_features = get_feature_types(X)

    preprocessor = ColumnTransformer(
        transformers=[
            ("num", StandardScaler(), numeric_features),
            ("cat", OneHotEncoder(handle_unknown="ignore", sparse_output=False), categorical_features),
        ],
        remainder="drop",
    )

    return preprocessor


def build_tree_preprocessor(X: pd.DataFrame):
    """
    Preprocessing pour modèles basés sur les arbres.
    Numériques :
        Pas de transformation
    Catégorielles :
        OneHotEncoder
    args:
        X : pd.DataFrame
    returns:
        preprocessor : ColumnTransformer
    """
    numeric_features, categorical_features = get_feature_types(X)

    preprocessor = ColumnTransformer(
        transformers=[
            ("num", "passthrough", numeric_features),
            ("cat", OneHotEncoder(handle_unknown="ignore", sparse_output=False), categorical_features),
        ],
        remainder="drop",
    )

    return preprocessor


def build_preprocessor(X: pd.DataFrame, model_type: str):
    """
    Construit un préprocesseur adapté au type de modèle.
    args:
        X : pd.DataFrame
        model_type : str, parmi 'linear', 'tree', 'boosting'
    returns:
        preprocessor : ColumnTransformer
    """
    if model_type == "linear":
        return build_linear_preprocessor(X)

    elif model_type in ["tree", "boosting"]:
        return build_tree_preprocessor(X)

    else:
        raise ValueError(
            "model_type doit être parmi : 'linear', 'tree', 'boosting'"
        )

def get_categorical_features(X):
    return X.select_dtypes(include=["object", "category"]).columns.tolist()