import os
import glob
import json
import numpy as np
import pandas as pd
from PIL import Image
from tqdm import tqdm
from datetime import datetime
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from scipy.ndimage import morphology

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts, ReduceLROnPlateau

from transformers import AutoModel
import albumentations as A
from albumentations.pytorch import ToTensorV2

import cv2
import warnings
import math
warnings.filterwarnings('ignore')

# Set seeds for reproducibility
def set_seed(seed=42):
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

set_seed(42)

# ============================================================================
# CONFIGURATION
# ============================================================================

class Config:
    # Model - USING YOUR LOCAL DOWNLOADED MODEL
    model_name = "facebook/dinov3-vitl16-pretrain-lvd1689m"
    local_model_path = "/data/F/VoiceNegar/models/pe_models/dino7b/checkpoints/initial_dinov3-vitl16-pretrain-lvd1689m_backbone"
    
    # Data paths
    dataset_path = "/home/PeBigModelForVilab/dinov3/toy-project/Kvasir-SEG/"
    image_size = 256
    patch_size = 16
    
    # Training
    batch_size = 96
    num_epochs = 150
    learning_rate = 1e-4
    min_lr = 1e-6
    weight_decay = 1e-4
    
    # Cosine Annealing with Warm Restarts
    T_0 = 10  # Initial restart period (epochs)
    T_mult = 2  # Period multiplier after each restart
    
    # Validation
    val_split = 0.1
    test_split = 0.05
    
    # Device
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    # Logging
    save_dir = "./checkpoints"
    log_interval = 10
    
    # Image normalization (ImageNet stats)
    mean = [0.485, 0.456, 0.406]
    std = [0.229, 0.224, 0.225]
    
    resume_from = None
    # Multi‑scale ViT layers
    multi_scale_layers = [5, 10, 16, 18, 20, 22, 23]
    # Loss parameters (Focal+Dice)
    focal_weight = 0.69
    dice_weight = 0.3
    boundary_weight = 0.01
    # HD95 parameter
    hd95_threshold = 0.5

config = Config()
os.makedirs(config.save_dir, exist_ok=True)
print(f"Using device: {config.device}")
print(f"Model: {config.model_name}")
print(f"Local model path: {config.local_model_path}")
print(f"Exists: {os.path.exists(config.local_model_path)}")

# ============================================================================
# DATASET CLASS
# ============================================================================

class PolypDataset(Dataset):
    """Kvasir-SEG dataset with manual preprocessing"""
    
    def __init__(self, image_paths, mask_paths, transform=None, target_size=(256, 256)):
        self.image_paths = image_paths
        self.mask_paths = mask_paths
        self.transform = transform
        self.target_size = target_size
        
        # ImageNet normalization values
        self.mean = torch.tensor(config.mean).view(3, 1, 1)
        self.std = torch.tensor(config.std).view(3, 1, 1)
        
    def __len__(self):
        return len(self.image_paths)
    
    def __getitem__(self, idx):
        # Load image
        image = cv2.imread(self.image_paths[idx])
        if image is None:
            raise ValueError(f"Could not load image: {self.image_paths[idx]}")
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        # Load mask
        mask = cv2.imread(self.mask_paths[idx], cv2.IMREAD_GRAYSCALE)
        if mask is None:
            raise ValueError(f"Could not load mask: {self.mask_paths[idx]}")
        mask = (mask > 127).astype(np.float32)
        
        # Apply augmentations
        if self.transform:
            augmented = self.transform(image=image, mask=mask)
            image, mask = augmented['image'], augmented['mask']
        else:
            image = cv2.resize(image, self.target_size)
            mask = cv2.resize(mask, self.target_size, interpolation=cv2.INTER_NEAREST)
        
        # Manual preprocessing
        if isinstance(image, np.ndarray):
            image = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0
        elif isinstance(image, torch.Tensor):
            image = image.float() / 255.0
        
        # Apply ImageNet normalization
        image = (image - self.mean) / self.std
        
        # Ensure mask is tensor
        if isinstance(mask, np.ndarray):
            mask = torch.from_numpy(mask).float()
        
        return image, mask.unsqueeze(0)

# ============================================================================
# FIXED DINOv3 ENCODER
# ============================================================================

