import os
import json
import tempfile
import numpy as np
import torch
import matplotlib.pyplot as plt
import segmentation_models_pytorch as smp
import mlflow
import re
import ast
import time

# Visualisation
import matplotlib.patches as mpatches
from matplotlib.colors import ListedColormap

# Preprocessing/Images
import albumentations as A
from albumentations.pytorch import ToTensorV2


def load_cityscapes_config(config_path="../cityscapes_config.json", verbose=True):
    """
    Charge la configuration Cityscapes depuis un fichier JSON.
    INCHANGÉ - Compatible PyTorch
    """
    with open(config_path, "r", encoding="utf-8") as f:
        config = json.load(f)

    class_groups = config["class_groups"]
    cityscapes_mapping = {int(k): v for k, v in config["cityscapes_mapping"].items()}
    id_to_label = {int(k): v for k, v in config["id_to_label"].items()}
    group_colors = np.array(config["group_colors"], dtype=np.uint8)
    group_names = config["group_names"]

    id_to_group = np.zeros(256, dtype=np.uint8)
    for id_, group_id in cityscapes_mapping.items():
        id_to_group[id_] = group_id

    num_classes = len(group_names)

    if verbose:
        print("== Groupes de classes Cityscapes ==")
        for group, classes in class_groups.items():
            print(f"  - {group} : {classes}")

        print("\n== Mapping ID Cityscapes → Groupe ==")
        for id_ in sorted(cityscapes_mapping.keys()):
            label = id_to_label.get(id_, "Label inconnu")
            group_id = cityscapes_mapping[id_]
            print(f"  ID {id_:2d} ('{label}') → groupe '{group_names[group_id]}' ({group_id})")

        print("\n== Liste des groupes et couleurs associées ==")
        for i, (name, color) in enumerate(zip(group_names, group_colors)):
            print(f"  Groupe {i}: {name} - Couleur RGB : {color}")

        print(f"\nNombre total de groupes/classes : {num_classes}")

    return {
        "class_groups": class_groups,
        "cityscapes_mapping": cityscapes_mapping,
        "id_to_label": id_to_label,
        "group_colors": group_colors,
        "group_names": group_names,
        "id_to_group": id_to_group,
        "num_classes": num_classes
    }

def map_mask_ids(mask, id_to_group_mapping):
    """
    Applique le mapping des IDs Cityscapes vers les groupes
    ADAPTÉ PYTORCH : Utilise torch au lieu de tf
    """
    if isinstance(mask, np.ndarray):
        mask = torch.from_numpy(mask)
    
    if isinstance(id_to_group_mapping, np.ndarray):
        if len(id_to_group_mapping) != 256:
            raise ValueError(f"id_to_group_mapping doit avoir 256 éléments, reçu: {len(id_to_group_mapping)}")
        mapping_tensor = torch.tensor(id_to_group_mapping, dtype=torch.long)
    else:
        mapping_tensor = torch.tensor(list(id_to_group_mapping), dtype=torch.long)
    
    # S'assurer que les valeurs du masque sont dans la plage valide [0, 255]
    mask_clipped = torch.clamp(mask, 0, 255)
    mapped_mask = mapping_tensor[mask_clipped.long()]
    return mapped_mask
    
def colorize_mask(mask, group_colors):
    """
    Applique la colorisation du masque avec les couleurs définies
    INCHANGÉ - Compatible PyTorch (prend numpy en entrée)
    """
    if isinstance(mask, torch.Tensor):
        mask = mask.numpy()
        
    colored_mask = np.zeros((mask.shape[0], mask.shape[1], 3), dtype=np.uint8)
    num_classes = len(group_colors)
    
    for group_id in range(num_classes):
        colored_mask[mask == group_id] = group_colors[group_id]
    
    return colored_mask


def get_preprocessing_fn(img_size=(512, 512), encoder=None):
    """
    Retourne une fonction qui prétraite à la fois les images et les masques
    pour l'inférence, en reproduisant exactement le pipeline CityscapesDataset.
    """
    
    if encoder is None:
        mean, std = [0.0, 0.0, 0.0], [1.0, 1.0, 1.0]
    else:
        mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]

    # Pipeline unifié (gérera image ET masque ensemble)
    transform = A.Compose([
        A.Resize(height=img_size[0], width=img_size[1]),
        A.Normalize(mean=mean, std=std),
        ToTensorV2()
    ])

    def preprocess(image_np, mask_np=None):
        if mask_np is not None:
            # Traitement simultané image + masque
            transformed = transform(image=image_np, mask=mask_np)
            return transformed['image'], transformed['mask'].long()
        else:
            # Image seule
            transformed = transform(image=image_np)
            return transformed['image'], None

    return preprocess
    
