# Detect environment (Colab vs Local)
import sys
import os
import shutil
import stat

IN_COLAB = 'google.colab' in sys.modules

# GitHub repository configuration
GITHUB_USER = "RiaanZoetmulder"
GITHUB_REPO = "deep_learning_course"
GITHUB_BRANCH = "main"

if IN_COLAB:
    print("Running in Google Colab")
    WORKING_DIR = "/content"
else:
    print("Running locally")
    WORKING_DIR = os.getcwd()

# Temporary clone location
REPO_DIR = os.path.join(WORKING_DIR, GITHUB_REPO)

# Target directories in working directory
SESSION_DIR = os.path.join(WORKING_DIR, 'session_008')
MEDIA_DIR = os.path.join(WORKING_DIR, 'media')
HELPERS_DIR = os.path.join(WORKING_DIR, 'general_helpers')

def force_remove_readonly(func, path, excinfo):
    """Handle Windows read-only files during shutil.rmtree"""
    os.chmod(path, stat.S_IWRITE)
    func(path)

# Clone repo and extract needed folders
if not os.path.exists(SESSION_DIR) or not os.path.exists(HELPERS_DIR):
    # Clone if not already cloned
    if not os.path.exists(REPO_DIR):
        !git clone https://github.com/{GITHUB_USER}/{GITHUB_REPO}.git {REPO_DIR}
        print(f"✓ Cloned {GITHUB_REPO}")

    # Copy lecture_6/session_008 → session_008
    if not os.path.exists(SESSION_DIR):
        shutil.copytree(os.path.join(REPO_DIR, 'lecture_6', 'session_008'), SESSION_DIR)
        print(f"✓ Copied session_008")

    # Copy lecture_6/media → media
    if not os.path.exists(MEDIA_DIR):
        shutil.copytree(os.path.join(REPO_DIR, 'lecture_6', 'media'), MEDIA_DIR)
        print(f"✓ Copied media")

    # Copy general_helpers → general_helpers
    if not os.path.exists(HELPERS_DIR):
        shutil.copytree(os.path.join(REPO_DIR, 'general_helpers'), HELPERS_DIR)
        print(f"✓ Copied general_helpers")

    # Clean up cloned repo (with Windows read-only file handling)
    shutil.rmtree(REPO_DIR, onerror=force_remove_readonly)
    print(f"✓ Cleaned up temporary clone")
else:
    print("✓ Code folders already exist")

# Add module paths
sys.path.insert(0, SESSION_DIR)
sys.path.insert(0, HELPERS_DIR)

print(f"Module paths configured for {'Colab' if IN_COLAB else 'local'} environment")



MEDIA_FOLDER_PARTS = ["media"]  # OS-agnostic path components

# Set up local media path based on environment (OS-agnostic)
if IN_COLAB:
    MEDIA_ROOT = os.path.join("/content", *MEDIA_FOLDER_PARTS)
else:
    MEDIA_ROOT = os.path.join(".", *MEDIA_FOLDER_PARTS)

print('Set up Media Root folder at: ', MEDIA_ROOT)

import sys
import os
import subprocess

# Check if we're in a virtual environment
IN_VENV = sys.prefix != sys.base_prefix
VENV_PATH = os.path.join(os.getcwd(), '.venv')

if not IN_VENV and not IN_COLAB:
    print("=" * 70)
    print("⚠️  NOT RUNNING IN A VIRTUAL ENVIRONMENT")
    print("=" * 70)
    print("\nFollow these steps to set up your environment:\n")
    print("1. Open a terminal in VSCode (Ctrl+` or Terminal > New Terminal)")
    print("2. Run these commands:\n")
    print(f'   python -m venv .venv')
    print(f'   .venv\\Scripts\\activate')
    print(f'   pip install ipykernel --index-url https://pypi.org/simple/')
    print(f'   python -m ipykernel install --user --name=deep_learning_course --display-name="Deep Learning Course"')
    print("\n3. In VSCode, click 'Select Kernel' (top right of notebook)")
    print("4. Choose 'Python Environments...' > '.venv (Python x.x.x)'")
    print("   OR choose 'Jupyter Kernel...' > 'Deep Learning Course'")
    print("\n5. Re-run this cell to install packages (I've had to restart the kernel, the notebook, or even VSCode after setting up the venv for it to recognize the new environment)")
    print('NOTE: Installation of the required packages takes about 4 minutes on my machine. Please be patient and wait for the "Environment ready!" message before running the rest of the notebook \n or smashing your computer in frustration.')
    print("=" * 70)

elif IN_VENV or IN_COLAB:
    print(f"✓ Running in virtual environment: {sys.prefix}")
    print(f"✓ Python executable: {sys.executable}")
    print("\nInstalling required packages...")

    # Install packages using subprocess to ensure we use the right pip
    packages = [
        "ipykernel",
        "scikit-learn",
        "opencv-python",
        "pydicom",
        "ipywidgets",
        "numpy",
        "matplotlib"
    ]

    # Install base packages
    subprocess.run([
        sys.executable, "-m", "pip", "install", "-q",
        "--index-url", "https://pypi.org/simple/",
        *packages
    ], capture_output=True, text=True)
    print("✓ Base packages installed")

    # Install PyTorch (separate due to size)
    subprocess.run([
        sys.executable, "-m", "pip", "install", "-q",
        "--index-url", "https://pypi.org/simple/",
        "torch", "torchvision"
    ], capture_output=True, text=True)
    print("✓ PyTorch installed")

    # Register kernel if not already registered
    subprocess.run([
        sys.executable, "-m", "ipykernel", "install",
        "--user", "--name=deep_learning_course",
        "--display-name=Deep Learning Course"
    ], capture_output=True, text=True)
    print("✓ Jupyter kernel registered")

    print("\n" + "=" * 70)
    print("✓ Environment ready! You can now run the rest of the notebook.")
    print("=" * 70)

# NOTE: After uploading data.zip, run this when in colab
# Make sure you upload the data zip file to google colab first.
if IN_COLAB:
  !unzip data.zip -d /content/

# Core imports

# HINT: If you have a module not found error, you have to rerun the first code cell in this notebook
# It adds all the paths to sys, such that custom modules are found.
import sys
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from typing import List, Tuple
from torchvision import transforms
import pydicom as pdm

import numpy as np
import matplotlib.pyplot as plt

# custom imports
from ct_visualization import load_dicom_series, show_normalization_comparison, create_normalization_slider
from display_media import show_figure, show_video
from github_assets import download_media_assets
from feature_extraction import extract_vgg_features
from video_segmentation import display_video_segmentation

# Set device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")

show_figure("classification_detection_segmentation_comparison.png", width = 1500, media_root = MEDIA_ROOT)

mssg = """
---- NOTE ----
Bear in mind that in NumPy (which we are using to plot the images), the convention for coordinates is (row, column) which corresponds to (y, x) in Cartesian coordinates.\n This can lead to confusion when interpreting the axes of the images, especially for those familiar with the (x, y) convention in Cartesian space.\n
Moreover, the coordinates in the images are often displayed with the origin (0, 0) at the top-left corner. So be careful when intereting the bounding box label.
\n
The segmentation label example is derived from the real segmentation max, but downsampled used nearest neighbour interpolation. The real label has the same resolution as the original image.
"""
print(mssg)

