# Librerías y configuración

import os
import time
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from PIL import Image

import torch
import torch.nn.functional as F
import torchvision
from torchvision import transforms
from torchvision.models.segmentation import (
    fcn_resnet50, FCN_ResNet50_Weights,
    deeplabv3_resnet50, DeepLabV3_ResNet50_Weights,
    deeplabv3_mobilenet_v3_large, DeepLabV3_MobileNet_V3_Large_Weights
)

# Selección de dispositivo
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Dispositivo: {device}')
print(f'PyTorch version: {torch.__version__}')
print(f'Torchvision version: {torchvision.__version__}')

plt.rcParams['figure.figsize'] = (12, 7)

# Descarga de Pascal VOC 2012
VOC_ROOT = '/tmp/voc_data'

# Descargamos el dataset de segmentación
try:
    voc_dataset = torchvision.datasets.VOCSegmentation(
        root=VOC_ROOT, year='2012', image_set='val', download=True
    )
    print(f'Pascal VOC 2012 (val): {len(voc_dataset)} imágenes')
except Exception as e:
    print(f'Error descargando VOC: {e}')
    print('Intentando con mirror alternativo...')
    # Si falla, usamos imágenes de ejemplo
    voc_dataset = None

# Clases de Pascal VOC
VOC_CLASSES = [
    'fondo', 'avión', 'bicicleta', 'pájaro', 'barco',
    'botella', 'autobús', 'coche', 'gato', 'silla',
    'vaca', 'mesa', 'perro', 'caballo', 'moto',
    'persona', 'planta', 'oveja', 'sofá', 'tren', 'monitor'
]

# Paleta de colores para visualización (colores de VOC)
VOC_COLORMAP = np.array([
    [0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0], [0, 0, 128],
    [128, 0, 128], [0, 128, 128], [128, 128, 128], [64, 0, 0], [192, 0, 0],
    [64, 128, 0], [192, 128, 0], [64, 0, 128], [192, 0, 128], [64, 128, 128],
    [192, 128, 128], [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
    [0, 64, 128]
], dtype=np.uint8)

print(f'\nClases ({len(VOC_CLASSES)}): {VOC_CLASSES}')

# Visualizamos ejemplos del dataset VOC
def decode_segmap(mask_tensor, n_classes=21):
    """Convierte máscara de índices a imagen RGB usando la paleta VOC."""
    mask = mask_tensor.numpy() if torch.is_tensor(mask_tensor) else np.array(mask_tensor)
    rgb = np.zeros((*mask.shape, 3), dtype=np.uint8)
    for c in range(n_classes):
        rgb[mask == c] = VOC_COLORMAP[c]
    return rgb


if voc_dataset is not None:
    fig, axes = plt.subplots(3, 3, figsize=(16, 14))
    
    # Seleccionamos imágenes interesantes (con múltiples objetos)
    sample_indices = [0, 15, 45, 80, 120, 200, 250, 300, 400]
    
    for i, idx in enumerate(sample_indices):
        row, col = divmod(i, 3)
        img, mask = voc_dataset[idx]
        mask_np = np.array(mask)
        
        # Superposición imagen + segmentación
        img_np = np.array(img)
        seg_rgb = decode_segmap(mask_np)
        
        # Blend
        valid = mask_np < 255  # Ignorar bordes (255)
        overlay = img_np.copy()
        overlay[valid] = (0.5 * img_np[valid] + 0.5 * seg_rgb[valid]).astype(np.uint8)
        
        axes[row, col].imshow(overlay)
        # Identificamos clases presentes
        classes_present = np.unique(mask_np)
        classes_present = classes_present[(classes_present != 0) & (classes_present != 255)]
        class_names = [VOC_CLASSES[c] for c in classes_present if c < len(VOC_CLASSES)]
        axes[row, col].set_title(', '.join(class_names), fontsize=10)
        axes[row, col].axis('off')
    
    plt.suptitle('Pascal VOC 2012: ejemplos con segmentación ground truth', fontsize=14)
    plt.tight_layout()
    plt.show()

# Cargamos los tres modelos
print('Cargando modelos preentrenados...')

models = {}

# 1) FCN-ResNet50
models['FCN-ResNet50'] = fcn_resnet50(weights=FCN_ResNet50_Weights.COCO_WITH_VOC_LABELS_V1)
print('  ✓ FCN-ResNet50')

# 2) DeepLabV3-ResNet50
models['DeepLabV3-ResNet50'] = deeplabv3_resnet50(weights=DeepLabV3_ResNet50_Weights.COCO_WITH_VOC_LABELS_V1)
print('  ✓ DeepLabV3-ResNet50')

# 3) DeepLabV3-MobileNetV3
models['DeepLabV3-MobileNet'] = deeplabv3_mobilenet_v3_large(
    weights=DeepLabV3_MobileNet_V3_Large_Weights.COCO_WITH_VOC_LABELS_V1
)
print('  ✓ DeepLabV3-MobileNetV3-Large')

# Ponemos en modo evaluación y movemos al dispositivo
for name, model in models.items():
    model.eval()
    model.to(device)

# Contamos parámetros
print(f'\nComparación de tamaño de modelos:')
for name, model in models.items():
    n_params = sum(p.numel() for p in model.parameters()) / 1e6
    print(f'  {name}: {n_params:.1f}M parámetros')

# Preprocesamiento estándar para modelos de torchvision
preprocess = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])