class DINOv3Encoder(nn.Module):
    """Frozen DINOv3 encoder that can return concatenated multi‑scale features."""
    
    def __init__(self, model_name="facebook/dinov3-vitl16-pretrain-lvd1689m", 
                 local_path=None, freeze=True, layers=None):
        super().__init__()
        
        # Load model
        if local_path and os.path.exists(local_path):
            print(f"Loading DINOv3 model from local path: {local_path}")
            self.model = AutoModel.from_pretrained(local_path, local_files_only=True)
        else:
            print(f"Loading DINOv3 from HuggingFace hub: {model_name}")
            self.model = AutoModel.from_pretrained(model_name)
        
        self.embed_dim = self.model.config.hidden_size
        self.patch_size = self.model.config.patch_size
        self.layers = layers
        
        if self.layers is not None:
            self.out_channels = self.embed_dim * len(self.layers)
        else:
            self.out_channels = self.embed_dim
        
        print(f"DINOv3 loaded - embed_dim: {self.embed_dim}, patch_size: {self.patch_size}")
        if self.layers:
            print(f"   Multi‑scale layers: {self.layers}, output channels: {self.out_channels}")
        
        if freeze:
            for param in self.model.parameters():
                param.requires_grad = False
    
    def _reshape_to_2d(self, patch_tokens, B):
        """Robust reshaping of patch tokens to 2D grid."""
        N = patch_tokens.shape[1]
        D = patch_tokens.shape[2]
        
        H_grid = int(math.sqrt(N))
        W_grid = H_grid
        
        while H_grid * W_grid != N:
            if H_grid * W_grid < N:
                W_grid += 1
            else:
                found = False
                for h in range(int(math.sqrt(N)), 0, -1):
                    if N % h == 0:
                        H_grid = h
                        W_grid = N // h
                        found = True
                        break
                if not found:
                    W_grid += 1
                else:
                    break
        
        if H_grid * W_grid != N:
            print(f"   Warning: Cannot reshape {N} patches into grid. Interpolating to square.")
            target_size = int(math.sqrt(N))
            patch_tokens_flat = patch_tokens.transpose(1, 2)
            patch_tokens_2d = F.interpolate(
                patch_tokens_flat.unsqueeze(-2) if patch_tokens_flat.dim() == 3 else patch_tokens_flat,
                size=target_size * target_size,
                mode='linear',
                align_corners=False
            ).reshape(B, D, target_size, target_size)
            return patch_tokens_2d
        
        feat_map = patch_tokens.transpose(1, 2).reshape(B, D, H_grid, W_grid)
        return feat_map
    
    def forward(self, pixel_values):
        B, C, H, W = pixel_values.shape
        
        if self.layers is not None:
            outputs = self.model(pixel_values, output_hidden_states=True)
            hidden_states = outputs.hidden_states
            
            feature_list = []
            for idx in self.layers:
                hidden = hidden_states[idx]
                patch_tokens = hidden[:, 1:, :]
                feat_map = self._reshape_to_2d(patch_tokens, B)
                feature_list.append(feat_map)
            
            target_h, target_w = feature_list[0].shape[-2:]
            
            resized_features = []
            for feat in feature_list:
                if feat.shape[-2:] != (target_h, target_w):
                    feat = F.interpolate(feat, size=(target_h, target_w), 
                                        mode='bilinear', align_corners=False)
                resized_features.append(feat)
            
            features = torch.cat(resized_features, dim=1)
        else:
            outputs = self.model(pixel_values, output_hidden_states=False)
            last_hidden = outputs.last_hidden_state[:, 1:, :]
            features = self._reshape_to_2d(last_hidden, B)
        
        return features

# ============================================================================
# SHALLOW STEM FOR SKIP CONNECTIONS
# ============================================================================
class ShallowStem(nn.Module):
    """Extracts multi‑scale features from the input image."""
    def __init__(self, in_channels=3, base_channels=64):
        super().__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels, base_channels, 3, padding=1, bias=False),
            nn.BatchNorm2d(base_channels),
            nn.ReLU(inplace=True)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(base_channels, base_channels*2, 3, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(base_channels*2),
            nn.ReLU(inplace=True)
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(base_channels*2, base_channels*4, 3, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(base_channels*4),
            nn.ReLU(inplace=True)
        )
        self.conv4 = nn.Sequential(
            nn.Conv2d(base_channels*4, base_channels*8, 3, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(base_channels*8),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x = self.conv1(x)
        f2 = self.conv2(x)
        f3 = self.conv3(f2)
        f4 = self.conv4(f3)
        return [f4, f3, f2]

# ============================================================================
# U‑Net DECODER WITH SKIP CONNECTIONS
# ============================================================================
class UNetDecoder(nn.Module):
    """Decoder that progressively upsamples ViT features."""
    def __init__(self, vit_channels=1024, stem_channels=[512,256,128], num_classes=1):
        super().__init__()
        self.up1 = self._up_block(vit_channels, 256)
        self.conv1 = self._conv_block(256 + stem_channels[0], 256)

        self.up2 = self._up_block(256, 128)
        self.conv2 = self._conv_block(128 + stem_channels[1], 128)

        self.up3 = self._up_block(128, 64)
        self.conv3 = self._conv_block(64 + stem_channels[2], 64)

        self.up4 = nn.UpsamplingBilinear2d(scale_factor=2)
        self.final = nn.Conv2d(64, num_classes, kernel_size=1)

    def _up_block(self, in_ch, out_ch):
        return nn.Sequential(
            nn.UpsamplingBilinear2d(scale_factor=2),
            nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True)
        )

    def _conv_block(self, in_ch, out_ch):
        return nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True)
        )

    def forward(self, vit_features, skip_features):
        x = self.up1(vit_features)
        
        if x.shape[-2:] != skip_features[0].shape[-2:]:
            x = F.interpolate(x, size=skip_features[0].shape[-2:], mode='bilinear', align_corners=False)
        
        x = torch.cat([x, skip_features[0]], dim=1)
        x = self.conv1(x)

        x = self.up2(x)
        if x.shape[-2:] != skip_features[1].shape[-2:]:
            x = F.interpolate(x, size=skip_features[1].shape[-2:], mode='bilinear', align_corners=False)
        
        x = torch.cat([x, skip_features[1]], dim=1)
        x = self.conv2(x)

        x = self.up3(x)
        if x.shape[-2:] != skip_features[2].shape[-2:]:
            x = F.interpolate(x, size=skip_features[2].shape[-2:], mode='bilinear', align_corners=False)
        
        x = torch.cat([x, skip_features[2]], dim=1)
        x = self.conv3(x)

        x = self.up4(x)
        return self.final(x)

