!pip install gdown==4.7.1
!pip install patool==1.12
!pip install numpy==1.24.3
!pip install pandas==1.5.3
!pip install matplotlib==3.7.1
!pip install opencv-python==4.7.0.72
!pip install pillow==9.5.0
!pip install scikit-learn==1.2.2
!pip install tqdm==4.65.0
!pip install torchsummary==1.5.1
!pip install tensorboard
!pip install openpyxl

# Libraries for operating system and file handling
import os                         # Interacting with the file system (directories, paths)
import time                       # Time measurements (e.g., training duration)
import random                     # Random number generation (e.g., for reproducibility)
from datetime import datetime     # Working with dates and timestamps (e.g., for naming logs or saving files)
import gdown                      # For downloading files from Google Drive using their shareable URL
import patoolib                   # For extracting archive files (e.g., .rar, .zip, .tar) using external tools like unrar or 7z
import re                         # Regular expressions for string pattern matching (e.g., extracting labels from filenames)
import requests                   # For making HTTP requests (e.g., checking internet connection or downloading via direct links)
import warnings                   # Suppressing or handling warning messages (e.g., SSL warnings, deprecated features)
import shutil                     # High-level file operations (e.g., copying, moving, deleting directories or files)

# Libraries for numerical and data analysis
import numpy as np                # NumPy for efficient numerical computations
import pandas as pd               # Pandas for handling tabular data (e.g., labels, metadata)
from math import log10            # Used for computing base-10 logarithmic values (e.g., for PSNR calculation)
import math                       # For calculating number of rows/columns in the subplot grid

# Libraries for visualization and progress tracking
import matplotlib.pyplot as plt   # Matplotlib for plotting images and performance metrics
from tqdm import tqdm             # TQDM for progress bars in loops and data loading

# Libraries for image processing and manipulation
from PIL import Image, ImageOps   # Pillow: Image for loading images, ImageOps for basic image operations (e.g., grayscale, inversion, padding)
import cv2                        # OpenCV for image processing (e.g., resizing, filtering, transformations)

# Libraries for machine learning utilities
from sklearn.model_selection import train_test_split         # Splitting data into train/validation/test sets
from sklearn.metrics import (                                          # Evaluation metrics for classification tasks
    classification_report, confusion_matrix, ConfusionMatrixDisplay,
    roc_auc_score, roc_curve
)
from sklearn.preprocessing import label_binarize      # Converts multiclass labels into a binary matrix for one-vs-rest ROC AUC computation


# PyTorch libraries for building and training models
import torch                      # Core PyTorch library for tensor computation and GPU acceleration
import torch.nn as nn             # Tools for building neural network layers (e.g., Conv2D, Linear)
import torch.nn.functional as F   # Functional API for operations like activations, pooling, etc.
from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler    # Dataset and batch loading utilities
from torch.utils.tensorboard import SummaryWriter            # Logging training progress to TensorBoard
from torchsummary import summary  # Layer-wise summary of model architecture (input/output shapes, parameter count)

# Torchvision libraries for pretrained models and transforms
import torchvision                # High-level computer vision utilities and pretrained models
from torchvision import transforms, models                   # Transformations for image preprocessing and model zoo
from torchvision.models.feature_extraction import create_feature_extractor  # For extracting intermediate features from models

def check_device():
    """
    Checks if a CUDA-capable GPU is available, prints device info,
    and runs a small matrix multiplication test to confirm GPU functionality.
    Also prints the number of available CPU workers.
    """
    if torch.cuda.is_available():
        device = torch.device("cuda")
        print(f"✅ GPU available: {torch.cuda.get_device_name(0)}")
    else:
        device = torch.device("cpu")
        print("⚠️ No GPU found. Using CPU.")

    num_workers = os.cpu_count()
    print(f"🧠 Number of available CPU workers: {num_workers}")
    print(f"🖥️ Using device: {device}")

    # Simple test calculation
    try:
        a = torch.rand((1000, 1000), device=device)
        b = torch.rand((1000, 1000), device=device)
        c = torch.matmul(a, b)
        torch.cuda.synchronize() if device.type == 'cuda' else None
        print("✅ GPU test calculation completed successfully.")
    except Exception as e:
        print(f"❌ GPU test failed: {e}")

    return device

# Run the device check
device = check_device()