def predict_segmentation(model, image_pil):
    """
    Ejecuta segmentación semántica sobre una imagen PIL.
    Retorna la máscara de predicción (H, W) con índices de clase.
    """
    input_tensor = preprocess(image_pil).unsqueeze(0).to(device)
    
    with torch.no_grad():
        output = model(input_tensor)['out']  # [1, 21, H, W]
    
    # Argmax para obtener la clase más probable por píxel
    pred = output.argmax(1).squeeze(0).cpu().numpy()  # [H, W]
    return pred


def predict_with_timing(model, image_pil, n_runs=5):
    """
    Ejecuta segmentación y mide el tiempo promedio de inferencia.
    """
    input_tensor = preprocess(image_pil).unsqueeze(0).to(device)
    
    # Warmup
    with torch.no_grad():
        _ = model(input_tensor)
    
    if device.type == 'cuda':
        torch.cuda.synchronize()
    
    times = []
    for _ in range(n_runs):
        start = time.time()
        with torch.no_grad():
            output = model(input_tensor)['out']
        if device.type == 'cuda':
            torch.cuda.synchronize()
        times.append(time.time() - start)
    
    pred = output.argmax(1).squeeze(0).cpu().numpy()
    avg_time = np.mean(times)
    return pred, avg_time

print('Pipeline de inferencia preparado.')

# Comparación visual en imágenes seleccionadas
if voc_dataset is not None:
    # Seleccionamos imágenes con objetos variados
    test_indices = [10, 50, 100, 150, 250]
    model_names = list(models.keys())
    
    fig, axes = plt.subplots(len(test_indices), len(model_names) + 2, 
                             figsize=(4 * (len(model_names) + 2), 4 * len(test_indices)))
    
    for row, idx in enumerate(test_indices):
        img_pil, mask_pil = voc_dataset[idx]
        img_np = np.array(img_pil)
        gt_mask = np.array(mask_pil)
        
        # Imagen original
        axes[row, 0].imshow(img_np)
        axes[row, 0].axis('off')
        if row == 0:
            axes[row, 0].set_title('Original', fontsize=12)
        
        # Ground Truth
        gt_rgb = decode_segmap(gt_mask)
        axes[row, 1].imshow(gt_rgb)
        axes[row, 1].axis('off')
        if row == 0:
            axes[row, 1].set_title('Ground Truth', fontsize=12)
        
        # Predicciones de cada modelo
        for col, name in enumerate(model_names):
            pred = predict_segmentation(models[name], img_pil)
            pred_rgb = decode_segmap(pred)
            axes[row, col + 2].imshow(pred_rgb)
            axes[row, col + 2].axis('off')
            if row == 0:
                axes[row, col + 2].set_title(name, fontsize=12)
    
    plt.suptitle('Comparación visual de modelos de segmentación en Pascal VOC', fontsize=15)
    plt.tight_layout()
    plt.show()