# ============================================================================
# LOSS FUNCTIONS
# ============================================================================

class DiceLoss(nn.Module):
    def __init__(self, smooth=1e-6):
        super().__init__()
        self.smooth = smooth
    
    def forward(self, pred, target):
        pred = torch.sigmoid(pred)
        pred_flat = pred.view(-1)
        target_flat = target.view(-1)
        
        intersection = (pred_flat * target_flat).sum()
        dice = (2. * intersection + self.smooth) / (pred_flat.sum() + target_flat.sum() + self.smooth)
        
        return 1 - dice


class FocalLoss(nn.Module):
    def __init__(self, alpha=0.25, gamma=2.0):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma

    def forward(self, pred, target):
        bce = F.binary_cross_entropy_with_logits(pred, target, reduction='none')
        pt = torch.exp(-bce)
        focal = self.alpha * (1 - pt) ** self.gamma * bce
        return focal.mean()
    
class BoundaryLoss(nn.Module):
    """Boundary loss using Sobel edge detection for sharper edges"""
    def __init__(self):
        super().__init__()
        # Sobel kernels for edge detection
        self.sobel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32).view(1, 1, 3, 3)
        self.sobel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32).view(1, 1, 3, 3)
    
    def forward(self, pred, target):
        device = pred.device
        self.sobel_x = self.sobel_x.to(device)
        self.sobel_y = self.sobel_y.to(device)
        
        # Get probabilities
        pred_prob = torch.sigmoid(pred)
        
        # Compute edges for prediction and target
        pred_edges_x = F.conv2d(pred_prob, self.sobel_x, padding=1)
        pred_edges_y = F.conv2d(pred_prob, self.sobel_y, padding=1)
        pred_edges = torch.sqrt(pred_edges_x**2 + pred_edges_y**2 + 1e-6)
        
        target_edges_x = F.conv2d(target, self.sobel_x, padding=1)
        target_edges_y = F.conv2d(target, self.sobel_y, padding=1)
        target_edges = torch.sqrt(target_edges_x**2 + target_edges_y**2 + 1e-6)
        
        # MSE between edge maps
        boundary_loss = F.mse_loss(pred_edges, target_edges)
        return boundary_loss

class FocalDiceBoundaryLoss(nn.Module):
    def __init__(self, focal_weight=0.6, dice_weight=0.3, boundary_weight=0.1):
        super().__init__()
        self.focal = FocalLoss()
        self.dice = DiceLoss()
        self.boundary = BoundaryLoss()
        self.w_f = focal_weight
        self.w_d = dice_weight
        self.w_b = boundary_weight

    def forward(self, pred, target):
        return (self.w_f * self.focal(pred, target) + 
                self.w_d * self.dice(pred, target) + 
                self.w_b * self.boundary(pred, target))
# ============================================================================
# METRICS
# ============================================================================

def compute_dice(pred, target, threshold=0.5):
    """Compute Dice score"""
    pred_binary = (torch.sigmoid(pred) > threshold).float()
    intersection = (pred_binary * target).sum()
    dice = (2. * intersection) / (pred_binary.sum() + target.sum() + 1e-6)
    return dice.item()


def compute_iou(pred, target, threshold=0.5):
    """Compute IoU (Jaccard index)"""
    pred_binary = (torch.sigmoid(pred) > threshold).float()
    intersection = (pred_binary * target).sum()
    union = pred_binary.sum() + target.sum() - intersection
    iou = intersection / (union + 1e-6)
    return iou.item()


def compute_precision_recall(pred, target, threshold=0.5):
    """Compute precision and recall"""
    pred_binary = (torch.sigmoid(pred) > threshold).float()
    tp = (pred_binary * target).sum()
    fp = (pred_binary * (1 - target)).sum()
    fn = ((1 - pred_binary) * target).sum()
    
    precision = tp / (tp + fp + 1e-6)
    recall = tp / (tp + fn + 1e-6)
    
    return precision.item(), recall.item()