# Set seeds for reproducibility across torch, numpy, and random
def set_seed(seed=42):
    """
    Sets global seeds for reproducibility.
    Ensures consistent behavior across runs by fixing random number generators.
    """
    torch.manual_seed(seed)                        # Seed PyTorch
    np.random.seed(seed)                           # Seed NumPy
    random.seed(seed)                              # Seed Python's random
    torch.backends.cudnn.deterministic = True      # Ensure deterministic results on GPU (slower)
    torch.backends.cudnn.benchmark = False         # Disable cuDNN auto-tuner for consistency

    
# Apply the global seed
seed = 42
set_seed(seed)

warnings.filterwarnings("ignore")  # Disable SSL warnings

# List of file IDs
file_ids = [
    "1_CRKD0LyEfsYbgWI7NAZIGgIV6vw1edE",
    "1Tse1Tqr2IC-9eYKcUpGCp76DZUs27QeT",
    "1DV3LtLYAaP3eUoL2SAYiEYuPs1vqA3d6"
]

# Target folders
final_output_folder = "retina_dataset/images"
os.makedirs(final_output_folder, exist_ok=True)

for idx, file_id in enumerate(file_ids):
    url = f"https://drive.google.com/uc?export=download&id={file_id}"
    output_file = f"images{idx}.rar"
    
    # Download the file
    response = requests.get(url, stream=True, verify=False)
    with open(output_file, "wb") as f:
        for chunk in tqdm(response.iter_content(chunk_size=1024), desc=f"Downloading file {idx+1}"):
            if chunk:
                f.write(chunk)
    
    # Temporary extraction folder
    temp_extract_folder = f"temp_extract_{idx}"
    os.makedirs(temp_extract_folder, exist_ok=True)
    
    # Extract the .rar file
    patoolib.extract_archive(output_file, outdir=temp_extract_folder)
    
    # Move all files to final output folder
    for root, _, files in os.walk(temp_extract_folder):
        for file in files:
            src = os.path.join(root, file)
            dst = os.path.join(final_output_folder, file)
            shutil.move(src, dst)
    
    # Clean up
    shutil.rmtree(temp_extract_folder)
    os.remove(output_file)

print("All files downloaded, extracted, and moved to retina_dataset/images.")

# Download the Annotations file (.xlsx) format and save it in retina_dataset

excel_file_id = "1UkTRJN8X8d6sjBtFuxVxWOnDzkUr80Jf"
excel_url = f"https://docs.google.com/spreadsheets/d/{excel_file_id}/export?format=xlsx"
excel_output_path = os.path.join("retina_dataset", "Annotations.xls")

# Ensure the folder exists
os.makedirs("retina_dataset", exist_ok=True)

# Download and save the file
response = requests.get(excel_url, stream=True, verify=False)
with open(excel_output_path, "wb") as f:
    for chunk in tqdm(response.iter_content(chunk_size=1024), desc="Downloading Excel file"):
        if chunk:
            f.write(chunk)

print(f"Metadata file saved to {excel_output_path}")

# Define paths to the image directory and annotation file

data_dir         = "C:/Users/andre/Documents/Work/Projects/Tutorials/Retina_classification/retina_dataset/images/" 
annotations_path = "C:/Users/andre/Documents/Work/Projects/Tutorials/Retina_classification/retina_dataset/Annotations.xls" 
health_condition = "Retinopathy grade" #"Risk of macular edema "

# Load the Excel file
df = pd.read_excel(annotations_path)

# Select only the necessary columns
df = df[["Image name", health_condition]]

# Rename columns for consistency
df.columns = ["image_name", "label"]

# Create full image paths
df["image_path"] = df["image_name"].apply(lambda x: os.path.join(data_dir, x))

# Generate sorted list of class labels (e.g., 0, 1, 2, 3)
unique_labels = sorted(df["label"].unique())

# Create class names like "Grade 0", "Grade 1", ...
class_names = [f"Grade {label}" for label in unique_labels]

# Display the  DataFrame
df

# Plot Class Distribution

# Compute class distribution and percentages
class_counts = df['label'].value_counts().sort_index()
class_percentages = (class_counts / class_counts.sum()) * 100

# Plot bar chart
plt.figure(figsize=(6, 4))
ax = class_counts.plot(kind='bar', color='skyblue')

# Add title and labels
plt.title("Class Distribution of " + health_condition)
plt.xlabel(health_condition)
plt.ylabel("Number of Images")
plt.grid(False)
plt.tight_layout()

# Annotate bars with percentage values
for i, (count, pct) in enumerate(zip(class_counts, class_percentages)):
    ax.text(i, count + 1, f'{pct:.1f}%', ha='center', va='bottom', fontsize=9)

