import tensorflow as tf
import numpy as np
from PIL import Image
import tensorflow as tf
import logging
import numpy as np
from PIL import Image
from keras.applications.efficientnet_v2 import preprocess_input as effnet_preprocess
from keras.applications.resnet_v2 import preprocess_input as resnet_preprocess
import io
from tf_keras_vis.gradcam import Gradcam,GradcamPlusPlus
from tf_keras_vis.utils import normalize

import numpy as np
import tensorflow as tf
from tf_keras_vis.saliency import Saliency
from tf_keras_vis.utils import normalize
import numpy as np
import tensorflow as tf
from tf_keras_vis.saliency import Saliency
from tf_keras_vis.utils import normalize
import logging
import time

from typing import TypedDict, Callable, Any
logging.basicConfig(
    level=logging.INFO,  # ou logging.DEBUG
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s"
)
logger = logging.getLogger(__name__)
confidence_threshold=0.55
entropy_threshold=2

class ModelStruct(TypedDict):
    model_name: str
    model: tf.keras.Model
    gradcam_model:tf.keras.Model
    preprocess_input: Callable[[np.ndarray], Any]
    target_size: tuple[int, int]
    last_conv_layer:str
    gradcam_type:str


_model_cache: list[ModelStruct] | None = None
def load_model() -> list[ModelStruct]:
    global _model_cache
    if _model_cache is None:
        print("📦 Chargement du modèle EfficientNetV2M...")
        model = tf.keras.models.load_model("model/best_efficientnetv2m_gradcam.keras", compile=False)

        _model_cache = [{
            "model_name": "EfficientNetV2M",
            "model": model,
            "gradcam_model": model, 
            "preprocess_input": effnet_preprocess,
            "target_size": (480, 480),
            "last_conv_layer": "block7a_expand_conv",
            "gradcam_type": "gradcam++"
        }]
    return _model_cache



def compute_gradcam(model, image_array, class_index=None, layer_name=None,gradcam_type="gradcam"):
    """
    Calcule la carte Grad-CAM pour une image et un modèle Keras.

    Args:
        model: tf.keras.Model.
        image_array: np.array (H, W, 3), float32, pré-traitée.
        class_index: int ou None, index de la classe cible. Si None, classe prédite.
        layer_name: str ou None, nom de la couche convolutionnelle à utiliser. Si None, dernière conv.

    Returns:
        gradcam_map: np.array (H, W), normalisée entre 0 et 1.
    """
    logging.info(f"Lancement calcul de la gradcam avec le type {gradcam_type}")

    if image_array.ndim == 3:
        input_tensor = np.expand_dims(image_array, axis=0)
    else:
        input_tensor = image_array
    if gradcam_type=="gradcam++":
        gradcam = GradcamPlusPlus(model, clone=False)
    else: 
        gradcam = Gradcam(model, clone=False)

    def loss(output):
        if class_index is None:
            class_index_local = tf.argmax(output[0])
        else:
            class_index_local = class_index
        return output[:, class_index_local]

    # Choisir la couche à utiliser pour GradCAM
    if layer_name is None:
        # Si non spécifié, chercher la dernière couche conv 2D
        for layer in reversed(model.layers):
            if 'conv' in layer.name and len(layer.output_shape) == 4:
                layer_name = layer.name
                break
        if layer_name is None:
            raise ValueError("Aucune couche convolutionnelle 2D trouvée dans le modèle.")

    cam = gradcam(loss, input_tensor, penultimate_layer=layer_name)
    cam = cam[0]

    # Normaliser entre 0 et 1
    cam = normalize(cam)

    return cam

 
def preprocess_image(image_bytes, target_size, preprocess_input):
    try:
        logger.info("📤 Lecture des bytes et conversion en image PIL")
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    except Exception as e:
        logger.exception("❌ Erreur lors de l'ouverture de l'image")
        raise ValueError("Impossible de décoder l'image") from e

    logger.info(f"📐 Redimensionnement de l'image à la taille {target_size}")
    image = image.resize(target_size)
    image_array = np.array(image).astype(np.float32)

    logger.debug(f"🔍 Shape de l'image après conversion en tableau : {image_array.shape}")

    if image_array.ndim != 3 or image_array.shape[-1] != 3:
        logger.error(f"❌ Image invalide : shape={image_array.shape}")
        raise ValueError("Image must have 3 channels (RGB)")

    logger.info("🎨 Conversion et prétraitement de l'image")

    # Préparation pour la prédiction
    preprocessed_input = preprocess_input(image_array.copy())
    preprocessed_input = np.expand_dims(preprocessed_input, axis=0)

    # Préparation pour Grad-CAM (non prétraitée, mais batchifiée et en float32)
    raw_input = np.expand_dims(image_array / 255.0, axis=0)  # Mise à l’échelle simple

    logger.debug(f"🧪 Shape après ajout de la dimension batch : {preprocessed_input.shape}")
    return preprocessed_input, raw_input



def compute_entropy_safe(probas):
    probas = np.array(probas)
    # On garde uniquement les probabilités strictement positives
    mask = probas > 0
    entropy = -np.sum(probas[mask] * np.log(probas[mask]))
    return entropy


def predict_with_model(config, image_bytes: bytes,show_heatmap=False):
   
    input_array,raw_input = preprocess_image(image_bytes,config["target_size"],config["preprocess_input"])

    logger.info("🤖 Lancement de la prédiction avec le modèle")
    preds = config["model"].predict(input_array)
    logger.debug(f"📈 Prédictions brutes : {preds[0].tolist()}")

    predicted_class_index = int(np.argmax(preds[0]))
    confidence = float(preds[0][predicted_class_index])
    entropy=float(compute_entropy_safe(preds))
    is_uncertain_model= (confidence<confidence_threshold) or (entropy>entropy_threshold)
    logger.info(f"✅ Prédiction : classe={predicted_class_index}, confiance={confidence:.4f},entropy={entropy:.4f},is_uncertain_model={is_uncertain_model}")

    result= {
        "preds": preds[0].tolist(),
        "predicted_class": predicted_class_index,
        "confidence": confidence,
        "entropy":entropy,
        "is_uncertain_model":is_uncertain_model
    }
    if show_heatmap and not is_uncertain_model:
        try:
            logger.info("✅ Début de la génération de la heatmap")
            start_time = time.time()

            # Vérification des entrées
            logger.info(f"🖼️ Image d'entrée shape: {raw_input.shape}")
            logger.info(f"🎯 Index de classe prédite: {predicted_class_index}")
            logger.info(f"🛠️ Dernière couche utilisée: {config['last_conv_layer']}")

            # Calcul de la heatmap
            heatmap = compute_gradcam(config["gradcam_model"], raw_input, class_index=predicted_class_index, layer_name=config["last_conv_layer"],gradcam_type=config["gradcam_type"])

            elapsed_time = time.time() - start_time
            logger.info(f"✅ Heatmap générée en {elapsed_time:.2f} secondes")

            # Conversion en liste pour le JSON
            result["heatmap"] = heatmap.tolist()

        except Exception as e:
            logger.error(f"❌ Erreur lors de la génération de la heatmap: {e}")
            result["heatmap"] = []
    else:
        logger.info("ℹ️ Heatmap non générée (option désactivée ou modèle incertain)")
        result["heatmap"] = []


    return result