show_figure("semantic_vs_instance_comparison.png", width = 1500, media_root = MEDIA_ROOT)

show_video("transposed_conv_2d.mp4", width = 900, media_root = MEDIA_ROOT)

# Example input tensor
input_tensor = torch.randn(1, 256, 8, 8)

print('Input tensor shape:', input_tensor.shape)

class TransposedConv2dExample(nn.Module):
    def __init__(self, in_channels: int=256, out_channels:int =128, kernel_size:int=3, stride:int=2, padding:int=1, output_padding:int=1):
        super(TransposedConv2dExample, self).__init__()
        """
        Single Layer Transpose Convolutional Network

        Bit superfluous to put it into an nn.Module. But we do want to reinforce the habit of using them :)
        """
        self.transposed_conv = nn.ConvTranspose2d(
            in_channels, out_channels, kernel_size, stride, padding, output_padding
        )

    def forward(self, x):
        return self.transposed_conv(x)

# Create an instance of the model and move it to the appropriate device
model = TransposedConv2dExample().to(device)

output_tensor = model(input_tensor.to(device))

print('Output tensor shape:', output_tensor.shape)

# the Input tensor, we are starting with a lot of channels, but very low height and width.
# This is a common scenario in the decoder part of a segmentation model,
# where we have a feature map with many channels (representing different features) but low spatial resolution (height and width).
input_tensor = torch.randn(1, 1024, 4, 4)

print('RECALL: The Dimensions are: (batch_size, channels, height, width)')
print('Input tensor shape:', input_tensor.shape)

class ExampleUpsamplingLayer(nn.Module):
    def __init__(
            self,
            in_channels: int=1024,
            out_channels:int=512,
            kernel_size:int=3,
            stride:int=2,
            padding:int=1,
            output_padding:int=1) -> None:
        super(ExampleUpsamplingLayer, self).__init__()
        """
        Example Upsampling Layer using Transposed Convolution

        Simple: transposed_convolution -> batch normalization -> ReLU
        """
        self.transposed_conv = nn.ConvTranspose2d(
            in_channels, out_channels, kernel_size, stride, padding, output_padding
        )

        self.relu = nn.ReLU()
        self.batch_norm = nn.BatchNorm2d(out_channels)  # Batch normalization

    def forward(self, x):
        x = self.transposed_conv(x)
        x = self.batch_norm(x)
        x = self.relu(x)
        return x