plt.show()

# Display Sample Images from Each Class
# Parameters
classes = sorted(df['label'].unique())
examples_per_class = 2
num_classes = len(classes)

# Set up the grid: rows = examples, columns = classes
fig, axes = plt.subplots(nrows=examples_per_class, ncols=num_classes, figsize=(num_classes * 4, examples_per_class * 3))
fig.suptitle("Example Images per Retinopathy Grade", fontsize=16, y=1.05)

print("--- Images Metadata ---")
# Plot examples_per_class random examples per class
for col, label in enumerate(classes):
    class_df = df[df['label'] == label]
    samples = class_df.sample(n=min(examples_per_class, len(class_df)), random_state=seed)

    for row in range(examples_per_class):
        ax = axes[row, col] if examples_per_class > 1 else axes[col]
        if row < len(samples):
            image_path = samples.iloc[row]['image_path']
            try:
                image = Image.open(image_path)
                img_array = np.array(image)

                # Extract image metadata
                width, height = image.size
                channels = 1 if len(img_array.shape) == 2 else img_array.shape[2]
                dtype = img_array.dtype

                if row == 0:
                    print(f"Image: {os.path.basename(image_path)}, grade:{label} {health_condition}, shape: {width}×{height}, # channels: C:{channels}, Image depth:{dtype}")

                # Show image
                ax.imshow(image)
                ax.axis('off')

                # Set column title in the top row
                if row == 0:
                    ax.set_title(f"Grade {label}", fontsize=12)

            except Exception as e:
                #ax.axis('off')
                ax.set_title("Error loading image")
                print(f"⚠️ Failed to load image {image_path}: {e}")
        else:
            ax.axis('off')
print("-----------------------")
plt.tight_layout()
plt.show()

# Image size and batch settings
image_size = (224, 224)
batch_size = 32
#num_workers = 8  # Adjust based on your hardware, suggeste half of teh avaible ones
imagenet_mean = [0.485, 0.456, 0.406]
imagenet_std = [0.229, 0.224, 0.225]

# Split proportions
test_size = 0.2       # 20% of the total dataset for testing
val_size = 0.2        # 20% of the training data for validation (i.e., 16% of total)

# First split: train + val and test
train_val_df, test_df = train_test_split(
    df,
    test_size=test_size,
    stratify=df['label'], # For unbalanced data to make the split even
    random_state=seed
)

# Second split: train and val from train_val_df
train_df, val_df = train_test_split(
    train_val_df,
    test_size=val_size,
    stratify=train_val_df['label'],  # For unbalanced data to make the split even
    random_state=seed
)


# Print Summary: Class Counts and Percentages

def print_distribution(name, subset):
    counts = subset['label'].value_counts().sort_index()
    percentages = subset['label'].value_counts(normalize=True).sort_index() * 100
    print(f"\n📊 {name} set:")
    for label in counts.index:
        print(f"  Class {label}: {counts[label]} images ({percentages[label]:.2f}%)")

print_distribution("Train", train_df)
print_distribution("Validation", val_df)
print_distribution("Test", test_df)