# Version simplifiée de preprocess_image_and_mask utilisant la nouvelle get_preprocessing_fn
def preprocess_image_and_mask(image_path, mask_path=None, img_size=(512, 512),
                                    encoder_name=None, id_to_group_mapping=None):
    """
    Version simplifiée utilisant get_preprocessing_fn pour la cohérence avec CityscapesDataset
    """
    import numpy as np
    from PIL import Image
    import torch
    
    print(f"🔄 Préprocessing de l'image: {os.path.basename(image_path)}")
    
    # Chargement des données
    image_pil = Image.open(image_path).convert('RGB')
    image_np = np.array(image_pil)
    img_display = np.array(image_pil.resize(img_size, Image.BILINEAR))  # Pour affichage
    
    mask_np = None
    if mask_path and os.path.exists(mask_path):
        print(f"🔄 Préprocessing du masque: {os.path.basename(mask_path)}")
        mask_pil = Image.open(mask_path).convert('L')
        mask_np = np.array(mask_pil)
        print(f"   Valeurs uniques avant mapping: {np.unique(mask_np)[:10]}...")
    
    # Utiliser notre preprocessing unifié
    preprocess_fn = get_preprocessing_fn(img_size=img_size, encoder=encoder_name)
    img_preprocessed, mask_preprocessed = preprocess_fn(image_np, mask_np)
    
    # Appliquer le mapping si fourni ET si on a un masque
    if mask_preprocessed is not None and id_to_group_mapping is not None:
        mask_preprocessed = map_mask_ids(mask_preprocessed, id_to_group_mapping)
        print(f"   Valeurs uniques après mapping: {torch.unique(mask_preprocessed)}")
    
    return img_preprocessed, img_display, mask_preprocessed