class ExampleUpsamplingModule(nn.Module):
    def __init__(self, in_channels: int=1024) -> None:
        super(ExampleUpsamplingModule, self).__init__()
        """
        Example Upsampling Module using Transposed Convolution

        We will create a simple upsampling module that takes an input tensor with 1024 channels and spatial dimension of 4.
        The module will use 6 layers of transposed convolutions to upsample the feature map to a spatial dimension of 256x256.

        The final layer will be a regular convolutional layer, which keeps the spatial dimensions the same but reduces the number of target channels.
        In case of binary segmentation (background vs foreground), this would just be the assignment of probability values to each pixels.
        """

        ## LAYER 1 - Upsampling from 1024 channels to 512 channels, and doubling the spatial dimensions to 8x8
        self.upsample1 = ExampleUpsamplingLayer(
            in_channels=in_channels,
            out_channels=512,
            kernel_size=3,
            stride=2,
            padding=1,
            output_padding=1
        )

        ## LAYER 2 - Upsampling from 512 channels to 256 channels, and doubling the spatial dimensions to 16x16
        self.upsample2 = ExampleUpsamplingLayer(
            in_channels=512,
            out_channels=256,
            kernel_size=3,
            stride=2,
            padding=1,
            output_padding=1
        )

        ## LAYER 3 - Upsampling from 256 channels to 128 channels, and doubling the spatial dimensions to 32x32
        self.upsample3 = ExampleUpsamplingLayer(
            in_channels=256,
            out_channels=128,
            kernel_size=3,
            stride=2,
            padding=1,
            output_padding=1
        )

        ## LAYER 4 - Upsampling from 128 channels to 64 channels, and doubling the spatial dimensions to 64x64
        self.upsample4 = ExampleUpsamplingLayer(
            in_channels=128,
            out_channels=64,
            kernel_size=3,
            stride=2,
            padding=1,
            output_padding=1
        )

        ## LAYER 5 - Upsampling from 64 channels to 32 channels, and doubling the spatial dimensions to 128x128
        self.upsample5 = ExampleUpsamplingLayer(
            in_channels=64,
            out_channels=32,
            kernel_size=3,
            stride=2,
            padding=1,
            output_padding=1
        )

        ## LAYER 6 - Upsampling from 32 channels to 16 channels, and doubling the spatial dimensions to 256x256
        self.upsample6 = ExampleUpsamplingLayer(
            in_channels=32,
            out_channels=16,
            kernel_size=3,
            stride=2,
            padding=1,
            output_padding=1
        )

        ## Out Layer - Final convolution to get the desired number of output channels (e.g., 1 for binary segmentation)
        # NOTE: Do not add an activation here, we would have to pass it to the loss function
        # (e.g., BCEWithLogitsLoss) which expects raw logits.
        self.out_layer = nn.Conv2d(16, 1, kernel_size=1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.upsample1(x)
        x = self.upsample2(x)
        x = self.upsample3(x)
        x = self.upsample4(x)
        x = self.upsample5(x)
        x = self.upsample6(x)
        x = self.out_layer(x)
        return x

# Create an instance of the upsampling module and move it to the appropriate device
upsampling_module = ExampleUpsamplingModule().to(device)

# Pass the input tensor through the upsampling module
output_tensor = upsampling_module(input_tensor.to(device))


print('\nOutput tensor shape:', output_tensor.shape)

"""
IMPORTANT NOTE

First we will show how to implement VGG! This is quite instructive
because it's been a while since we have implemented a convolutional neural network,
and it will be good to have a refresher on how to do that.
"""


class VGGConvolutionalBlock(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, blocks: int = 2) -> None:
        super(VGGConvolutionalBlock, self).__init__()
        """
        VGG Convolutional Block

        This block consists of two convolutional layers followed by ReLU activations.
        The number of input and output channels can be specified to match the architecture of VGG16.

        The number of blocks (convolutional layers) can be adjusted, but for VGG16, we typically
        have 2 or 3 convolutional layers per block.
        """

        layers = []
        for _ in range(blocks):
            conv_layer = nn.Conv2d(
                in_channels,
                out_channels,
                kernel_size=3,
                padding=1
            )
            layers.append(conv_layer)
            layers.append(nn.ReLU(inplace=True))

            in_channels = out_channels  # Update in_channels for the next block

        # add all of this to a sequential module for easier forward pass
        self.layers = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = self.layers(x)

        return out

# Illustrative Example of a VGG16-based Classifier, we will later modify this to be a segmentation model by adding a decoder part and
# removing the fully connected layers at the end.
class VGG16Example(nn.Module):
    def __init__(self, num_classes: int=1, input_size: int=224) -> None:
        super(VGG16Example, self).__init__()
        """
        VGG Classification model example, just so you know what it looks like!

        You do NOT have to train it, it is just here to show you how it is implemented.
        """

        # Layer 1 - Convolutional Block (input channels: 3 for RGB, output channels: 64)
        self.conv_block1 = VGGConvolutionalBlock(in_channels=3, out_channels=64, blocks = 2)
        self.max_pool_1 = nn.MaxPool2d(kernel_size=2, stride=2)

        # Layer 2 - Convolutional Block (input channels: 64, output channels: 128)
        self.conv_block2 = VGGConvolutionalBlock(in_channels=64, out_channels=128, blocks = 2)
        self.max_pool_2 = nn.MaxPool2d(kernel_size=2, stride=2)

        # Layer 3 - Convolutional Block (input channels: 128, output channels: 256)
        self.conv_block3 = VGGConvolutionalBlock(in_channels=128, out_channels=256, blocks = 3)
        self.max_pool_3 = nn.MaxPool2d(kernel_size=2, stride=2)

        # Layer 4 - Convolutional Block (input channels: 256, output channels: 512)
        self.conv_block4 = VGGConvolutionalBlock(in_channels=256, out_channels=512, blocks = 3)
        self.max_pool_4 = nn.MaxPool2d(kernel_size=2, stride=2)

        # Layer 5 - Convolutional Block (input channels: 512, output channels: 512)
        self.conv_block5 = VGGConvolutionalBlock(in_channels=512, out_channels=512, blocks = 3)
        self.max_pool_5 = nn.MaxPool2d(kernel_size=2, stride=2)

        # Fully Connected Layers (for classification, we will remove these for segmentation)
        # NOTE: Everything below this will be removed when we modify this to be a segmentation model.
        self.fc1 = nn.Linear(512 * (input_size // 32) * (input_size // 32), 4096)  # Adjust input size based on pooling
        self.fc2 = nn.Linear(4096, 4096)
        self.fc3 = nn.Linear(4096, num_classes)


    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]:
        skip_connections = []

        # Pass input through the encoder
        x = self.conv_block1(x)
        x = self.max_pool_1(x)

        x = self.conv_block2(x)
        x = self.max_pool_2(x)

        x = self.conv_block3(x)
        x = self.max_pool_3(x)
        skip_connections.append(x)

         # Save the output of the third block for skip connection in the decoder
        x = self.conv_block4(x)
        x = self.max_pool_4(x)
        skip_connections.append(x)

        x = self.conv_block5(x)
        x = self.max_pool_5(x)

        # though the fifth layer is not a skip connection, we will save it for later use during visualization of the feature maps.
        skip_connections.append(x)

        x = x.view(x.size(0), -1)  # Flatten for fully connected layers
        x = self.fc1(x)
        x = self.fc2(x)
        out = self.fc3(x)

        return out, skip_connections

# Parameters
num_classes = 20  # Example number of classes for classification
input_size = 224  # Example input size (height and width of the input image)
batch_size = 8   # Example batch size for testing
channels = 3
# Create an instance of the VGG16 example model and move it to the appropriate device
vgg16_example = VGG16Example(num_classes=num_classes, input_size=input_size).to(device)

# Example input tensor (batch size of 1, 3 channels for RGB, and spatial dimensions of 224x224)
input_tensor = torch.randn(batch_size, channels, input_size, input_size).to(device)

# Pass the input tensor through the VGG16 example model
output_tensor, skip_connections = vgg16_example(input_tensor)
print('Output tensor shape (classification logits):', output_tensor.shape)

## Load Pascal VOC Subsample for Transfer Learning Demo
import os
import xml.etree.ElementTree as ET
from PIL import Image
from torch.utils.data import Dataset

# VOC classes selected (similar to ImageNet classes)
VOC_CLASSES = ['dog', 'cat', 'bird', 'horse', 'cow', 'sheep', 'car', 'aeroplane', 'boat', 'train']
CLASS_TO_IDX = {cls: idx for idx, cls in enumerate(VOC_CLASSES)}

class VOCSubsampleDataset(Dataset):
    """
    Custom Dataset for loading the VOC subsample we created.
    Each image has a single object, and we use the XML annotations to get the class label.
    """
    def __init__(self, images_dir, annotations_dir, transform=None):
        self.images_dir = images_dir
        self.annotations_dir = annotations_dir
        self.transform = transform

        # Get all image files
        self.image_files = [f for f in os.listdir(images_dir) if f.endswith('.jpg')]

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

    def __getitem__(self, idx):
        # Load image
        img_name = self.image_files[idx]
        img_path = os.path.join(self.images_dir, img_name)
        image = Image.open(img_path).convert('RGB')

        # Load annotation to get class label
        xml_name = img_name.replace('.jpg', '.xml')
        xml_path = os.path.join(self.annotations_dir, xml_name)

        tree = ET.parse(xml_path)
        root = tree.getroot()

        # Get the class name from the first object (all our images have single objects)
        obj = root.find('object')
        class_name = obj.find('name').text
        label = CLASS_TO_IDX[class_name]

        # Apply transforms
        if self.transform:
            image = self.transform(image)

        return image, label


class VGG16Pretrained(nn.Module):
    def __init__(self, num_classes: int=1) -> None:
        super(VGG16Pretrained, self).__init__()
        """
        VGG16-based Classifier using Pretrained Weights

        This model uses the convolutional layers of VGG16 with pretrained weights on ImageNet.
        """

        # Load the pretrained VGG16 model
        vgg16 = torchvision.models.vgg16(pretrained=True)

        # Use the convolutional layers as the encoder
        self.encoder = vgg16.features
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Pass input through the encoder
        x = self.encoder(x)

        return x


# Preprocessing for VOC images:
# 1. CenterCrop to square (handles varying aspect ratios - crop using smaller dimension)
# 2. Resize to VGG16 input size (224x224)
# 3. Convert to tensor and normalize with ImageNet stats
#
# This ensures consistent preprocessing regardless of the original image resolution.
transform = transforms.Compose([
    transforms.Lambda(lambda img: transforms.CenterCrop(min(img.size))(img)),  # Crop to square using smaller dimension
    transforms.Resize((224, 224)),  # Resize to VGG16 input size
    transforms.ToTensor(),  # Convert to tensor
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # Normalize using ImageNet mean and std
])

# Set up paths for VOC subsample
VOC_SUBSAMPLE_DIR = os.path.join('data', 'VOC_subsample')
VOC_IMAGES_DIR = os.path.join(VOC_SUBSAMPLE_DIR, 'JPEGImages')
VOC_ANNOTATIONS_DIR = os.path.join(VOC_SUBSAMPLE_DIR, 'Annotations')

# Create dataset
voc_dataset = VOCSubsampleDataset(
    images_dir=VOC_IMAGES_DIR,
    annotations_dir=VOC_ANNOTATIONS_DIR,
    transform=transform
)

print(f"Loaded {len(voc_dataset)} images from VOC subsample")
print(f"Classes: {VOC_CLASSES}")

NUM_EXAMPLES = 400  # We have 400 images (40 per class × 10 classes)

# Create instances of the pretrained and untrained VGG16 models
vgg16_untrained = VGG16Example(num_classes=num_classes, input_size=input_size).to(device)
vgg16_pretrained = VGG16Pretrained().to(device)

# Extract features from both models
feature_maps_pretrained, feature_maps_untrained, labels = extract_vgg_features(
    dataset=voc_dataset,
    pretrained_model=vgg16_pretrained,
    untrained_model=vgg16_untrained,
    device=device,
    num_examples=NUM_EXAMPLES,
    print_interval=50,
    verbose=True
)

# Create two plots with t-SNE visualization (2D, side by side) of the feature maps
# from the pretrained and untrained models for comparison.
# We visualize the feature maps from the last convolutional block (before the FC layers).

from sklearn.manifold import TSNE
import matplotlib.pyplot as plt

# Flatten the feature maps from (1, 512, 7, 7) to 1D vectors
# Each feature map becomes a vector of size 512 * 7 * 7 = 25088
pretrained_flat = np.array([fm.flatten() for fm in feature_maps_pretrained])
untrained_flat = np.array([fm.flatten() for fm in feature_maps_untrained])

print(f"Pretrained feature shape: {pretrained_flat.shape}")
print(f"Untrained feature shape: {untrained_flat.shape}")

# Apply t-SNE with 2 components for 2D visualization
tsne = TSNE(n_components=2, random_state=42, perplexity=30, max_iter=2000)

print("Running t-SNE on pretrained features...")
pretrained_tsne = tsne.fit_transform(pretrained_flat)

print("Running t-SNE on untrained features...")
untrained_tsne = tsne.fit_transform(untrained_flat)

# Convert labels to numpy array
labels_array = np.array(labels[:len(pretrained_flat)])

# Create side-by-side 2D plots with extra space on the right for colorbar
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(18, 7))
fig.subplots_adjust(right=0.85)  # Make room for colorbar on the right

# Plot 1: Pretrained model features
scatter1 = ax1.scatter(pretrained_tsne[:, 0], pretrained_tsne[:, 1],
                       c=labels_array, cmap='tab10', s=20, alpha=0.7)
ax1.set_xlabel('t-SNE 1')
ax1.set_ylabel('t-SNE 2')
ax1.set_title('Pretrained VGG16 Features (ImageNet)')

# Plot 2: Untrained model features
scatter2 = ax2.scatter(untrained_tsne[:, 0], untrained_tsne[:, 1],
                       c=labels_array, cmap='tab10', s=20, alpha=0.7)
ax2.set_xlabel('t-SNE 1')
ax2.set_ylabel('t-SNE 2')
ax2.set_title('Untrained VGG16 Features (Random Weights)')

# Add colorbar in its own axes on the right side
cbar_ax = fig.add_axes([0.88, 0.15, 0.02, 0.7])  # [left, bottom, width, height]
cbar = fig.colorbar(scatter1, cax=cbar_ax)
cbar.set_ticks(range(10))
cbar.set_ticklabels(VOC_CLASSES)
cbar.set_label('Class')

plt.suptitle('t-SNE Visualization: Pretrained vs Untrained Feature Maps (Pascal VOC)', fontsize=14)
plt.show()

mssg = """
On the left you can see the Pre-trained VGG16 features. You can see that these are much more clustered and separated based on the class labels (color coded) compared to the untrained features on the right.
This is because the pretrained model has learned to extract meaningful features from the images that are relevant for classification, while the untrained model has random weights and therefore produces more random and less structured feature maps.\

It's not perfect (we are collapsing a very high dimensional space into 2 dimensions after all), but it's nifty for image data that the pre-trained model has never seen before!

"""

print(mssg)

import torch
import torch.nn as nn
from torchvision.models import vgg16, VGG16_Weights


class FCN8s_VGG16(nn.Module):
    """
    FCN-8s with VGG16 backbone for semantic segmentation.

    Architecture follows the original paper:
    "Fully Convolutional Networks for Semantic Segmentation" (Long et al., 2015)

    (1/N) refers to by how much it is downsampled.

    Key components:
    - VGG16 encoder (conv layers only, FC layers converted to conv)
    - Skip connections from pool3 (1/8) and pool4 (1/16)
    - FCN-8s combines predictions at 3 scales for fine-grained output
    It does so by adding in the predictions from the deeper layers (1/32) with the shallower layers (1/16 and 1/8) after upsampling.

    Args:
        n_classes: Number of output classes (21 for PASCAL VOC)
        pretrained_backbone: Whether to initialize VGG16 encoder with ImageNet weights.
            - Set to False (default) when loading full FCN weights from Caffe
            - Set to True only when training FCN from scratch (uses VGG16 as starting point)

    Usage:
        # For inference with pretrained FCN weights:
        model = FCN8s_VGG16(n_classes=21, pretrained_backbone=False)
        model = load_fcn8s_from_caffemodel(model, 'fcn8s-heavy-pascal.caffemodel')

        # For training from scratch (initialize encoder with ImageNet VGG16):
        model = FCN8s_VGG16(n_classes=21, pretrained_backbone=True)
    """

    def __init__(self, n_classes: int = 21, pretrained_backbone: bool = False):
        super().__init__()
        self.n_classes = n_classes

        # Load VGG16 architecture (weights will be replaced when loading Caffe model)
        weights = VGG16_Weights.IMAGENET1K_V1 if pretrained_backbone else None
        vgg = vgg16(weights=weights)
        features = list(vgg.features.children())

        # =================================================================
        # ENCODER: Split VGG16 into blocks for skip connections
        # =================================================================
        # This requires a little bit of explanation:
        # Basically what the  vgg16 class above does is it creates a sequential module which contains all the convolutional and pooling layers.
        # We know what the index is of vgg16 and we want to get intermediate outputs from the convolutional layers to use as our skip connections.
        # So we run the sequential model in 3 parts, obtaining the outputs at the end of each part for the skip connections.

        self.features_to_pool3 = nn.Sequential(*features[:17])  # Output: 1/8, 256 channels -> WILL BE A SKIP CONNECTION
        self.features_to_pool4 = nn.Sequential(*features[17:24]) # Output: 1/16, 512 channels -> WILL BE A SKIP CONNECTION
        self.features_to_pool5 = nn.Sequential(*features[24:])   # Output: 1/32, 512 channels

        # =================================================================
        # CLASSIFIER: Convert VGG FC layers to convolutional layers
        # =================================================================
        # Here we replace the fully connected layers of VGG16 with convolutional layers.

        self.fc6 = nn.Conv2d(512, 4096, kernel_size=7, padding=3)
        self.relu6 = nn.ReLU(inplace=True)
        self.drop6 = nn.Dropout2d()

        self.fc7 = nn.Conv2d(4096, 4096, kernel_size=1)
        self.relu7 = nn.ReLU(inplace=True)
        self.drop7 = nn.Dropout2d()

        # =================================================================
        # SCORING LAYERS: Project features to class scores
        # =================================================================
        self.score_fr = nn.Conv2d(4096, n_classes, kernel_size=1)    # From fc7 (1/32)
        self.score_pool4 = nn.Conv2d(512, n_classes, kernel_size=1)  # From pool4 (1/16)
        self.score_pool3 = nn.Conv2d(256, n_classes, kernel_size=1)  # From pool3 (1/8)

        # =================================================================
        # UPSAMPLING: Transposed convolutions for decoder
        # =================================================================
        # Upsample 2x: 1/32 -> 1/16 (to fuse with pool4)
        self.upscore2 = nn.ConvTranspose2d(
            n_classes, n_classes, kernel_size=4, stride=2, padding=1, bias=False
        )

        # Upsample 2x: 1/16 -> 1/8 (to fuse with pool3)
        self.upscore_pool4 = nn.ConvTranspose2d(
            n_classes, n_classes, kernel_size=4, stride=2, padding=1, bias=False
        )

        # Upsample 8x: 1/8 -> original size
        self.upscore8 = nn.ConvTranspose2d(
            n_classes, n_classes, kernel_size=16, stride=8, padding=4, bias=False
        )

        # Initialize weights (only used if not loading pretrained Caffe weights)
        self._initialize_weights()

    def _initialize_weights(self):
        """Initialize scoring and upsampling layers."""
        # Initialize scoring layers with zeros
        for layer in [self.score_fr, self.score_pool4, self.score_pool3]:
            nn.init.zeros_(layer.weight)
            nn.init.zeros_(layer.bias)

        # Initialize transposed convolutions with bilinear interpolation
        for layer in [self.upscore2, self.upscore_pool4, self.upscore8]:
            self._init_bilinear(layer)

    @staticmethod
    def _init_bilinear(conv: nn.ConvTranspose2d):
        """Initialize transposed conv with bilinear upsampling weights (diagonal/identity mapping)."""
        kernel_size = conv.kernel_size[0]
        factor = (kernel_size + 1) // 2
        center = factor - 1 if kernel_size % 2 == 1 else factor - 0.5

        og = torch.arange(kernel_size).float()
        filt = (1 - torch.abs(og - center) / factor)
        bilinear_kernel = filt.view(-1, 1) * filt.view(1, -1)

        # Initialize to zero, then set DIAGONAL only (channel i -> channel i)
        # This preserves class identity during upsampling
        conv.weight.data.zero_()
        for i in range(min(conv.in_channels, conv.out_channels)):
            conv.weight.data[i, i] = bilinear_kernel

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass with skip connections.

        Args:
            x: Input tensor of shape (N, 3, H, W)

        Returns:
            Output tensor of shape (N, n_classes, H, W)
        """
        input_size = x.shape[2:]  # Store for final crop if needed

        # Encoder path with skip connection outputs
        pool3 = self.features_to_pool3(x)        # 1/8, 256 channels
        pool4 = self.features_to_pool4(pool3)    # 1/16, 512 channels
        pool5 = self.features_to_pool5(pool4)    # 1/32, 512 channels

        # Classifier (converted FC layers) They are now convolutions with dropout.
        fc6 = self.drop6(self.relu6(self.fc6(pool5)))
        fc7 = self.drop7(self.relu7(self.fc7(fc6)))

        # Score from deepest layer (1/32)
        score_fr = self.score_fr(fc7)

        # First fusion: upsample 2x and add pool4 scores (1/16)
        upscore2 = self.upscore2(score_fr)
        score_pool4 = self.score_pool4(pool4)
        # Resize to match sizes (handles off-by-one from pooling/upsampling)
        upscore2 = nn.functional.interpolate(upscore2, size=score_pool4.shape[2:], mode='bilinear', align_corners=False)
        fuse_pool4 = upscore2 + score_pool4

        # Second fusion: upsample 2x and add pool3 scores (1/8)
        upscore_pool4 = self.upscore_pool4(fuse_pool4)
        score_pool3 = self.score_pool3(pool3)
        upscore_pool4 = nn.functional.interpolate(upscore_pool4, size=score_pool3.shape[2:], mode='bilinear', align_corners=False)
        fuse_pool3 = upscore_pool4 + score_pool3

        # Final upsampling 8x to original resolution
        upscore8 = self.upscore8(fuse_pool3)

        # Resize to input size (more robust than cropping)
        out = nn.functional.interpolate(upscore8, size=input_size, mode='bilinear', align_corners=False)

        return out

    @staticmethod
    def _crop_to_match(x: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        """Center crop x to match target's spatial dimensions."""
        _, _, h, w = target.shape
        _, _, xh, xw = x.shape
        dh, dw = (xh - h) // 2, (xw - w) // 2
        return x[:, :, dh:dh+h, dw:dw+w]

    @staticmethod
    def _crop_to_size(x: torch.Tensor, size: tuple) -> torch.Tensor:
        """Center crop x to specified size."""
        h, w = size
        _, _, xh, xw = x.shape
        dh, dw = (xh - h) // 2, (xw - w) // 2
        return x[:, :, dh:dh+h, dw:dw+w]


# Test the model (no pretrained backbone - will load Caffe weights separately)
model = FCN8s_VGG16(n_classes=21, pretrained_backbone=False)
print(f"FCN-8s VGG16 model created (random initialization)")
print(f"  - Number of parameters: {sum(p.numel() for p in model.parameters()):,}")

# Test forward pass
dummy_input = torch.randn(1, 3, 224, 224)
with torch.no_grad():
    output = model(dummy_input)
print(f"  - Input shape: {dummy_input.shape}")
print(f"  - Output shape: {output.shape}")
print(f"  - Output classes: {output.shape[1]} (PASCAL VOC)")

# Load pretrained FCN-8s weights and run inference
model = FCN8s_VGG16(n_classes=21, pretrained_backbone=False)
model.load_state_dict(torch.load('data/fcn8s_vgg16_pascal.pth', weights_only=True))
model.to(device)
model.eval()
print(f"✓ FCN-8s loaded from data/fcn8s_vgg16_pascal.pth")

VIDEO_INPUT = 'data/video_of_me.mov'

VOC_COLORMAP = np.array([
    (0, 0, 0),        # Background (black)
    (128, 0, 0),      # Aeroplane (dark red)
    (0, 128, 0),      # Bicycle (dark green)
    (128, 128, 0),    # Bird (olive)
    (0, 0, 128),      # Boat (dark blue)
    (128, 0, 128),    # Bottle (purple)
    (0, 128, 128),    # Bus (teal)
    (128, 128, 128),  # Car (gray)
    (64, 0, 0),       # Cat (maroon)
    (192, 0, 0),      # Chair (red)
    (64, 128, 0),     # Cow (dark olive)
    (192, 128, 0),    # Dining table (orange)
    (64, 0, 128),     # Dog (indigo)
    (192, 0, 128),    # Horse (magenta)
    (64, 128, 128),   # Motorbike (cyan)
    (192, 128, 128),   # Person (light pink)
    (0, 64, 0),       # Potted plant (dark green)
    (128, 64, 0),     # Sheep (brown)
    (0, 192, 0),      # Sofa (green)
    (128, 192, 0)     # Train (light green)
], np.uint8)

display_video_segmentation(
    video_path=VIDEO_INPUT,
    model=model,
    colormap=VOC_COLORMAP,
    device=device,
    target_fps=10.0,
    alpha=0.5,
    figsize=(15, 5),
    verbose=True
)

class InputBlock(nn.Module):
    def __init__(self, input_channels:int = 3, output_channels: int = 64, num_convolutions: int = 2, normalization: nn.Module = nn.BatchNorm2d, activation: nn.Module = nn.ReLU) -> None:
        """
        Input block to the U-Net architecture

        Args:
            input_channels (int, optional): Number of input channels. Defaults to 3.
            output_channels (int, optional): Number of output channels. Defaults to 64.
            num_convolutions (int, optional): Number of convolutional layers. Defaults to 2.
            normalization (nn.Module, optional): Normalization layer. Defaults to nn.BatchNorm2d.
            activation (nn.Module, optional): Activation function. Defaults to nn.ReLU.
        """
        super(InputBlock, self).__init__()

        layers = []
        #### IMPLEMENT HERE ####


        # save as module list for easier forward pass, we want to extract the last output before maxpooling for the skip connection (see figure above)
        self.block = nn.ModuleList(layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:

        skip_connection = None

        #### IMPLEMENT HERE ####

        return x, skip_connection


class DownsamplingBlock(nn.Module):
    def __init__(self, input_channels: int, output_channels: int, num_convolutions: int = 2, normalization: nn.Module = nn.BatchNorm2d, activation: nn.Module = nn.ReLU) -> None:
        """
        Single Block in the Encoder (Downsampling Path)


        Args:
            input_channels (int): Number of input channels.
            output_channels (int): Number of output channels.
            num_convolutions (int, optional): Number of convolutional layers. Defaults to 2.
            normalization (nn.Module, optional): Normalization layer. Defaults to nn.BatchNorm2d.
            activation (nn.Module, optional): Activation function. Defaults to nn.ReLU.
        """
        super(DownsamplingBlock, self).__init__()
        layers = []

        #### IMPLEMENT HERE ####


        self.block = nn.ModuleList(layers)


    def forward(self, x: torch.Tensor) -> torch.Tensor:

        skip_connection = None
        #### IMPLEMENT HERE ####


        return x, skip_connection


class BottleneckBlock(nn.Module):
    def __init__(self, input_channels: int, output_channels: int, num_convolutions: int = 2, normalization: nn.Module = nn.BatchNorm2d, activation: nn.Module = nn.ReLU) -> None:
        """
        Bottleneck block in the U-Net architecture
        Args:
            input_channels (int): Number of input channels.
            output_channels (int): Number of output channels.
            num_convolutions (int, optional): Number of convolutional layers. Defaults to 2.
            normalization (nn.Module, optional): Normalization layer. Defaults to nn.BatchNorm2d.
            activation (nn.Module, optional): Activation function. Defaults to nn.ReLU.
        """
        super(BottleneckBlock, self).__init__()
        layers = []
        #### IMPLEMENT HERE ####

        self.block = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = self.block(x)

        return out


class UpsamplingBlock(nn.Module):
    def __init__(self, input_channels: int, output_channels: int, num_convolutions: int = 2, normalization: nn.Module = nn.BatchNorm2d, activation: nn.Module = nn.ReLU) -> None:
        """
        Single Block in the Decoder (Upsampling Path)
        Args:
            input_channels (int): Number of input channels.
            output_channels (int): Number of output channels.
            num_convolutions (int, optional): Number of convolutional layers. Defaults to 2.
            normalization (nn.Module, optional): Normalization layer. Defaults to nn.BatchNorm2d.
            activation (nn.Module, optional): Activation function. Defaults to nn.ReLU. """

        super(UpsamplingBlock, self).__init__()
        layers = []
        #### IMPLEMENT HERE ####


        self.block = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor, skip_connection: torch.Tensor) -> torch.Tensor:

        # Concatenate the skip connection from the encoder with the output of the previous layer in the decoder
        # they are shaped (batch_size, channels, height, width) so we concatenate along the channel dimension (dim=1)
        #### IMPLEMENT HERE ####

        return out


class OutputBlock(nn.Module):
    def __init__(self, input_channels: int, num_classes: int, num_convolutions: int = 2, normalization: nn.Module = nn.BatchNorm2d, activation: nn.Module = nn.ReLU) -> None:
        """
        Output block of the U-Net architecture

        Args:
            input_channels (int): Number of input channels.
            num_classes (int): number of classes for segmentation.
            num_convolutions (int, optional): Number of convolutional layers. Defaults to 2.
            normalization (nn.Module, optional): Normalization layer. Defaults to nn.BatchNorm2d.
            activation (nn.Module, optional): Activation function. Not used for output layer. Defaults to nn.ReLU.
        """
        super(OutputBlock, self).__init__()
        layers = []
        #### IMPLEMENT HERE ####

        self.block = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor, skip_connection: torch.Tensor) -> torch.Tensor:

        #### IMPLEMENT HERE ####

        return out