def compute_hd95(pred, target, threshold=0.5, voxel_spacing=None):
    """
    Compute Hausdorff Distance 95th percentile.
    
    Args:
        pred: Tensor [B, 1, H, W] logits
        target: Tensor [B, 1, H, W] ground truth
        threshold: threshold for binarization
        voxel_spacing: not used for 2D but kept for compatibility
    
    Returns:
        hd95: 95th percentile Hausdorff distance
    """
    # Convert to numpy and binarize
    pred_binary = (torch.sigmoid(pred) > threshold).float().cpu().numpy().squeeze()
    target_binary = target.cpu().numpy().squeeze()
    
    # Handle batch dimension
    if pred_binary.ndim == 3:
        hd95_values = []
        for i in range(pred_binary.shape[0]):
            hd95_values.append(_compute_hd95_single(pred_binary[i], target_binary[i]))
        return np.mean(hd95_values)
    else:
        return _compute_hd95_single(pred_binary, target_binary)


def _compute_hd95_single(pred, target):
    """Compute HD95 for a single 2D image"""
    if pred.sum() == 0 or target.sum() == 0:
        return 100.0  # Return a high value if either is empty
    
    # Get surface points
    pred_border = pred - morphology.binary_erosion(pred)
    target_border = target - morphology.binary_erosion(target)
    
    if pred_border.sum() == 0 or target_border.sum() == 0:
        return 100.0
    
    # Get coordinates of border points
    pred_coords = np.argwhere(pred_border > 0)
    target_coords = np.argwhere(target_border > 0)
    
    # Compute pairwise distances
    distances_pred_to_target = []
    for p in pred_coords:
        dist = np.min(np.sqrt(np.sum((target_coords - p) ** 2, axis=1)))
        distances_pred_to_target.append(dist)
    
    distances_target_to_pred = []
    for t in target_coords:
        dist = np.min(np.sqrt(np.sum((pred_coords - t) ** 2, axis=1)))
        distances_target_to_pred.append(dist)
    
    # Get 95th percentile
    all_distances = distances_pred_to_target + distances_target_to_pred
    hd95 = np.percentile(all_distances, 95)
    
    return hd95


def compute_all_metrics(pred, target, threshold=0.5):
    """Compute all metrics at once"""
    dice = compute_dice(pred, target, threshold)
    iou = compute_iou(pred, target, threshold)
    precision, recall = compute_precision_recall(pred, target, threshold)
    hd95 = compute_hd95(pred, target, threshold)
    
    return {
        'dice': dice,
        'iou': iou,
        'precision': precision,
        'recall': recall,
        'hd95': hd95
    }


def evaluate(decoder, stem, encoder, loader, device):
    """Comprehensive evaluation"""
    decoder.eval()
    stem.eval()
    encoder.eval()
    
    all_metrics = {
        'dice': [], 'iou': [], 'precision': [], 'recall': [], 'hd95': []
    }
    
    with torch.no_grad():
        for images, masks in tqdm(loader, desc="Evaluating"):
            images, masks = images.to(device), masks.to(device)
            vit_features = encoder(images)
            skip = stem(images)
            logits = decoder(vit_features, skip)

            metrics = compute_all_metrics(logits, masks)
            
            for key in all_metrics:
                all_metrics[key].append(metrics[key])
    
    # Compute mean and std for each metric
    results = {}
    for key in all_metrics:
        results[key] = np.mean(all_metrics[key])
        results[f'{key}_std'] = np.std(all_metrics[key])
    
    return results

# ============================================================================
# TRAINING FUNCTION
# ============================================================================

