# Librerías y configuración

import os
import glob
import zipfile
import urllib.request
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from sklearn.model_selection import train_test_split

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms.functional as TF

# 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__}')

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

# Descarga del dataset Kvasir-SEG
import ssl

DATA_DIR = '/tmp/kvasir_seg'
os.makedirs(DATA_DIR, exist_ok=True)

ZIP_URL = 'https://datasets.simula.no/downloads/kvasir-seg.zip'
ZIP_PATH = os.path.join(DATA_DIR, 'kvasir-seg.zip')

# Crear contexto SSL sin verificación (necesario en algunos entornos)
ssl_context = ssl.create_default_context()
ssl_context.check_hostname = False
ssl_context.verify_mode = ssl.CERT_NONE

if not os.path.exists(os.path.join(DATA_DIR, 'Kvasir-SEG')):
    if not os.path.exists(ZIP_PATH):
        print('Descargando Kvasir-SEG dataset (~46 MB)...')
        with urllib.request.urlopen(ZIP_URL, context=ssl_context) as response, \
             open(ZIP_PATH, 'wb') as out_file:
            out_file.write(response.read())
        print('Descarga completada.')
    
    print('Extrayendo archivos...')
    with zipfile.ZipFile(ZIP_PATH, 'r') as zip_ref:
        zip_ref.extractall(DATA_DIR)
    print('Extracción completada.')
else:
    print('Dataset ya descargado.')

# Rutas de imágenes y máscaras
IMAGES_DIR = os.path.join(DATA_DIR, 'Kvasir-SEG', 'images')
MASKS_DIR = os.path.join(DATA_DIR, 'Kvasir-SEG', 'masks')

image_files = sorted(glob.glob(os.path.join(IMAGES_DIR, '*.jpg')))
mask_files = sorted(glob.glob(os.path.join(MASKS_DIR, '*.jpg')))

print(f'\nImágenes encontradas: {len(image_files)}')
print(f'Máscaras encontradas: {len(mask_files)}')

# Visualizamos algunas muestras del dataset
fig, axes = plt.subplots(3, 4, figsize=(16, 12))