class RetinaDataset(Dataset):
    """
    Custom Dataset for RGB retinal image classification.
    Performs bounding box cropping, 
    and applies additional transformations.
    """
    def __init__(self, dataframe, data_dir, transform=None):
        """
        Parameters:
        - dataframe (pd.DataFrame): Must contain 'image_name' and 'label' columns.
        - data_dir (str): Path to image directory.
        - transform (callable, optional): Torchvision transforms to apply.
        """
        self.df = dataframe.reset_index(drop=True)
        self.data_dir = data_dir
        self.transform = transform

    def __len__(self):
        return len(self.df)

    
    def bounding_box_crop(self, img):
        """
        Crops the smallest bounding box around non-black pixels and pads it with zeros
        to create a square image with side length equal to the max side + 5 pixels.
    
        Parameters:
        img (PIL.Image): Input RGB image.
    
        Returns:
        PIL.Image: Cropped and padded RGB image.
        """
        # Convert PIL to OpenCV format (RGB to BGR)
        img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
    
        # Convert to grayscale
        gray = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY)
    
        # Create binary mask of non-black areas
        _, mask = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

        # Ensure mask has the same number of channels as the image
        mask_3ch = cv2.merge([mask, mask, mask])  # Convert single-channel to 3-channel mask
        
        # Apply mask to image (zero out background)
        img_cv = cv2.bitwise_and(img_cv, mask_3ch)

        # Get non-zero (non-black) coordinates
        coords = cv2.findNonZero(mask)
    
        if coords is None:
            print("⚠️ Image is completely black. Skipping crop.")
            return img
    
        # Bounding box
        x, y, w, h = cv2.boundingRect(coords)
    
        # Crop to bounding box
        cropped = img_cv[y:y+h, x:x+w]

        # Determine new square size
        max_side = max(w, h) 
        max_side += int(0.1 * max_side) 
    
        # Create black square canvas
        padded = np.zeros((max_side, max_side, 3), dtype=np.uint8) 

    
        # Compute top-left coordinates to center the cropped image
        y_offset = (max_side - h) // 2
        x_offset = (max_side - w) // 2
    
        # Paste the cropped image onto the center of the black square
        padded[y_offset:y_offset+h, x_offset:x_offset+w] = cropped
    
        # Convert back to PIL format
        result_img = Image.fromarray(cv2.cvtColor(padded, cv2.COLOR_BGR2RGB))
    
        return result_img  

    def preprocessing(self, img, apply_clahe=True):
        """
        Applies Ben Graham preprocessing with optional CLAHE to a fundus image.
    
        Parameters:
        - img: Input image (RGB as NumPy array or PIL Image)
        - apply_clahe: Whether to apply CLAHE to enhance local contrast
    
        Returns:
        - Preprocessed image as PIL Image (RGB)
        """
        # Ensure image is a NumPy array in RGB
        if isinstance(img, Image.Image):
            img = np.array(img)

        mask = img > 0 
        mask = 255 * mask.astype(np.uint8)
        
       # Convert RGB to BGR for OpenCV operations
        img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    
        # Apply CLAHE 
        if apply_clahe:
            lab = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2LAB)
            l, a, b = cv2.split(lab)
            clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
            l_clahe = clahe.apply(l)
            a_clahe = clahe.apply(a)
            b_clahe = clahe.apply(b)
            lab = cv2.merge((l_clahe, a_clahe, b_clahe))
            img_bgr = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
    
        # Normalize to [0, 1] and convert to RGB for PIL
        img_bgr = np.clip(img_bgr, 0, 255).astype(np.uint8)
        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

        # Apply mask to image (zero out background)
        img_rgb = cv2.bitwise_and(img_rgb, mask)

       
        return Image.fromarray(img_rgb)

    
    def __getitem__(self, idx):
        """
        Loads an RGB image, applies bounding box crop, and transforms.
        Returns:
        - image (Tensor or PIL.Image): Preprocessed image
        - label (int): Corresponding class label
        """
        try:
            image_name = self.df.iloc[idx]['image_name']
            label = int(self.df.iloc[idx]['label'])

            img_path = os.path.join(self.data_dir, image_name)
            image = Image.open(img_path).convert("RGB")  # ensure RGB

            image = self.bounding_box_crop(image)        # crop to content       
            image = self.preprocessing(image, apply_clahe=True) # Image preprocessing
    
            if self.transform:
                image = self.transform(image)

            return image, label

        except Exception as e:
            print(f"⚠️ Error loading image at index {idx}: {e}")
            return None, None

# Data augmentation + preprocessing for training
train_transform = transforms.Compose([
    transforms.Resize(image_size),                     # Resize image to fixed size
    transforms.RandomHorizontalFlip(p=0.5),            # Random horizontal flip: Increases symmetry: retinas are often symmetric
    #transforms.RandomVerticalFlip(p=0.5),             # Random vertical flip    ! No physiological 
    transforms.RandomRotation(degrees=20),              # Random rotation: Simulates head tilt, camera rotation
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2), #Color Jitter simulates variation across devices
    transforms.ToTensor(),                             # Convert PIL Image to tensor
    transforms.Normalize(mean=imagenet_mean,          # Normalize for pretrained models
                         std=imagenet_std)
])

# Preprocessing only for validation and testing
val_test_transform = transforms.Compose([
    transforms.Resize(image_size),
    transforms.ToTensor(),
    transforms.Normalize(mean=imagenet_mean,
                        std=imagenet_std)
])