def train_model(decoder, stem, encoder, train_loader, val_loader, config):
    """Enhanced training loop with cosine annealing restarts and comprehensive logging"""
    device = config.device
    best_score = -float('inf')
    criterion = FocalDiceBoundaryLoss(focal_weight=config.focal_weight, dice_weight=config.dice_weight, boundary_weight=config.boundary_weight)
    
    # Optimizer includes both stem and decoder parameters
    optimizer = AdamW(
        list(decoder.parameters()) + list(stem.parameters()),
        lr=config.learning_rate,
        weight_decay=config.weight_decay
    )
    
    # Cosine Annealing with Warm Restarts
    scheduler = CosineAnnealingWarmRestarts(
        optimizer,
        T_0=config.T_0,
        T_mult=config.T_mult,
        eta_min=config.min_lr
    )

    
    history = {
        'train_loss': [],
        'val_metrics': [],  # Store full metrics dict per epoch
        'lr': []
    }

    for epoch in range(config.num_epochs):
        # Training
        decoder.train()
        stem.train()
        encoder.eval()

        epoch_loss = 0
        progress_bar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{config.num_epochs}")

        for batch_idx, (images, masks) in enumerate(progress_bar):
            images, masks = images.to(device), masks.to(device)

            # Frozen encoder
            with torch.no_grad():
                vit_features = encoder(images)

            # Trainable stem
            skip_features = stem(images)

            # Trainable decoder
            logits = decoder(vit_features, skip_features)
            loss = criterion(logits, masks)

            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(decoder.parameters(), max_norm=1.0)
            torch.nn.utils.clip_grad_norm_(stem.parameters(), max_norm=1.0)
            optimizer.step()

            # Step scheduler per batch for cosine annealing
            scheduler.step(epoch + batch_idx / len(train_loader))

            epoch_loss += loss.item()
            current_lr = optimizer.param_groups[0]['lr']
            progress_bar.set_postfix({'loss': loss.item(), 'lr': f'{current_lr:.2e}'})

        avg_loss = epoch_loss / len(train_loader)

        # Validation
        val_metrics = evaluate(decoder, stem, encoder, val_loader, device)

        # Store metrics
        history['train_loss'].append(avg_loss)
        history['val_metrics'].append(val_metrics)
        history['lr'].append(current_lr)

        # Save best model
        

        current_score = (0.6 * val_metrics['dice'] + 
                 0.3 * val_metrics['iou'] - 
                 0.1 * min(val_metrics['hd95'] / 100.0, 1.0))

        if current_score > best_score :  # Rename best_dice to best_score for clarity
            best_score  = current_score
            print(f"✓ Saved new best model with Dice: {val_metrics['dice']:.4f}, "
          f"IoU: {val_metrics['iou']:.4f}, HD95: {val_metrics['hd95']:.2f}")
            torch.save({
                'epoch': epoch,
                'decoder_state_dict': decoder.state_dict(),
                'stem_state_dict': stem.state_dict(),
                'encoder_state_dict': encoder.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'best_score': best_score,
                'config': config,
            }, os.path.join(config.save_dir, "best_unet_model.pth"))
            print(f"✓ Saved new best model with Score: {best_score:.4f}")

        # Print epoch summary
        print(f"\n{'='*60}")
        print(f"Epoch {epoch+1}/{config.num_epochs} Summary:")
        print(f"  Learning Rate: {current_lr:.6f}")
        print(f"  Train Loss: {avg_loss:.4f}")
        print(f"  Val Dice: {val_metrics['dice']:.4f} ± {val_metrics['dice_std']:.4f}")
        print(f"  Val IoU: {val_metrics['iou']:.4f} ± {val_metrics['iou_std']:.4f}")
        print(f"  Val Precision: {val_metrics['precision']:.4f} ± {val_metrics['precision_std']:.4f}")
        print(f"  Val Recall: {val_metrics['recall']:.4f} ± {val_metrics['recall_std']:.4f}")
        print(f"  Val HD95: {val_metrics['hd95']:.4f} ± {val_metrics['hd95_std']:.4f}")
        print(f"{'='*60}\n")

    return history,  best_score

# ============================================================================
# VISUALIZATION
# ============================================================================

def visualize_predictions(decoder, stem, encoder, dataset, device, num_samples=5, 
                         save_path="predictions.png", subset_name="Test"):
    """Visualize sample predictions with all metrics"""
    decoder.eval()
    stem.eval()
    encoder.eval()
    
    # Create a larger figure for 5 columns (image, mask, pred, overlay, metrics)
    fig, axes = plt.subplots(num_samples, 5, figsize=(20, 4*num_samples))
    
    if num_samples == 1:
        axes = axes.reshape(1, -1)
    
    indices = np.random.choice(len(dataset), num_samples, replace=False)
    
    with torch.no_grad():
        for i, idx in enumerate(indices):
            image, mask = dataset[idx]
            image_batch = image.unsqueeze(0).to(device)
            mask_np = mask.cpu().numpy().squeeze()
            
            vit_features = encoder(image_batch)
            skip = stem(image_batch)
            logits = decoder(vit_features, skip)
            pred = torch.sigmoid(logits).cpu().numpy().squeeze()
            pred_binary = (pred > 0.5).astype(np.float32)
            
            # Compute metrics
            metrics = compute_all_metrics(logits, mask.to(device))
            
            # Denormalize image for display
            img_display = image.cpu().squeeze().permute(1, 2, 0).numpy()
            mean = np.array(config.mean).reshape(1, 1, 3)
            std = np.array(config.std).reshape(1, 1, 3)
            img_display = img_display * std + mean
            img_display = np.clip(img_display, 0, 1)
            
            # Create overlay
            overlay = img_display.copy()
            overlay[pred_binary > 0.5] = [1, 0, 0]  # Red for predictions
            overlay = 0.7 * img_display + 0.3 * overlay
            
            # Plot images
            axes[i, 0].imshow(img_display)
            axes[i, 0].set_title("Input Image")
            axes[i, 0].axis('off')
            
            axes[i, 1].imshow(mask_np, cmap='gray')
            axes[i, 1].set_title("Ground Truth")
            axes[i, 1].axis('off')
            
            axes[i, 2].imshow(pred_binary, cmap='gray')
            axes[i, 2].set_title("Prediction")
            axes[i, 2].axis('off')
            
            axes[i, 3].imshow(overlay)
            axes[i, 3].set_title("Overlay")
            axes[i, 3].axis('off')
            
            # Display metrics in text
            metrics_text = f"Dice: {metrics['dice']:.3f}\nIoU: {metrics['iou']:.3f}\nHD95: {metrics['hd95']:.1f}"
            axes[i, 4].text(0.1, 0.5, metrics_text, fontsize=12, verticalalignment='center',
                           transform=axes[i, 4].transAxes)
            axes[i, 4].axis('off')
    
    plt.suptitle(f"{subset_name} Set - Sample Predictions", fontsize=16, y=1.02)
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"Visualization saved to {save_path}")