for i in range(3):
    idx = i * 100  # Muestreamos cada 100 imágenes
    img = np.array(Image.open(image_files[idx]))
    mask = np.array(Image.open(mask_files[idx]).convert('L'))
    
    axes[i, 0].imshow(img)
    axes[i, 0].set_title('Imagen original' if i == 0 else '')
    axes[i, 0].axis('off')
    
    axes[i, 1].imshow(mask, cmap='gray')
    axes[i, 1].set_title('Máscara GT' if i == 0 else '')
    axes[i, 1].axis('off')
    
    # Superposición
    axes[i, 2].imshow(img)
    axes[i, 2].imshow(mask, cmap='Reds', alpha=0.4)
    axes[i, 2].set_title('Superposición' if i == 0 else '')
    axes[i, 2].axis('off')
    
    # Contorno
    contour = np.zeros_like(mask)
    import cv2
    contours, _ = cv2.findContours((mask > 127).astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    cv2.drawContours(contour, contours, -1, 255, 2)
    axes[i, 3].imshow(img)
    axes[i, 3].imshow(contour, cmap='Greens', alpha=0.7)
    axes[i, 3].set_title('Contorno del pólipo' if i == 0 else '')
    axes[i, 3].axis('off')

plt.suptitle('Dataset Kvasir-SEG: segmentación de pólipos gastrointestinales', fontsize=14)
plt.tight_layout()
plt.show()

IMG_SIZE = 256


class KvasirSegDataset(Dataset):
    """Dataset para segmentación de pólipos Kvasir-SEG."""
    
    def __init__(self, image_paths, mask_paths, augment=False):
        self.image_paths = image_paths
        self.mask_paths = mask_paths
        self.augment = augment
    
    def __len__(self):
        return len(self.image_paths)
    
    def __getitem__(self, idx):
        # Cargamos imagen y máscara
        image = Image.open(self.image_paths[idx]).convert('RGB')
        mask = Image.open(self.mask_paths[idx]).convert('L')
        
        # Redimensionamos
        image = TF.resize(image, [IMG_SIZE, IMG_SIZE])
        mask = TF.resize(mask, [IMG_SIZE, IMG_SIZE], interpolation=Image.NEAREST)
        
        # Data augmentation
        if self.augment:
            if torch.rand(1) > 0.5:
                image = TF.hflip(image)
                mask = TF.hflip(mask)
            if torch.rand(1) > 0.5:
                image = TF.vflip(image)
                mask = TF.vflip(mask)
            if torch.rand(1) > 0.5:
                angle = float(torch.randint(-30, 30, (1,)))
                image = TF.rotate(image, angle)
                mask = TF.rotate(mask, angle)
        
        # Convertimos a tensores
        image = TF.to_tensor(image)  # [3, H, W], rango [0, 1]
        mask = TF.to_tensor(mask)    # [1, H, W], rango [0, 1]
        mask = (mask > 0.5).float()  # Binarizamos
        
        return image, mask


# División train/val/test (70/15/15)
train_imgs, test_imgs, train_masks, test_masks = train_test_split(
    image_files, mask_files, test_size=0.15, random_state=42
)
train_imgs, val_imgs, train_masks, val_masks = train_test_split(
    train_imgs, train_masks, test_size=0.176, random_state=42  # 0.176 de 0.85 ≈ 0.15 total
)

train_dataset = KvasirSegDataset(train_imgs, train_masks, augment=True)
val_dataset = KvasirSegDataset(val_imgs, val_masks, augment=False)
test_dataset = KvasirSegDataset(test_imgs, test_masks, augment=False)

BATCH_SIZE = 8
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)
val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)

print(f'Conjunto de entrenamiento: {len(train_dataset)} imágenes')
print(f'Conjunto de validación: {len(val_dataset)} imágenes')
print(f'Conjunto de test: {len(test_dataset)} imágenes')

# Verificamos un batch
imgs, masks = next(iter(train_loader))
print(f'\nBatch shape - imagen: {imgs.shape}, máscara: {masks.shape}')
print(f'Rango imagen: [{imgs.min():.2f}, {imgs.max():.2f}]')
print(f'Valores únicos máscara: {masks.unique().tolist()}')

class DoubleConv(nn.Module):
    """Bloque básico: (Conv3x3 -> BN -> ReLU) x 2"""
    
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True)
        )
    
    def forward(self, x):
        return self.double_conv(x)


class UNet(nn.Module):
    """
    U-Net para segmentación binaria.
    
    Arquitectura:
        Encoder: 4 niveles de downsampling
        Bottleneck: bloque convolucional central
        Decoder: 4 niveles de upsampling con skip connections
        Output: Conv 1x1 → Sigmoid
    """
    
    def __init__(self, in_channels=3, out_channels=1, features=[64, 128, 256, 512]):
        super().__init__()
        self.downs = nn.ModuleList()
        self.ups = nn.ModuleList()
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        
        # Encoder (camino de contracción)
        for feature in features:
            self.downs.append(DoubleConv(in_channels, feature))
            in_channels = feature
        
        # Bottleneck
        self.bottleneck = DoubleConv(features[-1], features[-1] * 2)
        
        # Decoder (camino de expansión)
        for feature in reversed(features):
            # ConvTranspose para upsampling
            self.ups.append(
                nn.ConvTranspose2d(feature * 2, feature, kernel_size=2, stride=2)
            )
            # DoubleConv después de concatenar
            self.ups.append(DoubleConv(feature * 2, feature))
        
        # Capa final: 1x1 conv
        self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)
    
    def forward(self, x):
        skip_connections = []
        
        # Encoder
        for down in self.downs:
            x = down(x)
            skip_connections.append(x)
            x = self.pool(x)
        
        # Bottleneck
        x = self.bottleneck(x)
        
        # Decoder (skip connections en orden inverso)
        skip_connections = skip_connections[::-1]
        for idx in range(0, len(self.ups), 2):
            x = self.ups[idx](x)  # Upsampling
            skip = skip_connections[idx // 2]
            
            # Ajuste de tamaño si es necesario (por resoluciones impares)
            if x.shape != skip.shape:
                x = F.interpolate(x, size=skip.shape[2:], mode='bilinear', align_corners=True)
            
            x = torch.cat([skip, x], dim=1)  # Concatenación (skip connection)
            x = self.ups[idx + 1](x)         # DoubleConv
        
        return torch.sigmoid(self.final_conv(x))


# Instanciamos el modelo
model = UNet(in_channels=3, out_channels=1).to(device)

# Resumen del modelo
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f'Parámetros totales: {total_params:,}')
print(f'Parámetros entrenables: {trainable_params:,}')