def create_balanced_sampler(labels):
    """
    Creates a WeightedRandomSampler to handle class imbalance.

    Parameters:
    - labels (list or np.array): Class labels for each sample in the dataset.

    Returns:
    - sampler: PyTorch WeightedRandomSampler for use in DataLoader
    """
    # Convert to numpy array if not already
    labels = np.array(labels)

    # Count how many samples belong to each class
    class_counts = np.bincount(labels)
    #print(f"📊 Class counts: {dict(enumerate(class_counts))}")

    # Assign higher weight to underrepresented classes
    class_weights = np.where(class_counts == 0, 0, 1. / class_counts)

    # Weight per sample
    sample_weights = class_weights[labels]

    # Define the sampler
    sampler = WeightedRandomSampler(
        weights=sample_weights,
        num_samples=len(sample_weights),
        replacement=True
    )

    return sampler

# Create datasets
train_dataset = RetinaDataset(train_df, data_dir, transform=train_transform)
val_dataset = RetinaDataset(val_df, data_dir, transform=val_test_transform)
test_dataset = RetinaDataset(test_df, data_dir, transform=val_test_transform)

# Create balanced sampler for training set
train_labels = train_df['label'].values
train_sampler = create_balanced_sampler(train_labels)

# Create DataLoaders
train_loader = DataLoader(
    train_dataset,
    batch_size=batch_size,
    sampler=train_sampler,
)

val_loader = DataLoader(
    val_dataset,
    batch_size=batch_size,
    shuffle= False,
)

test_loader = DataLoader(
    test_dataset,
    batch_size=batch_size,
    shuffle= False,

)

def plot_batch(loader, title="Batch from DataLoader", max_images=8):
    """
    Displays a batch of images from a DataLoader.

    Parameters:
    - loader: PyTorch DataLoader
    - title: Title for the plot
    - max_images: Maximum number of images to show (default: 8)
    """
    imagenet_mean = [0.485, 0.456, 0.406]
    imagenet_std = [0.229, 0.224, 0.225]

    images, labels = next(iter(loader))
    num_images = min(len(images), max_images)

    fig, axs = plt.subplots(1, num_images, figsize=(num_images * 2.5, 3))
    fig.suptitle(title, fontsize=16)

    for idx in range(num_images):
        ax = axs[idx] if num_images > 1 else axs

        img = images[idx]
        label = labels[idx]

        # Ensure mean and std are tensors shaped for broadcasting
        imagenet_mean = torch.tensor(imagenet_mean).view(-1, 1, 1)
        imagenet_std = torch.tensor(imagenet_std).view(-1, 1, 1)
        
        # Unnormalize and clip between 0 and 1
        img = img * imagenet_std + imagenet_mean
        img = torch.clamp(img, 0, 1)


        # Handle grayscale (1 channel) and RGB (3 channels)
        if img.shape[0] == 1:
            ax.imshow(img[0], cmap='gray')
        else:
            ax.imshow(img.permute(1, 2, 0))  # Convert from (C, H, W) to (H, W, C)

        ax.set_title(f'Label: {label.item()}')
        ax.axis('off')

    plt.tight_layout()
    plt.show()


# Visualize Batches from All Datasets

plot_batch(train_loader, title="Train Set Sample")
plot_batch(val_loader, title="Validation Set Sample")
plot_batch(test_loader, title="Test Set Sample")