class DownsamplingPath(nn.Module):
    def __init__(self, input_channels: int, base_channels: int = 64, num_blocks: int = 2, normalization: nn.Module = nn.BatchNorm2d, activation: nn.Module = nn.ReLU) -> None:
        """
        Downsampling path of the U-Net architecture (Encoder)
        Args:
            input_channels (int): Number of input channels.
            base_channels (int, optional): Number of output channels for the first block. Defaults to 64.
            num_blocks (int, optional): Number of downsampling blocks. Defaults to 4.
            num_convolutions (int, optional): Number of convolutional layers in each block. Defaults to 2.
            normalization (nn.Module, optional): Normalization layer. Defaults to nn.BatchNorm2d.
            activation (nn.Module, optional): Activation function. Defaults to nn.ReLU.
        """
        super(DownsamplingPath, self).__init__()

        print('INITIALIZING DOWNSAMPLING PATH (ENCODER)')

        blocks = []
        #### IMPLEMENT HERE ####


        print('---- ENCODER INITIALIZATION COMPLETE ----\n')

        self.blocks = nn.ModuleList(blocks)

    def forward(self, x: torch.Tensor) -> list:
        skip_connections = []
        #### IMPLEMENT HERE ####

        return x, skip_connections


class UpsamplingPath(nn.Module):
    def __init__(self, base_channels: int = 64, num_blocks: int = 2, num_classes: int = 1, normalization: nn.Module = nn.BatchNorm2d, activation: nn.Module = nn.ReLU) -> None:
        """
        Upsampling path of the U-Net architecture (Decoder)
        Args:
            base_channels (int, optional): Number of output channels for the first block. Defaults to 64.
            num_blocks (int, optional): Number of upsampling blocks. Defaults to 4.
            num_classes (int): Number of output classes for segmentation.
            normalization (nn.Module, optional): Normalization layer.Defaults to nn.BatchNorm2d
            activation (nn.Module, optional): Activation function. Defaults to nn.ReLU.
        """
        super(UpsamplingPath, self).__init__()

        print('INITIALIZING UPSAMPLING PATH (DECODER)')
        blocks = []
        #### IMPLEMENT HERE ####


        print('---- DECODER INITIALIZATION COMPLETE ----\n')
        self.blocks = nn.ModuleList(blocks)


    def forward(self, x: torch.Tensor, skip_connections: list) -> torch.Tensor:

        #### IMPLEMENT HERE ####

        return x