# Verificamos con una entrada dummy
dummy = torch.randn(1, 3, IMG_SIZE, IMG_SIZE).to(device)
output = model(dummy)
print(f'\nEntrada: {dummy.shape}')
print(f'Salida: {output.shape}')
print(f'Rango salida: [{output.min():.4f}, {output.max():.4f}]')

class DiceLoss(nn.Module):
    """Dice Loss para segmentación binaria."""
    def __init__(self, smooth=1e-6):
        super().__init__()
        self.smooth = smooth
    
    def forward(self, pred, target):
        pred_flat = pred.view(-1)
        target_flat = target.view(-1)
        intersection = (pred_flat * target_flat).sum()
        return 1 - (2. * intersection + self.smooth) / (
            pred_flat.sum() + target_flat.sum() + self.smooth
        )


class BCEDiceLoss(nn.Module):
    """Combinación de BCE y Dice Loss."""
    def __init__(self, bce_weight=0.5):
        super().__init__()
        self.bce = nn.BCELoss()
        self.dice = DiceLoss()
        self.bce_weight = bce_weight
    
    def forward(self, pred, target):
        return self.bce_weight * self.bce(pred, target) + (1 - self.bce_weight) * self.dice(pred, target)


def compute_metrics(pred, target, threshold=0.5):
    """Calcula métricas de segmentación binaria."""
    pred_bin = (pred > threshold).float()
    
    # Flatten
    pred_flat = pred_bin.view(-1)
    target_flat = target.view(-1)
    
    tp = (pred_flat * target_flat).sum()
    fp = (pred_flat * (1 - target_flat)).sum()
    fn = ((1 - pred_flat) * target_flat).sum()
    
    precision = tp / (tp + fp + 1e-8)
    recall = tp / (tp + fn + 1e-8)
    dice = 2 * tp / (2 * tp + fp + fn + 1e-8)
    iou = tp / (tp + fp + fn + 1e-8)
    
    return {
        'dice': dice.item(),
        'iou': iou.item(),
        'precision': precision.item(),
        'recall': recall.item()
    }

print('Funciones de pérdida y métricas definidas.')

# Configuración del entrenamiento
NUM_EPOCHS = 10
LEARNING_RATE = 1e-4

criterion = BCEDiceLoss(bce_weight=0.5)
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', patience=5, factor=0.5)

# Almacenamos métricas
history = {
    'train_loss': [], 'val_loss': [],
    'train_dice': [], 'val_dice': [],
    'train_iou': [], 'val_iou': []
}

best_val_dice = 0.0
best_model_state = None

print(f'Entrenando U-Net durante {NUM_EPOCHS} épocas...')
print(f'Dispositivo: {device}')
print(f'Tamaño del batch: {BATCH_SIZE}')
print(f'Learning rate inicial: {LEARNING_RATE}')
print('-' * 70)