# ============================================================================
# MAIN PIPELINE
# ============================================================================

def load_and_prepare_data(config):
    """Load Kvasir-SEG dataset and create train/val/test splits"""
    
    images_path = os.path.join(config.dataset_path, "images")
    masks_path = os.path.join(config.dataset_path, "masks")
    
    if not os.path.exists(images_path):
        images_path = config.dataset_path
        masks_path = config.dataset_path
    
    image_files = sorted(glob.glob(os.path.join(images_path, "*.jpg")))
    mask_files = sorted(glob.glob(os.path.join(masks_path, "*.jpg")))
    
    if len(image_files) == 0:
        image_files = sorted(glob.glob(os.path.join(images_path, "*.png")))
        mask_files = sorted(glob.glob(os.path.join(masks_path, "*.png")))
    
    print(f"Found {len(image_files)} images and {len(mask_files)} masks")
    
    if len(image_files) == 0:
        raise FileNotFoundError(f"No images found in {config.dataset_path}")
    
    assert len(image_files) == len(mask_files), f"Mismatch: {len(image_files)} images vs {len(mask_files)} masks"
    
    # Split into train/val/test
    train_files, temp_files = train_test_split(
        list(zip(image_files, mask_files)), 
        test_size=config.val_split + config.test_split, 
        random_state=42
    )
    val_files, test_files = train_test_split(
        temp_files, 
        test_size=config.test_split / (config.val_split + config.test_split), 
        random_state=42
    )
    
    train_images, train_masks = zip(*train_files) if train_files else ([], [])
    val_images, val_masks = zip(*val_files) if val_files else ([], [])
    test_images, test_masks = zip(*test_files) if test_files else ([], [])
    
    print(f"Train: {len(train_images)}, Val: {len(val_images)}, Test: {len(test_images)}")
    
    return (list(train_images), list(train_masks)), (list(val_images), list(val_masks)), (list(test_images), list(test_masks))


def plot_training_history(history, save_dir):
    """Plot training history"""
    epochs = range(1, len(history['train_loss']) + 1)
    
    # Extract validation metrics
    val_dice = [m['dice'] for m in history['val_metrics']]
    val_iou = [m['iou'] for m in history['val_metrics']]
    val_hd95 = [m['hd95'] for m in history['val_metrics']]
    val_precision = [m['precision'] for m in history['val_metrics']]
    val_recall = [m['recall'] for m in history['val_metrics']]
    
    fig, axes = plt.subplots(2, 3, figsize=(18, 10))
    
    # Loss
    axes[0, 0].plot(epochs, history['train_loss'], 'b-', label='Train Loss')
    axes[0, 0].set_title('Training Loss')
    axes[0, 0].set_xlabel('Epoch')
    axes[0, 0].set_ylabel('Loss')
    axes[0, 0].grid(True)
    axes[0, 0].legend()
    
    # Learning Rate
    axes[0, 1].plot(epochs, history['lr'], 'g-')
    axes[0, 1].set_title('Learning Rate')
    axes[0, 1].set_xlabel('Epoch')
    axes[0, 1].set_ylabel('LR')
    axes[0, 1].set_yscale('log')
    axes[0, 1].grid(True)
    
    # Dice
    axes[0, 2].plot(epochs, val_dice, 'r-', label='Val Dice')
    axes[0, 2].set_title('Validation Dice')
    axes[0, 2].set_xlabel('Epoch')
    axes[0, 2].set_ylabel('Dice')
    axes[0, 2].grid(True)
    axes[0, 2].legend()
    
    # IoU
    axes[1, 0].plot(epochs, val_iou, 'm-', label='Val IoU')
    axes[1, 0].set_title('Validation IoU')
    axes[1, 0].set_xlabel('Epoch')
    axes[1, 0].set_ylabel('IoU')
    axes[1, 0].grid(True)
    axes[1, 0].legend()
    
    # HD95
    axes[1, 1].plot(epochs, val_hd95, 'c-', label='Val HD95')
    axes[1, 1].set_title('Validation HD95')
    axes[1, 1].set_xlabel('Epoch')
    axes[1, 1].set_ylabel('HD95 (pixels)')
    axes[1, 1].grid(True)
    axes[1, 1].legend()
    
    # Precision & Recall
    axes[1, 2].plot(epochs, val_precision, 'orange', label='Precision')
    axes[1, 2].plot(epochs, val_recall, 'purple', label='Recall')
    axes[1, 2].set_title('Validation Precision & Recall')
    axes[1, 2].set_xlabel('Epoch')
    axes[1, 2].set_ylabel('Value')
    axes[1, 2].grid(True)
    axes[1, 2].legend()
    
    plt.tight_layout()
    plt.savefig(os.path.join(save_dir, 'training_history.png'), dpi=150, bbox_inches='tight')
    plt.close()
    
    # Save history to CSV
    history_df = pd.DataFrame({
        'epoch': epochs,
        'train_loss': history['train_loss'],
        'val_dice': val_dice,
        'val_iou': val_iou,
        'val_hd95': val_hd95,
        'val_precision': val_precision,
        'val_recall': val_recall,
        'lr': history['lr']
    })
    history_df.to_csv(os.path.join(save_dir, 'training_history.csv'), index=False)