from typing import Any


class BasicUNet(nn.Module):
    def __init__(self, in_channels: int = 3, num_classes: int = 1, base_channels: int = 64, num_blocks: int = 2, activation: nn.Module = nn.ReLU, normalization: nn.BatchNorm2d = nn.BatchNorm2d) -> None:
        """
        A very basic version of the U-Net architecture.

        Args:
            in_channels (int, optional): Number of input channels. Defaults to 3.
            num_classes (int, optional): Number of output classes for segmentation. Defaults to 1.
            base_channels (int, optional): Number of output channels for the first block. Defaults to 64.
            num_blocks (int, optional): Number of blocks in the encoder and decoder. Defaults to
            activation (nn.Module, optional): Activation function. Defaults to nn.ReLU.
            normalization (nn.Module, optional): Normalization layer. Defaults to nn.BatchNorm2d
        """
        super(BasicUNet, self).__init__()
        #### IMPLEMENT HERE ####

        # Apply Kaiming (He) initialization for ReLU activations
        self._initialize_weights()

    def _initialize_weights(self):
        """
        Initialize weights using Kaiming (He) initialization.

        For ReLU activations, Kaiming initialization maintains proper variance
        of activations throughout the network by accounting for the fact that
        ReLU zeroes out approximately half of its inputs.

        Reference: He et al., "Delving Deep into Rectifiers" (2015)
        """
        for m in self.modules():
            if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
                # Kaiming normal initialization for conv layers with ReLU
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.BatchNorm2d):
                # Standard initialization for batch norm
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:

        #### IMPLEMENT HERE ####

        return x