for epoch in range(NUM_EPOCHS):
    # ---------- ENTRENAMIENTO ----------
    model.train()
    train_loss = 0.0
    train_metrics = {'dice': 0, 'iou': 0}
    
    for images, masks in train_loader:
        images, masks = images.to(device), masks.to(device)
        
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, masks)
        loss.backward()
        optimizer.step()
        
        train_loss += loss.item()
        metrics = compute_metrics(outputs, masks)
        train_metrics['dice'] += metrics['dice']
        train_metrics['iou'] += metrics['iou']
    
    n_train = len(train_loader)
    train_loss /= n_train
    train_metrics = {k: v / n_train for k, v in train_metrics.items()}
    
    # ---------- VALIDACIÓN ----------
    model.eval()
    val_loss = 0.0
    val_metrics = {'dice': 0, 'iou': 0}
    
    with torch.no_grad():
        for images, masks in val_loader:
            images, masks = images.to(device), masks.to(device)
            outputs = model(images)
            loss = criterion(outputs, masks)
            
            val_loss += loss.item()
            metrics = compute_metrics(outputs, masks)
            val_metrics['dice'] += metrics['dice']
            val_metrics['iou'] += metrics['iou']
    
    n_val = len(val_loader)
    val_loss /= n_val
    val_metrics = {k: v / n_val for k, v in val_metrics.items()}
    
    # Scheduler
    scheduler.step(val_loss)
    
    # Guardamos métricas
    history['train_loss'].append(train_loss)
    history['val_loss'].append(val_loss)
    history['train_dice'].append(train_metrics['dice'])
    history['val_dice'].append(val_metrics['dice'])
    history['train_iou'].append(train_metrics['iou'])
    history['val_iou'].append(val_metrics['iou'])
    
    # Guardamos el mejor modelo
    if val_metrics['dice'] > best_val_dice:
        best_val_dice = val_metrics['dice']
        best_model_state = model.state_dict().copy()
    
    # Log cada 5 épocas
    if (epoch + 1) % 5 == 0 or epoch == 0:
        current_lr = optimizer.param_groups[0]['lr']
        print(f'Época {epoch+1:3d}/{NUM_EPOCHS} | '
              f'Loss: train={train_loss:.4f} val={val_loss:.4f} | '
              f'Dice: train={train_metrics["dice"]:.4f} val={val_metrics["dice"]:.4f} | '
              f'IoU: train={train_metrics["iou"]:.4f} val={val_metrics["iou"]:.4f} | '
              f'LR={current_lr:.2e}')

print('-' * 70)
print(f'Mejor Dice en validación: {best_val_dice:.4f}')

# Cargamos el mejor modelo
if best_model_state is not None:
    model.load_state_dict(best_model_state)

fig, axes = plt.subplots(1, 3, figsize=(18, 5))
epochs_range = range(1, NUM_EPOCHS + 1)

# Loss
axes[0].plot(epochs_range, history['train_loss'], 'b-', label='Train', linewidth=2)
axes[0].plot(epochs_range, history['val_loss'], 'r-', label='Validación', linewidth=2)
axes[0].set_title('Pérdida (BCE + Dice)', fontsize=13)
axes[0].set_xlabel('Época')
axes[0].set_ylabel('Loss')
axes[0].legend()
axes[0].grid(True, alpha=0.3)

# Dice Score
axes[1].plot(epochs_range, history['train_dice'], 'b-', label='Train', linewidth=2)
axes[1].plot(epochs_range, history['val_dice'], 'r-', label='Validación', linewidth=2)
axes[1].set_title('Dice Score', fontsize=13)
axes[1].set_xlabel('Época')
axes[1].set_ylabel('Dice')
axes[1].legend()
axes[1].grid(True, alpha=0.3)