def compute_iou_per_class(pred, target, n_classes=21, ignore_index=255):
    """
    Calcula IoU por clase entre predicción y ground truth.
    Ignora píxeles con valor ignore_index (bordes en VOC).
    """
    ious = []
    valid = target != ignore_index
    
    for c in range(n_classes):
        pred_c = (pred == c) & valid
        target_c = (target == c) & valid
        
        intersection = (pred_c & target_c).sum()
        union = (pred_c | target_c).sum()
        
        if union == 0:
            ious.append(float('nan'))  # Clase no presente
        else:
            ious.append(intersection / union)
    
    return ious


# Evaluamos cada modelo en un subconjunto de VOC val
N_EVAL = min(200, len(voc_dataset)) if voc_dataset is not None else 0

results = {}
timing_results = {}

if voc_dataset is not None and N_EVAL > 0:
    print(f'Evaluando {N_EVAL} imágenes de Pascal VOC 2012 val...')
    
    for name, model in models.items():
        print(f'\n  Evaluando {name}...')
        all_ious = []
        total_time = 0
        
        for i in range(N_EVAL):
            img_pil, mask_pil = voc_dataset[i]
            gt = np.array(mask_pil)
            
            pred, inf_time = predict_with_timing(model, img_pil, n_runs=1)
            total_time += inf_time
            
            # Redimensionamos predicción al tamaño del GT si es necesario
            if pred.shape != gt.shape:
                pred_tensor = torch.tensor(pred).unsqueeze(0).unsqueeze(0).float()
                pred_resized = F.interpolate(
                    pred_tensor, size=gt.shape, mode='nearest'
                ).squeeze().numpy().astype(int)
            else:
                pred_resized = pred
            
            ious = compute_iou_per_class(pred_resized, gt)
            all_ious.append(ious)
            
            if (i + 1) % 50 == 0:
                print(f'    {i+1}/{N_EVAL} imágenes procesadas')
        
        # Promediamos IoU por clase (ignorando NaN)
        all_ious = np.array(all_ious)
        class_ious = np.nanmean(all_ious, axis=0)
        miou = np.nanmean(class_ious)
        
        results[name] = {
            'mIoU': miou,
            'class_ious': class_ious,
            'avg_time': total_time / N_EVAL
        }
        timing_results[name] = total_time / N_EVAL
        
        print(f'    mIoU = {miou:.4f}, tiempo medio = {total_time/N_EVAL*1000:.1f} ms')

    print('\n¡Evaluación completada!')

if results:
    model_names = list(results.keys())
    
    fig, axes = plt.subplots(1, 3, figsize=(20, 6))
    colors_models = ['#3498db', '#e74c3c', '#2ecc71']
    
    # --- 1) mIoU global ---
    mious = [results[n]['mIoU'] for n in model_names]
    bars = axes[0].bar(model_names, mious, color=colors_models, edgecolor='white', width=0.6)
    axes[0].set_title('mIoU global', fontsize=13)
    axes[0].set_ylabel('mIoU')
    axes[0].set_ylim(0, 1)
    for bar, val in zip(bars, mious):
        axes[0].text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.02,
                     f'{val:.3f}', ha='center', fontsize=12, fontweight='bold')
    axes[0].tick_params(axis='x', rotation=15)
    
    # --- 2) Parámetros ---
    params = [sum(p.numel() for p in models[n].parameters()) / 1e6 for n in model_names]
    bars2 = axes[1].bar(model_names, params, color=colors_models, edgecolor='white', width=0.6)
    axes[1].set_title('Tamaño del modelo', fontsize=13)
    axes[1].set_ylabel('Millones de parámetros')
    for bar, val in zip(bars2, params):
        axes[1].text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.5,
                     f'{val:.1f}M', ha='center', fontsize=12, fontweight='bold')
    axes[1].tick_params(axis='x', rotation=15)
    
    # --- 3) Tiempo de inferencia ---
    times = [results[n]['avg_time'] * 1000 for n in model_names]  # ms
    bars3 = axes[2].bar(model_names, times, color=colors_models, edgecolor='white', width=0.6)
    axes[2].set_title('Tiempo de inferencia', fontsize=13)
    axes[2].set_ylabel('Milisegundos')
    for bar, val in zip(bars3, times):
        axes[2].text(bar.get_x() + bar.get_width()/2, bar.get_height() + 1,
                     f'{val:.0f} ms', ha='center', fontsize=12, fontweight='bold')
    axes[2].tick_params(axis='x', rotation=15)
    
    plt.suptitle('Comparación cuantitativa de modelos de segmentación', fontsize=15)
    plt.tight_layout()
    plt.show()