# Download ISBI 2012 EM Segmentation Challenge dataset
import os
import zipfile
import urllib.request
from pathlib import Path

def download_isbi2012(data_dir="data/isbi2012"):
    """
    Download the ISBI 2012 EM Segmentation Challenge dataset.

    The dataset contains:
    - train-volume.tif: 30 training images (512x512 each)
    - train-labels.tif: 30 corresponding label images
    """
    data_path = Path(data_dir)
    data_path.mkdir(parents=True, exist_ok=True)

    # Direct download links from the challenge website
    urls = {
        "train-volume.tif": "https://github.com/zhixuhao/unet/raw/master/data/membrane/train-volume.tif",
        "train-labels.tif": "https://github.com/zhixuhao/unet/raw/master/data/membrane/train-labels.tif",
        "test-volume.tif": "https://github.com/zhixuhao/unet/raw/master/data/membrane/test-volume.tif"
    }

    for filename, url in urls.items():
        filepath = data_path / filename
        if not filepath.exists():
            print(f"Downloading {filename}...")
            urllib.request.urlretrieve(url, filepath)
            print(f"  Saved to {filepath}")
        else:
            print(f"{filename} already exists")

    return data_path

# Download the dataset
isbi_path = download_isbi2012()
print(f"\nDataset downloaded to: {isbi_path}")