# IoU
axes[2].plot(epochs_range, history['train_iou'], 'b-', label='Train', linewidth=2)
axes[2].plot(epochs_range, history['val_iou'], 'r-', label='Validación', linewidth=2)
axes[2].set_title('IoU (Intersection over Union)', fontsize=13)
axes[2].set_xlabel('Época')
axes[2].set_ylabel('IoU')
axes[2].legend()
axes[2].grid(True, alpha=0.3)

plt.suptitle('Curvas de entrenamiento U-Net', fontsize=15)
plt.tight_layout()
plt.show()

# Evaluación final en test
model.eval()
test_metrics_accum = {'dice': 0, 'iou': 0, 'precision': 0, 'recall': 0}
n_batches = 0

all_preds = []
all_targets = []
all_images = []

with torch.no_grad():
    for images, masks in test_loader:
        images, masks = images.to(device), masks.to(device)
        outputs = model(images)
        
        metrics = compute_metrics(outputs, masks)
        for k in test_metrics_accum:
            test_metrics_accum[k] += metrics[k]
        n_batches += 1
        
        all_images.append(images.cpu())
        all_preds.append(outputs.cpu())
        all_targets.append(masks.cpu())

test_metrics = {k: v / n_batches for k, v in test_metrics_accum.items()}

print('Métricas en el conjunto de TEST:')
print(f'  Dice Score:  {test_metrics["dice"]:.4f}')
print(f'  IoU:         {test_metrics["iou"]:.4f}')
print(f'  Precision:   {test_metrics["precision"]:.4f}')
print(f'  Recall:      {test_metrics["recall"]:.4f}')

# Concatenamos todo
all_images = torch.cat(all_images, dim=0)
all_preds = torch.cat(all_preds, dim=0)
all_targets = torch.cat(all_targets, dim=0)

# Visualizamos predicciones en test
n_samples = 6
indices = np.linspace(0, len(all_images) - 1, n_samples, dtype=int)

fig, axes = plt.subplots(n_samples, 4, figsize=(16, n_samples * 3.5))
col_titles = ['Imagen original', 'Ground Truth', 'Predicción U-Net', 'Superposición']

for row, idx in enumerate(indices):
    img = all_images[idx].permute(1, 2, 0).numpy()
    gt = all_targets[idx, 0].numpy()
    pred = all_preds[idx, 0].numpy()
    pred_bin = (pred > 0.5).astype(float)
    
    # Dice individual
    dice_i = 2 * (pred_bin * gt).sum() / (pred_bin.sum() + gt.sum() + 1e-8)
    
    axes[row, 0].imshow(img)
    axes[row, 0].axis('off')
    
    axes[row, 1].imshow(gt, cmap='gray')
    axes[row, 1].axis('off')
    
    axes[row, 2].imshow(pred_bin, cmap='gray')
    axes[row, 2].set_title(f'Dice={dice_i:.3f}', fontsize=10)
    axes[row, 2].axis('off')
    
    # Superposición: verde=TP, rojo=FP, azul=FN
    overlay = img.copy()
    tp_mask = (pred_bin == 1) & (gt == 1)
    fp_mask = (pred_bin == 1) & (gt == 0)
    fn_mask = (pred_bin == 0) & (gt == 1)
    overlay[tp_mask] = overlay[tp_mask] * 0.5 + np.array([0, 1, 0]) * 0.5  # Verde: acierto
    overlay[fp_mask] = overlay[fp_mask] * 0.5 + np.array([1, 0, 0]) * 0.5  # Rojo: falso positivo
    overlay[fn_mask] = overlay[fn_mask] * 0.5 + np.array([0, 0, 1]) * 0.5  # Azul: falso negativo
    axes[row, 3].imshow(overlay)
    axes[row, 3].axis('off')
    
    if row == 0:
        for col, title in enumerate(col_titles):
            axes[row, col].set_title(title, fontsize=12)

plt.suptitle('Predicciones U-Net en el conjunto de test\nVerde=TP, Rojo=FP, Azul=FN', fontsize=14)
plt.tight_layout()
plt.show()

