src/features/build_feature.py

"""Module pour la construction des features et la séparation des données en ensembles d'entraînement et de test.
Ce module contient des fonctions pour séparer les features de la target et pour diviser les données en ensembles d'entraînement et de test, avec une option de stratification.
"""
import pandas as pd
from sklearn.model_selection import train_test_split
from src.config import RANDOM_STATE, TEST_SIZE, STRATIFY

def split_features_target(df,target,features):
    """
    Sépare les features de la target
    Args:
        df (pd.DataFrame): Le DataFrame contenant les données.
        target (str): Le nom de la colonne cible (target).
        features (list): La liste des noms des colonnes à utiliser comme features.
    Returns:
        X (pd.DataFrame): Le DataFrame contenant les features.
        y (pd.Series): La Series contenant la target.
    """
    try:
        X = df[features]
        y = df[target]
        return X, y
    except KeyError as e:
        print(f"Erreur : La colonne {e} est introuvable dans le DataFrame.")
        return None, None
    except Exception as e:
        print(f"Une erreur est survenue lors de la séparation des features et de la target : {e}")
        return None, None
    



def split_train_test(X, y, test_size=TEST_SIZE, random_state=RANDOM_STATE, stratify=STRATIFY):
    """
    Sépare les données en ensembles d'entraînement et de test.
    Args:
        X (pd.DataFrame): Les features.
        y (pd.Series): La target.
        test_size (float): La proportion des données à utiliser pour le test (par défaut 0.2).
        random_state (int): La graine pour la reproductibilité (par défaut 42).
        stratify (pd.Series): La variable à utiliser pour la stratification (par défaut None)
        stratify peut être dans X ou y
    Returns:
        X_train (pd.DataFrame): Les features d'entraînement.
        X_test (pd.DataFrame): Les features de test.
        y_train (pd.Series): La target d'entraînement.
        y_test (pd.Series): La target de test.
    """
    try:
        if stratify is not None:
            if stratify in X.columns:
                stratify_data = X[stratify]
            elif stratify in y.name:
                stratify_data = y
            else:
                raise ValueError(f"La variable de stratification '{stratify}' n'est pas trouvée dans les features ou la target.")
        else:
            stratify_data = None

        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=test_size, random_state=random_state, stratify=stratify_data)
        return X_train, X_test, y_train, y_test
    except ValueError as e:
        print(f"Erreur : {e}")
        return None, None, None, None
    except Exception as e:
        print(f"Une erreur est survenue lors de la séparation des données en ensembles d'entraînement et de test : {e}")
        return None, None, None, None