# ===========================================
# Fichier: utilities/cuml_pyfunc_wrapper.py
# ===========================================

from typing import Any, Dict, Optional
import pandas as pd
import numpy as np
import cudf
import cupy as cp
import pickle
import os
import mlflow

from interfaces.vectorizer import Vectorizer


class CuMLPyFuncWrapper(mlflow.pyfunc.PythonModel):
    """
    Classe wrapper pour intégration de modèles cuML dans MLflow PyFunc,
    permettant le chargement et l'inférence.
    """

    def __init__(self, vectorizer: Vectorizer, classifier: object) -> None:
        """
        Initialise le wrapper avec un vectorizer cuDF/cupy et le classifieur cuML.

        :param vectorizer: Instance implémentant l'interface Vectorizer.
        :param classifier: Modèle cuML déjà entraîné (ex. SVC, RandomForest, etc.).
        """
        self.vectorizer: Vectorizer = vectorizer
        self.classifier: object = classifier

    def load_context(self, context: Dict[str, Any]) -> None:
        """
        Chargé par MLflow PyFunc lors du rechargement du modèle,
        par exemple pour initialiser l'environnement.

        :param context: Contexte de chargement contenant d'éventuelles informations
                        sur l'environnement ou l'emplacement d'artefacts.
        """
        

        # Charger le vectorizer depuis les artefacts
        vectorizer_path = os.path.join(context.artifacts["vectorizer"])
        with open(vectorizer_path, "rb") as f:
            self.vectorizer = pickle.load(f)

        # Charger le classifier depuis les artefacts
        classifier_path = os.path.join(context.artifacts["classifier"])
        with open(classifier_path, "rb") as f:
            self.classifier = pickle.load(f)

    def predict(self, context: Dict[str, Any],
                model_input: pd.DataFrame) -> np.ndarray:
        """
        Fonction de prédiction, appelée par MLflow PyFunc.
        Convertit model_input en cudf, vectorise, puis appelle le modèle cuML.

        :param context: Contexte d'exécution éventuel.
        :param model_input: Données d'entrée sous forme de DataFrame pandas.
        :return: Un vecteur numpy des prédictions.
        """
        

        # Convertir le DataFrame pandas en DataFrame cuDF
        cudf_input = cudf.DataFrame(model_input)

        # Vectoriser les données d'entrée
        X_vectorized = self.vectorizer.transform(cudf_input)

        # Effectuer la prédiction avec le classifieur cuML
        predictions_gpu = self.classifier.predict(X_vectorized)

        # Convertir les prédictions GPU en array numpy pour MLflow
        return cp.asnumpy(predictions_gpu)