# Version mise à jour de run_inference_and_visualize
def run_inference_and_visualize(image_path, model, encoder_name, img_size, mapping_config, 
                               mask_path=None, save_dir=None, show_stats=True, show_plot=True):
    """
    Version simplifiée utilisant preprocess_image_and_mask
    """
    print(f"=== INFÉRENCE PYTORCH SUR {os.path.basename(image_path)} ===")

    # Utiliser la version simplifiée du preprocessing
    img_preprocessed, img_display, mask_gt = preprocess_image_and_mask(
        image_path=image_path,
        mask_path=mask_path,
        img_size=img_size or (512, 512),
        encoder_name=encoder_name,
        id_to_group_mapping=mapping_config.get('id_to_group')
    )

    print("🔮 Exécution de l'inférence PyTorch...")
    model.eval()
    device = next(model.parameters()).device

    with torch.no_grad():
        # Ajouter dimension batch
        img_batch = img_preprocessed.unsqueeze(0).to(device)  # [1, 3, H, W]
        
        # ⏱️ MESURE DU TEMPS D'INFÉRENCE
        start_time = time.time()
        predictions = model(img_batch)  # [1, num_classes, H, W]
        end_time = time.time()
        
        inference_time = end_time - start_time
        
        predicted_mask = torch.argmax(predictions[0], dim=0).cpu().numpy()  # [H, W]

    print(f"✅ Inférence terminée - Shape: {predicted_mask.shape}")
    print(f"⏱️ Temps d'inférence: {inference_time*1000:.1f} ms ({inference_time:.4f} s)")

    # [Le reste du code de visualisation reste identique]
    # Visualisation
    num_plots = 2 if mask_gt is None else 4
    fig, axes = plt.subplots(1, num_plots, figsize=(5*num_plots, 5))
    if num_plots == 1:
        axes = [axes]

    axes[0].imshow(img_display)
    axes[0].set_title("Image Originale")
    axes[0].axis('off')

    predicted_colored = colorize_mask(predicted_mask, mapping_config['group_colors'])
    axes[1].imshow(predicted_colored)
    axes[1].set_title("Masque Prédit")
    axes[1].axis('off')

    if mask_gt is not None:
        mask_gt_np = mask_gt.numpy() if isinstance(mask_gt, torch.Tensor) else mask_gt
        gt_colored = colorize_mask(mask_gt_np, mapping_config['group_colors'])
        axes[2].imshow(gt_colored)
        axes[2].set_title("Masque de Vérité")
        axes[2].axis('off')

        # Visualisation des erreurs
        diff_mask = (mask_gt_np != predicted_mask).astype(float)
        axes[3].imshow(img_display)
        axes[3].imshow(diff_mask, alpha=0.5, cmap='Reds')
        axes[3].set_title("Erreurs (en rouge)")
        axes[3].axis('off')

    colors_normalized = mapping_config['group_colors'] / 255.0
    custom_cmap = ListedColormap(colors_normalized)

    legend_elements = []
    for i, (name, color) in enumerate(zip(mapping_config['group_names'], colors_normalized)):
        legend_elements.append(mpatches.Patch(color=color, label=f'{i}: {name}'))

    fig.legend(handles=legend_elements, 
              loc='center left', 
              bbox_to_anchor=(1.01, 0.5),
              fontsize=9,
              frameon=True,
              fancybox=True,
              shadow=True)

    plt.tight_layout()
    fig.subplots_adjust(right=0.95)

    if save_dir:
        os.makedirs(save_dir, exist_ok=True)
        img_name = os.path.splitext(os.path.basename(image_path))[0]
        fig_path = os.path.join(save_dir, f'{img_name}_inference.png')
        fig.savefig(fig_path, dpi=150, bbox_inches='tight')
        print(f"💾 Résultat sauvegardé: {fig_path}")


    # Statistiques
    if show_stats:
        print("\n📊 Statistiques du masque prédit:")
        unique_classes, counts = np.unique(predicted_mask, return_counts=True)
        for class_id, count in zip(unique_classes, counts):
            if class_id < len(mapping_config['group_names']):
                percentage = (count / predicted_mask.size) * 100
                print(f"   {mapping_config['group_names'][class_id]}: {count} pixels ({percentage:.1f}%)")

        if mask_gt is not None:
            mask_gt_np = mask_gt.numpy() if isinstance(mask_gt, torch.Tensor) else mask_gt
            print("\n📊 Statistiques du masque de vérité:")
            unique_classes, counts = np.unique(mask_gt_np, return_counts=True)
            for class_id, count in zip(unique_classes, counts):
                if class_id < len(mapping_config['group_names']):
                    percentage = (count / mask_gt_np.size) * 100
                    print(f"   {mapping_config['group_names'][class_id]}: {count} pixels ({percentage:.1f}%)")

    if show_plot:
        plt.show()
        return predicted_mask, inference_time
    else:
        return predicted_mask, inference_time, fig