# Visualize ISBI 2012 samples
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path

# Load the TIFF stacks
isbi_path = Path("data/isbi2012")
train_volume = Image.open(isbi_path / "train-volume.tif")
train_labels = Image.open(isbi_path / "train-labels.tif")

# Display first 4 image-mask pairs
fig, axes = plt.subplots(2, 4, figsize=(14, 7))
fig.suptitle("ISBI 2012 EM Dataset: Cell Membrane Segmentation", fontsize=14)

for i in range(4):
    # Navigate to the i-th frame
    train_volume.seek(i)
    train_labels.seek(i)

    # Convert to numpy
    img = np.array(train_volume)
    mask = np.array(train_labels)

    # Display image
    axes[0, i].imshow(img, cmap='gray')
    axes[0, i].set_title(f"Image {i+1}")
    axes[0, i].axis('off')

    # Display mask (invert so membranes are white)
    axes[1, i].imshow(mask, cmap='gray')
    axes[1, i].set_title(f"Mask {i+1}")
    axes[1, i].axis('off')

plt.tight_layout()
plt.savefig("media/isbi2012_samples.png", dpi=150, bbox_inches='tight')
plt.show()

# Print dataset info
print(f"\nDataset info:")
print(f"  Image size: {img.shape}")
print(f"  Number of training images: 30")
print(f"  Mask values: {np.unique(mask)}")  # Should be 0 and 255

# ISBI 2012 Dataset class and DataLoaders
import torch
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import numpy as np
from pathlib import Path
import torchvision.transforms.functional as TF
import random

class ISBI2012Dataset(Dataset):
    """
    PyTorch Dataset for ISBI 2012 EM Segmentation Challenge.

    The dataset is stored as multi-page TIFF files where each page is a 512x512 image.
    Supports both full resolution (512x512) and downsampled versions for faster training.
    """
    def __init__(self, data_dir="data/isbi2012", split="train", transform=True,
                 val_indices=None, train_indices=None, resolution=512):
        """
        Args:
            data_dir: Path to the ISBI2012 dataset
            split: "train" or "val"
            transform: Whether to apply data augmentation (only for training)
            val_indices: List of indices to use for validation (default: last 5)
            train_indices: List of indices to use for training (default: first 25)
            resolution: Target resolution (512 for full, 128/256 for downsampled)
        """
        self.data_dir = Path(data_dir)
        self.split = split
        self.transform = transform and (split == "train")
        self.resolution = resolution

        # Load the TIFF stacks
        self.volume = Image.open(self.data_dir / "train-volume.tif")
        self.labels = Image.open(self.data_dir / "train-labels.tif")

        # Get number of images in stack
        n_images = 30  # Known for ISBI 2012

        # Define train/val split (25 train, 5 val by default)
        if val_indices is None:
            val_indices = list(range(25, 30))  # Last 5 for validation
        if train_indices is None:
            train_indices = list(range(25))    # First 25 for training

        self.indices = train_indices if split == "train" else val_indices

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

    def __getitem__(self, idx):
        # Get the actual index in the TIFF stack
        tiff_idx = self.indices[idx]

        # Navigate to the correct frame
        self.volume.seek(tiff_idx)
        self.labels.seek(tiff_idx)

        # Convert to PIL Images for resizing
        image_pil = self.volume.copy()
        mask_pil = self.labels.copy()

        # Downsample if needed
        if self.resolution != 512:
            image_pil = image_pil.resize((self.resolution, self.resolution), Image.BILINEAR)
            mask_pil = mask_pil.resize((self.resolution, self.resolution), Image.NEAREST)

        # Convert to numpy arrays
        image = np.array(image_pil, dtype=np.float32)
        mask = np.array(mask_pil, dtype=np.float32)

        # Normalize image to [0, 1]
        image = image / 255.0

        # Convert mask to binary (0 or 1) - membranes are originally 0 (black)
        # We want membranes to be 1 (positive class)
        mask = (mask < 128).astype(np.float32)  # Invert: black (0) -> 1

        # Convert to tensors
        image = torch.from_numpy(image).unsqueeze(0)  # Add channel dim: [1, H, W]
        mask = torch.from_numpy(mask).unsqueeze(0)    # Add channel dim: [1, H, W]

        # Apply data augmentation
        if self.transform:
            image, mask = self._augment(image, mask)

        return image, mask

    def _augment(self, image, mask):
        """Apply random augmentations to image and mask."""
        # Random horizontal flip
        if random.random() > 0.5:
            image = TF.hflip(image)
            mask = TF.hflip(mask)

        # Random vertical flip
        if random.random() > 0.5:
            image = TF.vflip(image)
            mask = TF.vflip(mask)

        # Random rotation (0, 90, 180, or 270 degrees)
        angle = random.choice([0, 90, 180, 270])
        if angle != 0:
            image = TF.rotate(image, angle)
            mask = TF.rotate(mask, angle)

        return image, mask