def main():
    print("=" * 60)
    print("DINOv3 Polyp Segmentation Training - With HD95 & Cosine Annealing")
    print("=" * 60)
    
    # Load data
    print("\n1. Loading dataset...")
    train_data, val_data, test_data = load_and_prepare_data(config)
    
    # Data augmentations
    train_transform = A.Compose([
        A.Resize(config.image_size, config.image_size),
        A.RandomRotate90(p=0.5),
        A.HorizontalFlip(p=0.5),
        A.VerticalFlip(p=0.5),
        A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.05, rotate_limit=15, p=0.5),
        A.OneOf([
            A.MotionBlur(p=0.2),
            A.GaussianBlur(blur_limit=3, p=0.2),
        ], p=0.3),
        A.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.05, p=0.3),
        ToTensorV2(),
    ])
    
    val_transform = A.Compose([
        A.Resize(config.image_size, config.image_size),
        ToTensorV2(),
    ])
    
    # Create datasets
    train_dataset = PolypDataset(
        train_data[0], train_data[1], 
        transform=train_transform,
        target_size=(config.image_size, config.image_size)
    )
    val_dataset = PolypDataset(
        val_data[0], val_data[1],
        transform=val_transform,
        target_size=(config.image_size, config.image_size)
    )
    test_dataset = PolypDataset(
        test_data[0], test_data[1],
        transform=val_transform,
        target_size=(config.image_size, config.image_size)
    )
    
    # Dataloaders
    train_loader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True, num_workers=4, pin_memory=True)
    val_loader = DataLoader(val_dataset, batch_size=config.batch_size, shuffle=False, num_workers=4, pin_memory=True)
    test_loader = DataLoader(test_dataset, batch_size=config.batch_size, shuffle=False, num_workers=4, pin_memory=True)
    
    print(f"\n2. Initializing DINOv3 encoder...")
    
    encoder = DINOv3Encoder(
        model_name=config.model_name, 
        local_path=config.local_model_path,
        freeze=True,
        layers=config.multi_scale_layers
    ).to(config.device)
    
    # Test encoder
    print("   Testing encoder with sample batch...")
    sample_images, _ = next(iter(train_loader))
    sample_images = sample_images.to(config.device)
    with torch.no_grad():
        sample_features = encoder(sample_images)
    print(f"   Encoder output shape: {sample_features.shape}")
    
    print("\n3. Building U‑Net decoder with skip connections...")
    
    stem = ShallowStem(in_channels=3, base_channels=64).to(config.device)
    decoder = UNetDecoder(
        vit_channels=encoder.out_channels,
        stem_channels=[512, 256, 128],
        num_classes=1
    ).to(config.device)
    
    trainable = sum(p.numel() for p in decoder.parameters()) + sum(p.numel() for p in stem.parameters())
    print(f"   Trainable parameters (stem + decoder): {trainable:,}")
    
    print("\n4. Starting training with Cosine Annealing Warm Restarts...")
    print(f"   Initial LR: {config.learning_rate:.6f}")
    print(f"   T_0: {config.T_0}, T_mult: {config.T_mult}")
    print(f"   Min LR: {config.min_lr:.6f}")
    
    history, best_score = train_model(decoder, stem, encoder, train_loader, val_loader, config)
    
    print(f"\n✓ Training complete! Best validation Score: {best_score:.4f}")
    
    # Final evaluation on all sets
    print("\n5. Final evaluation on all sets...")
    
    # Load best model for final evaluation
    checkpoint = torch.load(os.path.join(config.save_dir, "best_unet_model.pth"),weights_only=False)
    decoder.load_state_dict(checkpoint['decoder_state_dict'])
    stem.load_state_dict(checkpoint['stem_state_dict'])
    
    # Evaluate on all splits
    print("\nEvaluating on Training Set...")
    train_metrics = evaluate(decoder, stem, encoder, train_loader, config.device)
    
    print("Evaluating on Validation Set...")
    val_metrics = evaluate(decoder, stem, encoder, val_loader, config.device)
    
    print("Evaluating on Test Set...")
    test_metrics = evaluate(decoder, stem, encoder, test_loader, config.device)
    
    # Print comprehensive results
    print("\n" + "=" * 80)
    print("FINAL RESULTS - ALL METRICS")
    print("=" * 80)
    
    print(f"\n{'Metric':<15} {'Train':<20} {'Validation':<20} {'Test':<20}")
    print("-" * 75)
    
    for metric in ['dice', 'iou', 'precision', 'recall', 'hd95']:
        print(f"{metric.upper():<15} "
              f"{train_metrics[metric]:.4f} ± {train_metrics[f'{metric}_std']:.4f}    "
              f"{val_metrics[metric]:.4f} ± {val_metrics[f'{metric}_std']:.4f}    "
              f"{test_metrics[metric]:.4f} ± {test_metrics[f'{metric}_std']:.4f}")
    
    print("=" * 80)
    
    # Plot training history
    print("\n6. Plotting training history...")
    plot_training_history(history, config.save_dir)
    
    # Visualize predictions for all subsets
    print("\n7. Generating visualizations for all subsets...")
    visualize_predictions(decoder, stem, encoder, train_dataset, config.device, 
                         num_samples=5, save_path=os.path.join(config.save_dir, "train_predictions.png"),
                         subset_name="Training")
    visualize_predictions(decoder, stem, encoder, val_dataset, config.device,
                         num_samples=5, save_path=os.path.join(config.save_dir, "val_predictions.png"),
                         subset_name="Validation")
    visualize_predictions(decoder, stem, encoder, test_dataset, config.device,
                         num_samples=5, save_path=os.path.join(config.save_dir, "test_predictions.png"),
                         subset_name="Test")
    
    # Save comprehensive results
    results = {
        'best_val_score': float(best_score),
        'final_epoch': len(history['train_loss']),
        'train_metrics': {k: float(v) for k, v in train_metrics.items()},
        'val_metrics': {k: float(v) for k, v in val_metrics.items()},
        'test_metrics': {k: float(v) for k, v in test_metrics.items()},
        'training_history': {
            'train_loss': [float(x) for x in history['train_loss']],
            'lr': [float(x) for x in history['lr']],
            'val_metrics': [{k: float(v) for k, v in m.items()} for m in history['val_metrics']]
        },
        'config': {
            'model_name': config.model_name,
            'image_size': config.image_size,
            'batch_size': config.batch_size,
            'num_epochs': config.num_epochs,
            'learning_rate': config.learning_rate,
            'min_lr': config.min_lr,
            'T_0': config.T_0,
            'T_mult': config.T_mult,
            'scheduler': 'CosineAnnealingWarmRestarts',
            'focal_weight': config.focal_weight,
            'dice_weight': config.dice_weight,
            'multi_scale_layers': config.multi_scale_layers
        }
    }
    
    # Save as JSON
    with open(os.path.join(config.save_dir, "comprehensive_results.json"), 'w') as f:
        json.dump(results, f, indent=2)
    
    # Save as formatted text report
    with open(os.path.join(config.save_dir, "results_report.txt"), 'w') as f:
        f.write("=" * 80 + "\n")
        f.write("DINOv3 POLYP SEGMENTATION - FINAL REPORT\n")
        f.write("=" * 80 + "\n\n")
        f.write(f"Date: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")
        
        f.write("CONFIGURATION:\n")
        f.write("-" * 40 + "\n")
        for key, value in results['config'].items():
            f.write(f"  {key}: {value}\n")
        
        f.write("\n\nFINAL METRICS:\n")
        f.write("-" * 40 + "\n")
        f.write(f"{'Metric':<15} {'Train':<25} {'Validation':<25} {'Test':<25}\n")
        f.write("-" * 90 + "\n")
        
        for metric in ['dice', 'iou', 'precision', 'recall', 'hd95']:
            f.write(f"{metric.upper():<15} "
                   f"{train_metrics[metric]:.4f} ± {train_metrics[f'{metric}_std']:.4f}        "
                   f"{val_metrics[metric]:.4f} ± {val_metrics[f'{metric}_std']:.4f}        "
                   f"{test_metrics[metric]:.4f} ± {test_metrics[f'{metric}_std']:.4f}\n")
        
        f.write("\n\nBest Validation Score (Dice+IoU-HD95/100): {:.4f}\n".format(best_score))
        f.write("Training completed at epoch: {}\n".format(len(history['train_loss'])))
    
    print(f"\n✓ Comprehensive results saved to {config.save_dir}/")
    print(f"  - comprehensive_results.json")
    print(f"  - results_report.txt")
    print(f"  - training_history.csv")
    print(f"  - training_history.png")
    print(f"  - train_predictions.png")
    print(f"  - val_predictions.png")
    print(f"  - test_predictions.png")
    print("\n🎉 Enhanced training pipeline complete!")


if __name__ == "__main__":
    main()