from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split

class EnsembleModel:
    def __init__(self):

        self.group_classifier = RandomForestClassifier()
        self.product_classifier = RandomForestClassifier()

    def fit(self, X_train, y_group_train, y_product_train):

        self.group_classifier.fit(X_train, y_group_train)
        pred_group = self.group_classifier.predict(X_train)
        print(pred_group)
        liste = list(set(pred_group))
        for i in liste :
           X_filtered = X_train[y_group_train == i]
           y_product_filtered = y_product_train[y_group_train == i]
           self.product_classifier.fit(X_filtered, y_product_filtered)

    def predict(self, X):
        y_group_pred = self.group_classifier.predict(X)

        X_filtered = X[self.group_classifier.predict(X) == y_group_pred] 
        y_product_pred = self.product_classifier.predict(X_filtered)

        return y_group_pred, y_product_pred


# Exemple d'utilisation de la classe EnsembleModel
if __name__ == "__main__":
    # Exemple de données d'entraînement (vous devriez utiliser vos propres données réelles)
    X = [[1, 2], [2, 3], [3, 4], [4, 5], [5, 6]]
    y_group = [0, 1, 0, 1, 0]
    y_product = [10, 11, 12, 13, 14]

    # Création et entraînement du modèle ensembliste
    ensemble_model = EnsembleModel()
    ensemble_model.fit(X, y_group, y_product)

    # Exemple de données de test (vous devriez utiliser vos propres données réelles)
    X_test = [[3, 5], [1, 3]]
    y_group_pred, y_product_pred = ensemble_model.predict(X_test)

    print("Prédictions du groupe de classe :", y_group_pred)
    print("Prédictions de la classe du produit :", y_product_pred)