# IoU por clase para cada modelo
if results:
    fig, ax = plt.subplots(figsize=(16, 8))
    
    x = np.arange(len(VOC_CLASSES))
    width = 0.25
    
    for i, name in enumerate(model_names):
        class_ious = results[name]['class_ious']
        # Reemplazamos NaN por 0 para visualización
        class_ious_plot = np.nan_to_num(class_ious, nan=0.0)
        bars = ax.bar(x + i * width, class_ious_plot, width, 
                      label=name, color=colors_models[i], alpha=0.8, edgecolor='white')
    
    ax.set_xticks(x + width)
    ax.set_xticklabels(VOC_CLASSES, rotation=45, ha='right', fontsize=9)
    ax.set_ylabel('IoU')
    ax.set_title('IoU por clase - Comparación de modelos', fontsize=14)
    ax.legend(fontsize=11)
    ax.set_ylim(0, 1)
    ax.grid(True, axis='y', alpha=0.3)
    
    plt.tight_layout()
    plt.show()
    
    # Tabla resumen por clase
    print(f'{"Clase":<12}', end='')
    for name in model_names:
        print(f'{name:<22}', end='')
    print()
    print('-' * 78)
    for c, cls_name in enumerate(VOC_CLASSES):
        print(f'{cls_name:<12}', end='')
        for name in model_names:
            iou = results[name]['class_ious'][c]
            if np.isnan(iou):
                print(f'{"  -":<22}', end='')
            else:
                print(f'{iou:<22.4f}', end='')
        print()

if results:
    fig, ax = plt.subplots(figsize=(10, 7))
    
    for i, name in enumerate(model_names):
        miou = results[name]['mIoU']
        n_params = sum(p.numel() for p in models[name].parameters()) / 1e6
        inf_time = results[name]['avg_time'] * 1000
        
        # Tamaño del punto proporcional al tiempo de inferencia
        ax.scatter(n_params, miou, s=inf_time * 5, color=colors_models[i], 
                   edgecolors='black', linewidth=1.5, zorder=5, alpha=0.8)
        ax.annotate(name, (n_params, miou), 
                    xytext=(10, 10), textcoords='offset points',
                    fontsize=11, fontweight='bold',
                    arrowprops=dict(arrowstyle='->', color='gray'))
    
    ax.set_xlabel('Millones de parámetros', fontsize=12)
    ax.set_ylabel('mIoU', fontsize=12)
    ax.set_title('Trade-off Precisión vs Tamaño del modelo\n(tamaño del punto ∝ tiempo de inferencia)', fontsize=14)
    ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.show()
    
    # Eficiencia: mIoU por millón de parámetros
    print('\nEficiencia (mIoU por millón de parámetros):')
    for name in model_names:
        n_params = sum(p.numel() for p in models[name].parameters()) / 1e6
        efficiency = results[name]['mIoU'] / n_params
        print(f'  {name}: {efficiency:.4f} mIoU/M_params')

if results:
    print('Análisis de clases más difíciles y más fáciles:')
    print('=' * 70)
    
    for name in model_names:
        class_ious = results[name]['class_ious']
        
        # Filtramos NaN
        valid_classes = [(i, iou) for i, iou in enumerate(class_ious) if not np.isnan(iou)]
        valid_classes.sort(key=lambda x: x[1])
        
        print(f'\n{name} (mIoU = {results[name]["mIoU"]:.4f}):')
        
        print('  Top 3 más fáciles:')
        for idx, iou in valid_classes[-3:]:
            print(f'    {VOC_CLASSES[idx]}: IoU = {iou:.4f}')
        
        print('  Top 3 más difíciles:')
        for idx, iou in valid_classes[:3]:
            print(f'    {VOC_CLASSES[idx]}: IoU = {iou:.4f}')