# Visualizamos feature maps del encoder
model.eval()

# Tomamos una imagen de ejemplo
sample_img = all_images[0:1].to(device)

# Extraemos feature maps del encoder
encoder_features = []
x = sample_img
with torch.no_grad():
    for down in model.downs:
        x = down(x)
        encoder_features.append(x.cpu())
        x = model.pool(x)

# Visualizamos el primer canal de cada nivel del encoder
fig, axes = plt.subplots(1, 5, figsize=(20, 4))

axes[0].imshow(sample_img[0].cpu().permute(1, 2, 0))
axes[0].set_title(f'Imagen original\n{sample_img.shape[2]}×{sample_img.shape[3]}', fontsize=11)
axes[0].axis('off')

for i, feat in enumerate(encoder_features):
    # Promedio de los canales para visualizar
    feat_map = feat[0].mean(dim=0).numpy()
    axes[i + 1].imshow(feat_map, cmap='viridis')
    axes[i + 1].set_title(f'Encoder nivel {i+1}\n{feat.shape[2]}×{feat.shape[3]}, {feat.shape[1]} canales', fontsize=11)
    axes[i + 1].axis('off')

plt.suptitle('Feature maps del encoder: de resolución alta a baja', fontsize=14)
plt.tight_layout()
plt.show()

print('Los niveles más profundos capturan semántica (qué es un pólipo).')
print('Los niveles superficiales capturan detalles (bordes, texturas).')
print('Las skip connections combinan ambos → segmentación precisa.')

# Calculamos Dice score por imagen individual
individual_dices = []

for i in range(len(all_preds)):
    pred_i = (all_preds[i, 0] > 0.5).float().numpy()
    gt_i = all_targets[i, 0].numpy()
    dice_i = 2 * (pred_i * gt_i).sum() / (pred_i.sum() + gt_i.sum() + 1e-8)
    individual_dices.append(dice_i)

individual_dices = np.array(individual_dices)

fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# Histograma
axes[0].hist(individual_dices, bins=25, color='steelblue', edgecolor='white', alpha=0.8)
axes[0].axvline(x=individual_dices.mean(), color='red', linestyle='--', linewidth=2,
                label=f'Media = {individual_dices.mean():.3f}')
axes[0].axvline(x=np.median(individual_dices), color='orange', linestyle='--', linewidth=2,
                label=f'Mediana = {np.median(individual_dices):.3f}')
axes[0].set_title('Distribución de Dice Score por imagen', fontsize=12)
axes[0].set_xlabel('Dice Score')
axes[0].set_ylabel('Número de imágenes')
axes[0].legend()

# Sorted plot
sorted_dices = np.sort(individual_dices)
colors = ['crimson' if d < 0.5 else 'orange' if d < 0.7 else 'forestgreen' for d in sorted_dices]
axes[1].bar(range(len(sorted_dices)), sorted_dices, color=colors, width=1)
axes[1].axhline(y=0.5, color='gray', linestyle='--', alpha=0.5)
axes[1].axhline(y=0.7, color='gray', linestyle=':', alpha=0.5)
axes[1].set_title('Dice Score ordenado por imagen', fontsize=12)
axes[1].set_xlabel('Imagen (ordenada)')
axes[1].set_ylabel('Dice Score')

plt.tight_layout()
plt.show()

print(f'\nEstadísticas del Dice Score por imagen:')
print(f'  Media:    {individual_dices.mean():.4f}')
print(f'  Mediana:  {np.median(individual_dices):.4f}')
print(f'  Std:      {individual_dices.std():.4f}')
print(f'  Min:      {individual_dices.min():.4f}')
print(f'  Max:      {individual_dices.max():.4f}')
print(f'  % > 0.7:  {(individual_dices > 0.7).mean():.1%}')
print(f'  % > 0.5:  {(individual_dices > 0.5).mean():.1%}')