def load_model(run_id=None, experiment_name="OC Projet 9", 
               metric="test_mean_iou", top_n=1):
    """
    Charge un modèle PyTorch depuis MLflow avec plusieurs méthodes de sélection
    
    Args:
        run_id (str, optional): Run ID spécifique (prioritaire si fourni)
        experiment_name (str): Nom de l'expérience MLflow
        metric (str): Métrique pour classer les runs
        top_n (int): Position dans le classement (1=meilleur, 2=deuxième, etc.)
        
    Returns:
        tuple: (model, source_run_id, encoder_name, script_img_size)
    """
    
    print(f"\n🔍 Chargement d'un modèle PyTorch depuis MLflow...")
    
    client = mlflow.tracking.MlflowClient()
    
    # 1. Sélection du run selon les paramètres
    if run_id:
        print(f"🎯 Mode run_id spécifique: {run_id}")
        try:
            run = mlflow.get_run(run_id)
            source_run_id = run_id
            print(f"✅ Run trouvé: {run.data.tags.get('mlflow.runName', 'Unknown')}")
        except Exception as e:
            raise RuntimeError(f"❌ Run {run_id} introuvable: {e}")
    
    else:
        print(f"🏆 Mode classement: top_{top_n} de '{metric}' dans '{experiment_name}'")
        try:
            # Récupérer l'expérience
            experiment = client.get_experiment_by_name(experiment_name)
            if experiment is None:
                raise RuntimeError(f"❌ Expérience '{experiment_name}' non trouvée")
            
            # Rechercher les runs avec la métrique, triés par ordre décroissant
            filter_string = f"status = 'FINISHED' and metrics.{metric} >= 0"
            order_by = f"metrics.{metric} DESC"
            
            runs_df = mlflow.search_runs(
                experiment_ids=[experiment.experiment_id],
                filter_string=filter_string,
                order_by=[order_by],
                max_results=top_n
            )
            
            if runs_df.empty:
                raise RuntimeError(f"❌ Aucun run avec '{metric}' trouvé dans '{experiment_name}'")
            
            if len(runs_df) < top_n:
                raise RuntimeError(f"❌ Seulement {len(runs_df)} runs trouvés, impossible d'obtenir le top_{top_n}")
            
            # Prendre le run à la position top_n
            selected_run = runs_df.iloc[top_n - 1]
            source_run_id = selected_run['run_id']
            metric_value = selected_run[f'metrics.{metric}']
            run_name = selected_run.get('tags.mlflow.runName', 'Unknown')
            
            print(f"✅ Run sélectionné (position #{top_n}):")
            print(f"   • Nom: {run_name}")
            print(f"   • Run ID: {source_run_id}")
            print(f"   • {metric}: {metric_value:.4f}")
            
            # Afficher le contexte
            print(f"📊 Contexte du classement:")
            for i, (_, row) in enumerate(runs_df.head(min(5, len(runs_df))).iterrows()):
                indicator = "👑" if i == (top_n - 1) else f"{i+1}."
                name = row.get('tags.mlflow.runName', 'Unknown')[:30]
                value = row[f'metrics.{metric}']
                print(f"   {indicator} {name}: {value:.4f}")
            
            run = mlflow.get_run(source_run_id)
            
        except Exception as e:
            raise RuntimeError(f"❌ Erreur lors de la sélection: {e}")
    
    # 2-5. [Reste du code identique - extraction params, détection variante, chargement...]
    try:
        params = run.data.params
        
        encoder_name = params.get("encoder_name", None)
        architecture = params.get("architecture", "FPN")
        num_classes = int(params.get("num_classes", 8))
        script_img_size_raw = params.get("script_img_size", None)
        run_name = run.data.tags.get('mlflow.runName', 'Unknown')
        
        # Parser script_img_size
        script_img_size = None
        if script_img_size_raw:
            try:
                script_img_size = ast.literal_eval(script_img_size_raw)
                if isinstance(script_img_size, (tuple, list)) and len(script_img_size) == 2:
                    script_img_size = tuple(script_img_size)
            except:
                script_img_size = None
        
        print(f"\n📋 Paramètres du modèle:")
        print(f"🏗️ Architecture: {architecture}")
        print(f"🔧 Encodeur: {encoder_name}")
        print(f"📊 Nombre de classes: {num_classes}")
        print(f"🖼️ Script img size: {script_img_size}")
        
    except Exception as e:
        raise RuntimeError(f"❌ Erreur lors de la récupération des paramètres: {e}")
    
    # Détecter la variante SegFormer
    def detect_segformer_variant_from_name(run_name, params):
        segformer_pattern = r'SegFormer[_-]?b(\d+)'
        match = re.search(segformer_pattern, run_name, re.IGNORECASE)
        if match:
            variant_num = match.group(1)
            detected = f"mit_b{variant_num}"
            print(f"🔍 Variante détectée: 'SegFormer_b{variant_num}' → {detected}")
            return detected
        
        if encoder_name and encoder_name.startswith("mit_b"):
            print(f"🔍 Variante depuis paramètres: {encoder_name}")
            return encoder_name
            
        b_pattern = r'\bb(\d+)\b'
        match = re.search(b_pattern, run_name, re.IGNORECASE)
        if match:
            variant_num = match.group(1)
            detected = f"mit_b{variant_num}"
            print(f"🔍 Pattern détecté: 'b{variant_num}' → {detected}")
            return detected
        
        print(f"🔍 Aucune variante détectée, utilisation de mit_b0")
        return "mit_b0"
    
    detected_variant = detect_segformer_variant_from_name(run_name, params)
    
    # Télécharger et charger le modèle
    with tempfile.TemporaryDirectory() as tmpdir:
        try:
            artifacts = client.list_artifacts(source_run_id)
            model_artifact_candidates = [a.path for a in artifacts if a.path.endswith('.pth')]
            
            if not model_artifact_candidates:
                raise RuntimeError("❌ Aucun fichier .pth trouvé")
            
            model_artifact = model_artifact_candidates[0]
            print(f"⬇️ Téléchargement de {model_artifact}...")
            
            local_path = mlflow.artifacts.download_artifacts(
                run_id=source_run_id,
                artifact_path=model_artifact,
                dst_path=tmpdir
            )
            
            checkpoint = torch.load(local_path, map_location='cpu')
            if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
                state_dict = checkpoint['state_dict']
            else:
                state_dict = checkpoint
                
        except Exception as e:
            raise RuntimeError(f"❌ Erreur lors du téléchargement: {e}")
    
    # Créer le modèle
    try:
        is_transformer = any('patch_embed' in key for key in state_dict.keys())
        
        if is_transformer:
            print(f"🔧 Création SegFormer ({detected_variant})")
            model = smp.Segformer(encoder_name=detected_variant, classes=num_classes)
            encoder_name = detected_variant
        else:
            print("🔧 Création modèle CNN")
            if encoder_name is None or encoder_name == "None":
                encoder_name = "efficientnet-b0"
            
            if architecture.lower() == "fpn":
                model = smp.FPN(encoder_name=encoder_name, classes=num_classes)
            elif architecture.lower() == "unet":
                model = smp.Unet(encoder_name=encoder_name, classes=num_classes)
            elif architecture.lower() == "deeplabv3plus":
                model = smp.DeepLabV3Plus(encoder_name=encoder_name, classes=num_classes)
            else:
                raise ValueError(f"Architecture {architecture} non supportée")
        
        print("🔄 Chargement des poids...")
        model.load_state_dict(state_dict, strict=True)
        model.eval()
        print("✅ Modèle chargé avec succès!")
        
    except Exception as e:
        try:
            model.load_state_dict(state_dict, strict=False)
            model.eval()
            print("✅ Chargement partiel réussi!")
        except Exception as e2:
            raise RuntimeError(f"❌ Erreur lors du chargement: {e2}")
    
    print(f"🎉 Modèle chargé depuis le run {source_run_id}")
    return model, source_run_id, encoder_name, script_img_size


