import joblib
import numpy as np
import mne

def extract_features(raw, channels=None):
    """
    Fonction de base pour extraire des features simples à partir du signal EEG.
    Ici, on calcule la moyenne par canal.
    À améliorer avec tes features plus avancées !
    """
    if channels:
        raw = raw.copy().pick_channels(channels)
    data = raw.get_data()  # shape : (n_channels, n_times)
    features = np.mean(data, axis=1)  # Moyenne par canal
    return features

def predict_epilepsy(raw, model_path, channels=None):
    """
    Prédiction de l'état épileptique à partir du signal brut et d'un modèle sauvegardé.
    Retourne aussi la probabilité associée.
    """
    try:
        model = joblib.load(model_path)
        if channels:
            raw = raw.copy().pick_channels(channels)

        features = extract_features(raw).reshape(1, -1)
        prediction = model.predict(features)[0]

        if hasattr(model, "predict_proba"):
            prob = model.predict_proba(features)[0][1]  # proba que ce soit 1 (epilepsie)
        else:
            prob = None

        label = "Épilepsie détectée" if prediction == 1 else "Pas d’épilepsie détectée"
        return label, prob

    except Exception as e:
        return f"Erreur : {e}", None