# Comparación detallada con overlay de errores
if voc_dataset is not None:
    # Seleccionamos 3 imágenes interesantes
    detail_indices = [25, 75, 180]
    
    for idx in detail_indices:
        img_pil, mask_pil = voc_dataset[idx]
        img_np = np.array(img_pil)
        gt = np.array(mask_pil)
        
        fig, axes = plt.subplots(2, len(model_names) + 1, figsize=(5 * (len(model_names) + 1), 9))
        
        # Fila superior: imagen + predicciones coloreadas
        axes[0, 0].imshow(img_np)
        axes[0, 0].set_title('Original', fontsize=12)
        axes[0, 0].axis('off')
        
        # GT en fila inferior izquierda
        gt_rgb = decode_segmap(gt)
        axes[1, 0].imshow(gt_rgb)
        axes[1, 0].set_title('Ground Truth', fontsize=12)
        axes[1, 0].axis('off')
        
        for col, name in enumerate(model_names):
            pred = predict_segmentation(models[name], img_pil)
            
            # Redimensionamos si es necesario
            if pred.shape != gt.shape:
                pred_tensor = torch.tensor(pred).unsqueeze(0).unsqueeze(0).float()
                pred = F.interpolate(
                    pred_tensor, size=gt.shape, mode='nearest'
                ).squeeze().numpy().astype(int)
            
            # Predicción coloreada
            pred_rgb = decode_segmap(pred)
            axes[0, col + 1].imshow(pred_rgb)
            
            # IoU para esta imagen
            ious = compute_iou_per_class(pred, gt)
            miou_img = np.nanmean([iou for iou in ious if not np.isnan(iou)])
            axes[0, col + 1].set_title(f'{name}\nmIoU={miou_img:.3f}', fontsize=11)
            axes[0, col + 1].axis('off')
            
            # Mapa de errores (rojo = error)
            valid_mask = gt != 255
            error_map = np.zeros((*gt.shape, 3), dtype=np.uint8)
            correct = (pred == gt) & valid_mask
            wrong = (pred != gt) & valid_mask
            error_map[correct] = [0, 180, 0]   # Verde: correcto
            error_map[wrong] = [220, 0, 0]      # Rojo: error
            error_map[~valid_mask] = [128, 128, 128]  # Gris: ignorado
            
            axes[1, col + 1].imshow(error_map)
            pixel_acc = correct.sum() / valid_mask.sum() if valid_mask.sum() > 0 else 0
            axes[1, col + 1].set_title(f'Errores ({pixel_acc:.1%} acc)', fontsize=11)
            axes[1, col + 1].axis('off')
        
        plt.suptitle(f'Imagen #{idx}: comparación detallada', fontsize=14)
        plt.tight_layout()
        plt.show()

import pandas as pd

if results:
    summary_data = []
    
    for name in model_names:
        n_params = sum(p.numel() for p in models[name].parameters()) / 1e6
        summary_data.append({
            'Modelo': name,
            'Backbone': 'ResNet50' if 'ResNet' in name else 'MobileNetV3',
            'Parámetros (M)': f'{n_params:.1f}',
            'mIoU': f'{results[name]["mIoU"]:.4f}',
            'Tiempo inf. (ms)': f'{results[name]["avg_time"]*1000:.1f}',
            'FPS': f'{1/results[name]["avg_time"]:.1f}',
            'Eficiencia (mIoU/M)': f'{results[name]["mIoU"]/n_params:.4f}'
        })
    
    df_summary = pd.DataFrame(summary_data)
    display(df_summary)
    
    print('\nObservaciones:')
    
    # Encontramos el mejor modelo por mIoU
    best_miou_name = max(model_names, key=lambda n: results[n]['mIoU'])
    best_efficiency_name = max(model_names, key=lambda n: results[n]['mIoU'] / 
                               (sum(p.numel() for p in models[n].parameters()) / 1e6))
    fastest_name = min(model_names, key=lambda n: results[n]['avg_time'])
    
    print(f'  Mejor mIoU: {best_miou_name}')
    print(f'  Más eficiente: {best_efficiency_name}')
    print(f'  Más rápido: {fastest_name}')