def load_best_model_from_registry(model_name):
    """
    Charge le meilleur modèle depuis le MLflow Model Registry
    (Version refactorisée - wrapper autour de load_model)
    
    Args:
        model_name (str): Nom du modèle dans le Registry
        
    Returns:
        tuple: (model, source_run_id, encoder_name, script_img_size)
    """
    print(f"🏛️ Récupération du modèle '{model_name}' depuis le Registry...")
    
    client = mlflow.tracking.MlflowClient()
    
    try:
        # Récupérer la version en Production
        versions = client.search_model_versions(f"name='{model_name}'")
        prod_versions = [v for v in versions if v.current_stage == "Production"]
        
        if not prod_versions:
            raise RuntimeError(f"❌ Aucune version en Production trouvée pour {model_name}")
        
        latest_prod = max(prod_versions, key=lambda v: int(v.version))
        print(f"📦 Version Registry: {latest_prod.version} (Stage: {latest_prod.current_stage})")
        
        # Récupérer le source_run_id depuis le tag du Registry
        source_run_id = latest_prod.tags.get("source_run_id")
        if not source_run_id:
            raise RuntimeError(f"❌ Tag 'source_run_id' absent pour la version {latest_prod.version}")
        
        print(f"🔗 Source run détecté: {source_run_id}")
        
    except Exception as e:
        raise RuntimeError(f"❌ Erreur lors de l'accès au Registry: {e}")
    
    # Déléguer le chargement à load_model avec le run_id trouvé
    print("➡️ Délégation vers load_model...")
    return load_model(run_id=source_run_id)


# Fonctions de convenance
def load_best_model(experiment_name="OC Projet 9", metric="test_mean_iou"):
    """Charge le meilleur modèle selon une métrique"""
    return load_model(experiment_name=experiment_name, metric=metric, top_n=1)

def load_second_best_model(experiment_name="OC Projet 9", metric="test_mean_iou"):
    """Charge le deuxième meilleur modèle selon une métrique"""
    return load_model(experiment_name=experiment_name, metric=metric, top_n=2)

def load_model_by_run_id(run_id):
    """Charge un modèle par run_id spécifique"""
    return load_model(run_id=run_id)