# =============================================================================
# Create DOWNSAMPLED datasets (128x128) - for fast initial training
# =============================================================================
DOWNSAMPLED_RES = 128

train_dataset_low = ISBI2012Dataset(split="train", transform=True, resolution=DOWNSAMPLED_RES)
val_dataset_low = ISBI2012Dataset(split="val", transform=False, resolution=DOWNSAMPLED_RES)

print(f"=== Downsampled Dataset ({DOWNSAMPLED_RES}x{DOWNSAMPLED_RES}) ===")
print(f"Training samples: {len(train_dataset_low)}")
print(f"Validation samples: {len(val_dataset_low)}")

train_loader_low = DataLoader(train_dataset_low, batch_size=8, shuffle=True, num_workers=0)
val_loader_low = DataLoader(val_dataset_low, batch_size=8, shuffle=False, num_workers=0)

# =============================================================================
# Create FULL RESOLUTION datasets (512x512) - for high-res experiments
# =============================================================================
train_dataset_full = ISBI2012Dataset(split="train", transform=True, resolution=512)
val_dataset_full = ISBI2012Dataset(split="val", transform=False, resolution=512)

print(f"\n=== Full Resolution Dataset (512x512) ===")
print(f"Training samples: {len(train_dataset_full)}")
print(f"Validation samples: {len(val_dataset_full)}")

train_loader_full = DataLoader(train_dataset_full, batch_size=2, shuffle=True, num_workers=0)
val_loader_full = DataLoader(val_dataset_full, batch_size=2, shuffle=False, num_workers=0)

# Verify shapes for both resolutions
sample_low, mask_low = train_dataset_low[0]
sample_full, mask_full = train_dataset_full[0]

print(f"\n=== Sample Shapes ===")
print(f"Downsampled: image {sample_low.shape}, mask {mask_low.shape}")
print(f"Full res:    image {sample_full.shape}, mask {mask_full.shape}")

# Compare downsampled vs full resolution
import matplotlib.pyplot as plt

fig, axes = plt.subplots(2, 4, figsize=(14, 7))
fig.suptitle("Downsampled (128×128) vs Full Resolution (512×512)", fontsize=14)

# Get samples from both datasets (same index for comparison)
for i in range(2):
    img_low, mask_low = train_dataset_low[i]
    img_full, mask_full = train_dataset_full[i]

    # Downsampled images
    axes[0, i*2].imshow(img_low[0].numpy(), cmap='gray')
    axes[0, i*2].set_title(f"128×128 Image {i+1}")
    axes[0, i*2].axis('off')

    axes[0, i*2+1].imshow(mask_low[0].numpy(), cmap='gray')
    axes[0, i*2+1].set_title(f"128×128 Mask {i+1}")
    axes[0, i*2+1].axis('off')

    # Full resolution images
    axes[1, i*2].imshow(img_full[0].numpy(), cmap='gray')
    axes[1, i*2].set_title(f"512×512 Image {i+1}")
    axes[1, i*2].axis('off')

    axes[1, i*2+1].imshow(mask_full[0].numpy(), cmap='gray')
    axes[1, i*2+1].set_title(f"512×512 Mask {i+1}")
    axes[1, i*2+1].axis('off')

plt.tight_layout()
plt.show()

# Memory comparison
print("\n=== Memory Usage per Batch ===")
print(f"Downsampled (batch=8): {8 * 1 * 128 * 128 * 4 / 1024:.1f} KB")
print(f"Full res (batch=2):    {2 * 1 * 512 * 512 * 4 / 1024:.1f} KB")
print(f"Full res (batch=8):    {8 * 1 * 512 * 512 * 4 / 1024:.1f} KB (if you have enough GPU memory)")

# Step 1: initialize the U-Net model and move it to the GPU if available
print(f"Model initialized on {device}")

# Step 2: Create an Adam optimizer with the learning rate
learning_rate = 0.001

# Step 3: Implement the weighted cross entropy, use a weight of 9 and reduction mean.

print(f"Optimizer: Adam with learning rate {learning_rate}")

# Training loop for downsampled dataset, train for number of epochs
num_epochs = 200
print(f"\n=== Training on Downsampled Dataset (128×128) for {num_epochs} epochs ===")
for epoch in range(num_epochs):

    # Step 4.1: Set model to train mode


    # Step 4.2 initialize the loss at 0.0 at the start of each epoch


    # Step 4.3: Iterate over the train loader
    for images, masks in train_loader_low:

        # Step 4.4 Move the images and masks to the device


        # Step 4.5: Before each training step, remove the old gradients


        # Step 4.6 Calculate the predictions and the loss


        # Step 4.7: Update the model with the calculated gradients & increment the step in the optimizer


        # Step 4.8: Add the loss to the total train loss and multiply by the batch size


    # Step 5: Average train loss


    # Step 6: Implement Validation (You get this one for free!)
    # Same as with training, except in eval mode and without gradient updates or optimizer.step()
    model.eval()
    val_loss = 0.0
    with torch.no_grad():
        for images, masks in val_loader_low:
            images = images.to(device)
            masks = masks.to(device)
            outputs = model(images)
            loss = loss_function(outputs, masks)
            val_loss += loss.item() * images.size(0)

    val_loss /= len(val_loader_low.dataset)

    print(f"Epoch {epoch+1}/{num_epochs} - Train Loss: {train_loss:.4f} - Val Loss: {val_loss:.4f}")