class SimpleCNN(nn.Module):
    def __init__(self, num_classes):
        super(SimpleCNN, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(32 * 56 * 56, 16),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(16, num_classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = self.classifier(x)
        return x


class SimpleCNNv2(nn.Module):
    def __init__(self, num_classes):
        super(SimpleCNNv2, self).__init__()
        
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=1),  # 3xHxW -> 32xHxW
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Conv2d(32, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(2),  # Downsample: H/2 x W/2

            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(2),  # Downsample again: H/4 x W/4

            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.MaxPool2d(2)  # Final downsampling: H/8 x W/8
        )

        self.global_pool = nn.AdaptiveAvgPool2d((1, 1))  # Output shape: [batch, 128, 1, 1]
        
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(128, 128),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(128, num_classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = self.global_pool(x)
        x = self.classifier(x)
        return x

def get_pretrained_model(name, num_classes, pretrained=True):
    if name == "resnet18":
        model = models.resnet18(pretrained=pretrained)
        model.fc = nn.Linear(model.fc.in_features, num_classes)
    elif name == "resnet50":
        model = models.resnet50(pretrained=pretrained)
        model.fc = nn.Linear(model.fc.in_features, num_classes)
    elif name == "resnet152":
        model = models.resnet152(pretrained=pretrained)
        model.fc = nn.Linear(model.fc.in_features, num_classes)       
    else:
        raise ValueError("Model not supported.")
    return model

# Define model dictionary

model_dict = {
    "01_SimpleCNN_scratch ": lambda: SimpleCNN(num_classes).to("cuda"),
    "02_SimpleCNNv2_scratch ": lambda: SimpleCNNv2(num_classes).to("cuda"),
    "03_ResNet18_scratch ": lambda: get_pretrained_model("resnet18", num_classes, pretrained=False).to("cuda"),
    "04_ResNet18_pretrained": lambda: get_pretrained_model("resnet18", num_classes, pretrained=True).to("cuda"),
    "05_ResNet50_scratch": lambda: get_pretrained_model("resnet50", num_classes, pretrained=False).to("cuda"),
    "06_ResNet50_pretrained": lambda: get_pretrained_model("resnet50", num_classes, pretrained=False).to("cuda"),
    "07_ResNet152_scratch": lambda: get_pretrained_model("resnet152", num_classes, pretrained=True).to("cuda"),
    "08_ResNet152_pretrained": lambda: get_pretrained_model("resnet152", num_classes, pretrained=True).to("cuda"),
}

'''
for model_name, model_fn in model_dict.items():
    print(f"\n🔍 Summary for: {model_name}")
    try:
        model = model_fn()
        summary(model, input_size=(3, *image_size))
    except Exception as e:
        print(f"⚠️ Could not generate summary: {e}")
     '''

class EarlyStopping:
    """
    Early stopping utility to stop training when validation loss doesn't improve.
    """
    def __init__(self, patience=5, min_delta=0.0, verbose=True, save_best_only=True):
        self.patience = patience
        self.min_delta = min_delta
        self.verbose = verbose
        self.counter = 0
        self.best_loss = float('inf')
        self.early_stop = False
        self.save_best_only = save_best_only
        self.best_model = None

    def __call__(self, val_loss: float, model: nn.Module):
        if val_loss < self.best_loss - self.min_delta:
            if self.verbose:
                print(f"Validation loss improved from {self.best_loss:.4f} to {val_loss:.4f}. Saving model...")
            self.best_loss = val_loss
            self.best_model = model.state_dict()
            self.counter = 0
        else:
            self.counter += 1
            if self.verbose:
                print(f"No improvement. EarlyStopping counter: {self.counter}/{self.patience}")
            if self.counter >= self.patience:
                self.early_stop = True

def train_model(model, train_loader, val_loader, num_epochs, model_name, learning_rate=1e-4):
    """
    Train a model with early stopping and TensorBoard logging.

    Parameters:
    - model: the PyTorch model to train
    - train_loader, val_loader: DataLoaders
    - num_epochs: number of max training epochs
    - model_name: name for logging/saving
    - learning_rate: optimizer learning rate
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)

    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    early_stopper = EarlyStopping(patience=5, min_delta=1e-4, verbose=True)

    writer = SummaryWriter(log_dir=f'runs/{model_name}_{int(time.time())}')

    for epoch in range(num_epochs):
        start_time = time.time()

        # Training loop
        model.train()
        train_loss = 0
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            train_loss += loss.item()

        # Validation loop
        model.eval()
        val_loss = 0
        with torch.no_grad():
            for images, labels in val_loader:
                images, labels = images.to(device), labels.to(device)
                outputs = model(images)
                loss = criterion(outputs, labels)
                val_loss += loss.item()

        avg_train_loss = train_loss / len(train_loader)
        avg_val_loss = val_loss / len(val_loader)
        epoch_time = time.time() - start_time

        # Logging
        writer.add_scalars('Loss', {
            'train': avg_train_loss,
            'val': avg_val_loss
        }, epoch)

        print(f"Epoch {epoch+1}/{num_epochs} - "
              f"Train Loss: {avg_train_loss:.4f}, "
              f"Val Loss: {avg_val_loss:.4f}, "
              f"Time: {epoch_time:.2f}s")

        # Check early stopping
        early_stopper(avg_val_loss, model)
        if early_stopper.early_stop:
            print("🛑 Early stopping triggered.")
            break

    # Restore best model weights
    model.load_state_dict(early_stopper.best_model)
    torch.save(model.state_dict(), f"{model_name}.pth")
    writer.close()

    return model

# Define base log directory
log_root = "runs_all_models"
os.makedirs(log_root, exist_ok=True)

# Update model dictionary to avoid overwriting logs
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

for model_name, model_fn in model_dict.items():
    print(f"\n🚀 Training model: {model_name}")
    try:
        model = model_fn()  # Already moved to CUDA
        log_name = f"{model_name.replace(' ', '_')}_{timestamp}"
        trained_model = train_model(
            model=model,
            train_loader=train_loader,
            val_loader=val_loader,
            num_epochs=100,
            model_name=os.path.join(log_root, log_name),
            learning_rate=1e-5
        )
    except Exception as e:
        print(f"⚠️ Failed to train {model_name}: {e}")

# In a Jupyter notebook cell
%reload_ext tensorboard
%tensorboard --logdir runs/runs_all_models --port 6007

def evaluate_model(model, model_name, dataloader, class_names):
    """
    Evaluate a multi-class model and generate reports, AUC-ROC curves, and confusion matrix.

    Args:
        model (nn.Module): Trained model.
        model_name (str): Name for display in plots and logs.
        dataloader (DataLoader): Dataloader for test or validation data.
        class_names (list[str]): List of class names.

    Returns:
        tuple: (confusion matrix, classification report as dict, auc_macro)
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    model.eval()

    y_true, y_pred, y_score = [], [], []

    with torch.no_grad():
        for images, labels in dataloader:
            images = images.to(device)
            outputs = model(images)
            probs = torch.softmax(outputs, dim=1).cpu().numpy()

            y_score.extend(probs)
            y_pred.extend(np.argmax(probs, axis=1))
            y_true.extend(labels.cpu().numpy())

    y_true = np.array(y_true)
    y_pred = np.array(y_pred)
    y_score = np.array(y_score)

    # Classification report
    print(f"\n📊 Classification Report for: {model_name}")
    print(classification_report(y_true, y_pred, target_names=class_names))

    # Confusion matrix
    cm = confusion_matrix(y_true, y_pred, normalize='true')
    disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=class_names)
    disp.plot(cmap='Blues', xticks_rotation=45, colorbar=False)
    plt.title(f"Confusion Matrix - {model_name}")
    plt.grid(False)
    plt.tight_layout()
    plt.show()

    # ROC-AUC (One-vs-Rest, macro average)
    y_true_bin = label_binarize(y_true, classes=range(len(class_names)))
    auc_macro = roc_auc_score(y_true_bin, y_score, multi_class='ovr', average='macro')
    print(f"\n📈 Macro AUC-ROC (OvR): {auc_macro:.4f}")

    # Plot ROC curves per class
    fpr = {}
    tpr = {}
    for i in range(len(class_names)):
        fpr[i], tpr[i], _ = roc_curve(y_true_bin[:, i], y_score[:, i])

    plt.figure(figsize=(8, 6))
    for i, class_name in enumerate(class_names):
        plt.plot(fpr[i], tpr[i], label=f'{class_name} (AUC = {roc_auc_score(y_true_bin[:, i], y_score[:, i]):.2f})')

    plt.plot([0, 1], [0, 1], 'k--')
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title(f'ROC Curves - {model_name}')
    plt.legend()
    plt.grid(True)
    plt.show()

    return cm, classification_report(y_true, y_pred, target_names=class_names, output_dict=True), auc_macro

evaluation_results = {}

# Scan all .pth files in log_root
for file in os.listdir(log_root):
    if not file.endswith(".pth"):
        continue

    model_path = os.path.join(log_root, file)
    filename = os.path.basename(model_path)
    
    # Regex captures number prefix + model name
    match = re.match(r"(\d+_.+?)_+(\d{8}_\d{6})\.pth$", filename)
    
    if match:
        model_key = match.group(1).strip()  # e.g., '01_ResNet18_pretrained'
    else:
        model_key = filename.replace(".pth", "").strip()

    # Find matching model in dictionary
    matching_keys = [key for key in model_dict.keys() if key.startswith(model_key)]
    if not matching_keys:
        print(f"⚠️ No matching model found in model_dict for: {model_key}")
        continue

    model_name = matching_keys[0]
    print(f"\n🔍 Loading and evaluating: {model_name} from  {file}")

    # Load model
    model = model_dict[model_name]()
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.to(device)

    # Evaluate model
    cm, report, auc_macro = evaluate_model(model, model_name, test_loader, class_names=class_names)
    evaluation_results[model_name] = {
        "confusion_matrix": cm,
        "report": report,
        "auc_macro": auc_macro,
        "model_path": model_path,
    }

# Find best model by AUC-ROC (macro average)
best_model = None
best_auc = -1.0

for model_name, metrics in evaluation_results.items():
    auc = metrics.get("auc_macro", 0)
    print(f"🔎 {model_name} - AUC-ROC (macro): {auc:.4f}")
    
    if auc > best_auc:
        best_auc = auc
        best_model = model_name

print("\n🏆 Best model based on AUC-ROC (macro):")
print(f"👉 {best_model} with AUC = {best_auc:.4f}")

def generate_cams(model, dataloader, target_layer='layer4', num_samples=10, class_names=None):
    """
    Generate and visualize CAMs (Class Activation Maps) for a given model and data loader.
    Displays them all in a grid of subplots.
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    model.eval()

    feature_extractor = create_feature_extractor(model, return_nodes={target_layer: "features"})

    try:
        if hasattr(model, 'fc'):
            softmax_weights = model.fc.weight.detach()
        elif hasattr(model, 'classifier'):
            if isinstance(model.classifier, nn.Sequential):
                softmax_weights = model.classifier[-1].weight.detach()
            else:
                softmax_weights = model.classifier.weight.detach()
        else:
            raise ValueError("Model must have a recognizable classifier layer (fc or classifier).")
    except Exception as e:
        raise ValueError(f"❌ Unable to extract classifier weights: {e}")
    
    overlays = []
    titles = []

    count = 0
    for images, labels in dataloader:
        images, labels = images.to(device), labels.to(device)

        with torch.no_grad():
            features = feature_extractor(images)["features"]
            outputs = model(images)
            preds = torch.argmax(outputs, dim=1)

        batch_size = images.size(0)

        for i in range(batch_size):
            if count >= num_samples:
                break

            feature_map = features[i]  # [C, H, W]
            pred_class = preds[i].item()
            label_class = labels[i].item()

            cam = torch.matmul(softmax_weights[pred_class], feature_map.view(feature_map.shape[0], -1))
            cam = cam.view(feature_map.shape[1], feature_map.shape[2]).cpu().numpy()
            cam = cv2.resize(cam, (images.shape[3], images.shape[2]))
            cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-5)
            cam = np.uint8(255 * cam)
            heatmap = cv2.applyColorMap(cam, cv2.COLORMAP_JET)

            img_np = images[i].detach().cpu().numpy().transpose(1, 2, 0)
            img_np = np.clip(img_np, 0, 1)
            img_np = np.uint8(255 * img_np)

            # Convert to grayscale and then back to 3-channel image for overlay
            gray_img = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
            img_np = cv2.cvtColor(gray_img, cv2.COLOR_GRAY2BGR)

            overlay = cv2.addWeighted(img_np, 0.7, heatmap, 0.3, 0)
            overlays.append(overlay[..., ::-1])  # Convert to RGB
            pred_label = class_names[pred_class] if class_names else f"Class {pred_class}"
            true_label = class_names[label_class] if class_names else f"Class {label_class}"
            titles.append(f"Pred: {pred_label}\nTrue: {true_label}")
            count += 1

        if count >= num_samples:
            break

    # 📊 Plot all images in a grid
    cols = min(5, num_samples)
    rows = math.ceil(num_samples / cols)
    fig, axes = plt.subplots(rows, cols, figsize=(4*cols, 4*rows))

    for idx, ax in enumerate(axes.flat):
        if idx < len(overlays):
            ax.imshow(overlays[idx])
            ax.set_title(titles[idx], fontsize=10)
            ax.axis('off')
        else:
            ax.axis('off')

    plt.tight_layout()
    plt.show()

# ✅ Get Top 3 Models Based on AUC-ROC (macro)
top_3_models = sorted(
    evaluation_results.items(),
    key=lambda x: x[1].get("auc_macro", 0),
    reverse=True
)[:3]

# 🔁 Loop over top 3 models only
for model_name, _ in top_3_models:
    print(f"\n🔍 Generating CAMs for: {model_name}")

    if model_name not in model_dict:
        print(f"⚠️ Model '{model_name}' not found in all_models dictionary.")
        continue

    # Retrieve model
    model = model_dict[model_name]()
    model_path = evaluation_results[model_name]["model_path"]
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.to(device)

    # Determine appropriate target layer
    
    if 'resnet' in model_name.lower():
        target_layer = 'layer4'
    else:  # For custom CNNs
        target_layer = 'features.3'

    # Generate CAMs
    generate_cams(model, test_loader, target_layer=target_layer, num_samples=5, class_names=class_names)