import os
import cv2
import random
import glob
import PIL
import shutil
import numpy as np
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
from skimage import data
from skimage.util import montage
import skimage.transform as skTrans
from skimage.transform import rotate
from skimage.transform import resize
from PIL import Image, ImageOps
import nibabel as nib
import keras
import keras.backend as K
from keras.callbacks import CSVLogger
import tensorflow as tf
from tensorflow.keras.utils import plot_model
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
from tensorflow.keras.models import *
from tensorflow.keras.layers import *
from tensorflow.keras.optimizers import *
from tensorflow.keras.callbacks import ModelCheckpoint, ReduceLROnPlateau, EarlyStopping, TensorBoard
from tensorflow.keras.layers.experimental import preprocessing

TRAIN_DATASET_PATH = "/kaggle/input/brats20-dataset-training-validation/BraTS2020_TrainingData/MICCAI_BraTS2020_TrainingData/"

# load .nii file as a numpy array
test_image_flair = nib.load(TRAIN_DATASET_PATH + "BraTS20_Training_355/BraTS20_Training_355_flair.nii").get_fdata()
print("Shape: ", test_image_flair.shape)
print("Dtype: ", test_image_flair.dtype)

print("Min: ", test_image_flair.min())
print("Max: ", test_image_flair.max())

scaler = MinMaxScaler()

# Scale the test_image_flair array and then reshape it back to its original dimensions.
# This ensures the data is normalized/standardized for model input without altering its spatial structure.
test_image_flair = scaler.fit_transform(test_image_flair.reshape(-1, test_image_flair.shape[-1])).reshape(test_image_flair.shape)

print("Min: ", test_image_flair.min())
print("Max: ", test_image_flair.max())

# rescaling t1
test_image_t1 = nib.load(TRAIN_DATASET_PATH + 'BraTS20_Training_354/BraTS20_Training_354_t1.nii').get_fdata()
test_image_t1 = scaler.fit_transform(test_image_t1.reshape(-1, test_image_t1.shape[-1])).reshape(test_image_t1.shape)

# rescaling t1ce
test_image_t1ce = nib.load(TRAIN_DATASET_PATH + 'BraTS20_Training_354/BraTS20_Training_354_t1ce.nii').get_fdata()
test_image_t1ce = scaler.fit_transform(test_image_t1ce.reshape(-1, test_image_t1ce.shape[-1])).reshape(test_image_t1ce.shape)

# rescaling t2
test_image_t2 = nib.load(TRAIN_DATASET_PATH + 'BraTS20_Training_354/BraTS20_Training_354_t2.nii').get_fdata()
test_image_t2 = scaler.fit_transform(test_image_t2.reshape(-1, test_image_t2.shape[-1])).reshape(test_image_t2.shape)

# we will not rescale the mask
test_image_seg = nib.load(TRAIN_DATASET_PATH + 'BraTS20_Training_354/BraTS20_Training_354_seg.nii').get_fdata()

slice = 95

print("Slice Number: " + str(slice))

plt.figure(figsize=(12, 8))

# T1
plt.subplot(2, 3, 1)
plt.imshow(test_image_t1[:,:,slice], cmap='gray')
plt.title('T1')

# T1ce
plt.subplot(2, 3, 2)
plt.imshow(test_image_t1ce[:,:,slice], cmap='gray')
plt.title('T1ce')

# T2
plt.subplot(2, 3, 3)
plt.imshow(test_image_t2[:,:,slice], cmap='gray')
plt.title('T2')

# Flair
plt.subplot(2, 3, 4)
plt.imshow(test_image_flair[:,:,slice], cmap='gray')
plt.title('FLAIR')

# Mask
plt.subplot(2, 3, 5)
plt.imshow(test_image_seg[:,:,slice])
plt.title('Mask')
plt.show()

# Modality shape
print("Modality: ", test_image_t1.shape)

# Segmentation shape
print("Segmentation: ", test_image_seg.shape)

slice = 95

print("Slice number: " + str(slice))

plt.figure(figsize=(12, 8))

# Apply a 90° rotation with an automatic resizing, otherwise the display is less obvious to analyze
# T1 - Transverse View
plt.subplot(1, 3, 1)
plt.imshow(test_image_t1ce[:,:,slice], cmap='gray')
plt.title('T1 - Transverse View')

# T1 - Frontal View
plt.subplot(1, 3, 2)
plt.imshow(rotate(test_image_t1ce[:,slice,:], 90, resize=True), cmap='gray')
plt.title('T1 - Frontal View')

# T1 - Sagittal View
plt.subplot(1, 3, 3)
plt.imshow(rotate(test_image_t1ce[slice,:,:], 90, resize=True), cmap='gray')
plt.title('T1 - Sagittal View')
plt.show()

plt.figure(figsize=(10, 10))
plt.subplot(1, 1, 1)

# montage allows us to concatenate multiple images of the same size horizontally and vertically
plt.imshow(rotate(montage(test_image_t1ce[:,:,:]), 90, resize=True), cmap ='gray');

# Skip 50:-50 slices since there is not much to see
plt.figure(figsize=(10, 10))
plt.subplot(1, 1, 1)
plt.imshow(rotate(montage(test_image_t1ce[50:-50,:,:]), 90, resize=True), cmap ='gray');

# Skip 50:-50 slices since there is not much to see
plt.figure(figsize=(10, 10))
plt.subplot(1, 1, 1)
plt.imshow(rotate(montage(test_image_seg[50:-50,:,:]), 90, resize=True), cmap ='gray');

# Plotting the segmantation
cmap = matplotlib.colors.ListedColormap(['#440054', '#3b528b', '#18b880', '#e6d74f'])
norm = matplotlib.colors.BoundaryNorm([-0.5, 0.5, 1.5, 2.5, 3.5], cmap.N)

# plotting the 95th slice
plt.imshow(test_image_seg[:,:,95], cmap=cmap, norm=norm)
plt.colorbar()
plt.show()

# Isolation of class 0
seg_0 = test_image_seg.copy()
seg_0[seg_0 != 0] = np.nan

# Isolation of class 1
seg_1 = test_image_seg.copy()
seg_1[seg_1 != 1] = np.nan

# Isolation of class 2
seg_2 = test_image_seg.copy()
seg_2[seg_2 != 2] = np.nan

# Isolation of class 4
seg_4 = test_image_seg.copy()
seg_4[seg_4 != 4] = np.nan

# Define legend
class_names = ['class 0', 'class 1', 'class 2', 'class 4']
legend = [plt.Rectangle((0, 0), 1, 1, color=cmap(i), label=class_names[i]) for i in range(len(class_names))]

fig, ax = plt.subplots(1, 5, figsize=(20, 20))

ax[0].imshow(test_image_seg[:,:, slice], cmap=cmap, norm=norm)
ax[0].set_title('Original Segmentation')
ax[0].legend(handles=legend, loc='lower left')

ax[1].imshow(seg_0[:,:, slice], cmap=cmap, norm=norm)
ax[1].set_title('Not Tumor (class 0)')

ax[2].imshow(seg_1[:,:, slice], cmap=cmap, norm=norm)
ax[2].set_title('Non-Enhancing Tumor (class 1)')

ax[3].imshow(seg_2[:,:, slice], cmap=cmap, norm=norm)
ax[3].set_title('Edema (class 2)')

ax[4].imshow(seg_4[:,:, slice], cmap=cmap, norm=norm)
ax[4].set_title('Enhancing Tumor (class 4)')

plt.show()

# lists of directories with studies
train_and_val_directories = [f.path for f in os.scandir(TRAIN_DATASET_PATH) if f.is_dir()]

def pathListIntoIds(dirList):
    x = []
    for i in range(0,len(dirList)):
        x.append(dirList[i][dirList[i].rfind('/')+1:])
    return x

train_and_test_ids = pathListIntoIds(train_and_val_directories);

train_test_ids, val_ids = train_test_split(train_and_test_ids,test_size=0.2)
train_ids, test_ids = train_test_split(train_test_ids,test_size=0.15)

# Print data distribution (Train: 68%, Test: 12%, Val: 20%)
print(f"Train length: {len(train_ids)}")
print(f"Validation length: {len(val_ids)}")
print(f"Test length: {len(test_ids)}")

plt.bar(["Train","Valid","Test"],
        [len(train_ids), len(val_ids), len(test_ids)],
        align='center',
        color=[ 'green','red', 'blue'],
        label=["Train", "Valid", "Test"]
       )

plt.legend()

plt.ylabel('Number of Images')
plt.title('Data Distribution')

plt.show()

# Define seg-areas
SEGMENT_CLASSES = {
    0 : 'NOT tumor',
    1 : 'NECROTIC/CORE', # or NON-ENHANCING tumor CORE
    2 : 'EDEMA',
    3 : 'ENHANCING' # original 4 -> converted into 3
}

# Select Slices and Image Size
VOLUME_SLICES = 100
VOLUME_START_AT = 22 # first slice of volume that we will include
IMG_SIZE=128

import os

# Function to filter out any patients missing their segmentation mask
def clean_id_list(id_list):
    cleaned_list = []
    for patient_id in id_list:
        # Construct the exact path the generator will look for
        case_path = os.path.join(TRAIN_DATASET_PATH, patient_id)
        seg_path = os.path.join(case_path, f'{patient_id}_seg.nii')
        
        # Only keep the ID if the file actually exists on the hard drive
        if os.path.exists(seg_path):
            cleaned_list.append(patient_id)
        else:
            print(f"🗑️ Removed corrupted/missing patient: {patient_id}")
            
    return cleaned_list

# Clean all three lists
print("Scanning lists for corrupted files...")
train_ids = clean_id_list(train_ids)
val_ids = clean_id_list(val_ids)
test_ids = clean_id_list(test_ids)

# Re-initialize the generators with the cleaned lists
training_generator = DataGenerator(train_ids)
valid_generator = DataGenerator(val_ids)
test_generator = DataGenerator(test_ids)

class DataGenerator(keras.utils.Sequence):
    'Generates data for Keras'
    def __init__(self, list_IDs, dim=(IMG_SIZE,IMG_SIZE), batch_size=1, n_channels=2, shuffle=True):
        'Initialization'
        self.dim = dim
        self.batch_size = batch_size
        self.list_IDs = list_IDs
        self.n_channels = n_channels
        self.shuffle = shuffle
        self.on_epoch_end()

    def __len__(self):
        'Denotes the number of batches per epoch'
        return int(np.floor(len(self.list_IDs) / self.batch_size))

    def __getitem__(self, index):
        'Generate one batch of data'
        indexes = self.indexes[index*self.batch_size:(index+1)*self.batch_size]
        Batch_ids = [self.list_IDs[k] for k in indexes]
        X, y = self.__data_generation(Batch_ids)
        return X, y

    def on_epoch_end(self):
        'Updates indexes after each epoch'
        self.indexes = np.arange(len(self.list_IDs))
        if self.shuffle == True:
            np.random.shuffle(self.indexes)

    def __data_generation(self, Batch_ids):
        'Generates data containing batch_size samples' 
        
        # FIX 1: Explicitly initialize arrays as float32
        X = np.zeros((self.batch_size*VOLUME_SLICES, *self.dim, self.n_channels), dtype=np.float32)
        y = np.zeros((self.batch_size*VOLUME_SLICES, 240, 240), dtype=np.float32)

        # Generate data
        for c, i in enumerate(Batch_ids):
            case_path = os.path.join(TRAIN_DATASET_PATH, i)

            data_path = os.path.join(case_path, f'{i}_flair.nii')
            flair = nib.load(data_path).get_fdata()

            data_path = os.path.join(case_path, f'{i}_t1ce.nii')
            t1ce = nib.load(data_path).get_fdata()

            data_path = os.path.join(case_path, f'{i}_seg.nii')
            seg = nib.load(data_path).get_fdata()

            for j in range(VOLUME_SLICES):
                 X[j +VOLUME_SLICES*c,:,:,0] = cv2.resize(flair[:,:,j+VOLUME_START_AT], (IMG_SIZE, IMG_SIZE))
                 X[j +VOLUME_SLICES*c,:,:,1] = cv2.resize(t1ce[:,:,j+VOLUME_START_AT], (IMG_SIZE, IMG_SIZE))
                 y[j +VOLUME_SLICES*c] = seg[:,:,j+VOLUME_START_AT]

        # Generate masks
        y[y==4] = 3
        mask = tf.one_hot(y, 4)
        Y = tf.image.resize(mask, (IMG_SIZE, IMG_SIZE))
        
        # FIX 2: Safely normalize X to prevent Division by Zero (NaNs)
        max_val = np.max(X)
        if max_val > 0:
            X = X / max_val
            
        # FIX 3: Convert TF Tensor 'Y' back to a standard numpy float32 array
        return np.array(X, dtype=np.float32), np.array(Y, dtype=np.float32)

# Reinitialize generators
training_generator = DataGenerator(train_ids)
valid_generator = DataGenerator(val_ids)
test_generator = DataGenerator(test_ids)

# Define a function to display one slice and its segmentation
def display_slice_and_segmentation(flair, t1ce, segmentation):
    fig, axes = plt.subplots(1, 3, figsize=(10, 5))

    axes[0].imshow(flair, cmap='gray')
    axes[0].set_title('Flair')
    axes[0].axis('off')

    axes[1].imshow(t1ce, cmap='gray')
    axes[1].set_title('T1CE')
    axes[1].axis('off')

    axes[2].imshow(segmentation) # Displaying segmentation
    axes[2].set_title('Segmentation')
    axes[2].axis('off')

    plt.tight_layout()
    plt.show()


# Retrieve the batch from the training generator
X_batch, Y_batch = training_generator[8]

# Extract Flair, T1CE, and segmentation from the batch
flair_batch = X_batch[:, :, :, 0]
t1ce_batch = X_batch[:, :, :, 1]
segmentation_batch = np.argmax(Y_batch, axis=-1)  # Convert one-hot encoded to categorical

# Extract the 50th slice from Flair, T1CE, and segmentation
slice_index = 60  # Indexing starts from 0
slice_flair = flair_batch[slice_index]
slice_t1ce = t1ce_batch[slice_index]
slice_segmentation = segmentation_batch[slice_index]

# Display the 50th slice and its segmentation
display_slice_and_segmentation(slice_flair, slice_t1ce, slice_segmentation)

import tensorflow as tf

# dice loss as defined above for 4 classes
def dice_coef(y_true, y_pred, smooth=1.0):
    class_num = 4
    total_loss = 0.0 # Initialize total_loss outside the loop for cleaner syntax
    
    for i in range(class_num):
        # Replace K.flatten with tf.reshape to flatten the arrays
        y_true_f = tf.reshape(y_true[:,:,:,i], [-1])
        y_pred_f = tf.reshape(y_pred[:,:,:,i], [-1])
        
        # Replace K.sum with tf.reduce_sum
        intersection = tf.reduce_sum(y_true_f * y_pred_f)
        loss = ((2. * intersection + smooth) / (tf.reduce_sum(y_true_f) + tf.reduce_sum(y_pred_f) + smooth))
        
        total_loss = total_loss + loss
        
    total_loss = total_loss / class_num
    return total_loss

import tensorflow as tf

# Define per class evaluation of dice coef
def dice_coef_necrotic(y_true, y_pred, epsilon=1e-6):
    # Replace K.sum, K.abs, and K.square with their tf equivalents
    intersection = tf.reduce_sum(tf.abs(y_true[:,:,:,1] * y_pred[:,:,:,1]))
    return (2. * intersection) / (tf.reduce_sum(tf.square(y_true[:,:,:,1])) + tf.reduce_sum(tf.square(y_pred[:,:,:,1])) + epsilon)

def dice_coef_edema(y_true, y_pred, epsilon=1e-6):
    intersection = tf.reduce_sum(tf.abs(y_true[:,:,:,2] * y_pred[:,:,:,2]))
    return (2. * intersection) / (tf.reduce_sum(tf.square(y_true[:,:,:,2])) + tf.reduce_sum(tf.square(y_pred[:,:,:,2])) + epsilon)

def dice_coef_enhancing(y_true, y_pred, epsilon=1e-6):
    intersection = tf.reduce_sum(tf.abs(y_true[:,:,:,3] * y_pred[:,:,:,3]))
    return (2. * intersection) / (tf.reduce_sum(tf.square(y_true[:,:,:,3])) + tf.reduce_sum(tf.square(y_pred[:,:,:,3])) + epsilon)

import tensorflow as tf
import keras

# Computing Precision
def precision(y_true, y_pred):
    y_true_f = tf.cast(y_true, tf.float32)
    y_pred_f = tf.cast(y_pred, tf.float32)
    
    true_positives = tf.reduce_sum(tf.round(tf.clip_by_value(y_true_f * y_pred_f, 0.0, 1.0)))
    predicted_positives = tf.reduce_sum(tf.round(tf.clip_by_value(y_pred_f, 0.0, 1.0)))
    
    return true_positives / (predicted_positives + keras.backend.epsilon())


# Computing Sensitivity (Recall)
def sensitivity(y_true, y_pred):
    y_true_f = tf.cast(y_true, tf.float32)
    y_pred_f = tf.cast(y_pred, tf.float32)
    
    true_positives = tf.reduce_sum(tf.round(tf.clip_by_value(y_true_f * y_pred_f, 0.0, 1.0)))
    possible_positives = tf.reduce_sum(tf.round(tf.clip_by_value(y_true_f, 0.0, 1.0)))
    
    return true_positives / (possible_positives + keras.backend.epsilon())


# Computing Specificity
def specificity(y_true, y_pred):
    y_true_f = tf.cast(y_true, tf.float32)
    y_pred_f = tf.cast(y_pred, tf.float32)
    
    true_negatives = tf.reduce_sum(tf.round(tf.clip_by_value((1.0 - y_true_f) * (1.0 - y_pred_f), 0.0, 1.0)))
    possible_negatives = tf.reduce_sum(tf.round(tf.clip_by_value(1.0 - y_true_f, 0.0, 1.0)))
    
    return true_negatives / (possible_negatives + keras.backend.epsilon())

def build_unet(inputs, ker_init, dropout):
    conv1 = Conv2D(32, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(inputs)
    conv1 = Conv2D(32, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(conv1)

    pool = MaxPooling2D(pool_size=(2, 2))(conv1)
    conv = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(pool)
    conv = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(conv)

    pool1 = MaxPooling2D(pool_size=(2, 2))(conv)
    conv2 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(pool1)
    conv2 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(conv2)

    pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)
    conv3 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(pool2)
    conv3 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(conv3)


    pool4 = MaxPooling2D(pool_size=(2, 2))(conv3)
    conv5 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(pool4)
    conv5 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(conv5)
    drop5 = Dropout(dropout)(conv5)

    up7 = Conv2D(256, 2, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(UpSampling2D(size = (2,2))(drop5))
    merge7 = concatenate([conv3,up7], axis = 3)
    conv7 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(merge7)
    conv7 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(conv7)

    up8 = Conv2D(128, 2, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(UpSampling2D(size = (2,2))(conv7))
    merge8 = concatenate([conv2,up8], axis = 3)
    conv8 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(merge8)
    conv8 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(conv8)

    up9 = Conv2D(64, 2, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(UpSampling2D(size = (2,2))(conv8))
    merge9 = concatenate([conv,up9], axis = 3)
    conv9 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(merge9)
    conv9 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(conv9)

    up = Conv2D(32, 2, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(UpSampling2D(size = (2,2))(conv9))
    merge = concatenate([conv1,up], axis = 3)
    conv = Conv2D(32, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(merge)
    conv = Conv2D(32, 3, activation = 'relu', padding = 'same', kernel_initializer = ker_init)(conv)

    conv10 = Conv2D(4, (1,1), activation = 'softmax')(conv)

    return Model(inputs = inputs, outputs = conv10)

input_layer = Input((IMG_SIZE, IMG_SIZE, 2))

generator = build_unet(input_layer, 'he_normal', 0.2)

plot_model(generator,
           show_shapes = True,
           show_dtype=False,
           show_layer_names = True,
           rankdir = 'TB',
           expand_nested = False,
           dpi = 70)

import tensorflow as tf
from tensorflow.keras.layers import Input, concatenate, Conv2D, BatchNormalization, LeakyReLU, MaxPooling2D
from tensorflow.keras.models import Model

def build_proper_3x3_patchgan(image_shape=(240, 240, 2), mask_channels=4):
    """
    A proper 3x3 PatchGAN Discriminator.
    Uses Stride=1 for clean feature extraction, and MaxPooling for downsampling
    to completely avoid 3x3 stride artifacts.
    """
    initializer = tf.random_normal_initializer(0., 0.02)

    # 1. Inputs
    inp = Input(shape=image_shape, name='mri_input')
    tar = Input(shape=(image_shape[0], image_shape[1], mask_channels), name='mask_input')

    # Concatenate the MRI and the Mask
    x = concatenate([inp, tar])

    # 2. Block 1: 3x3 Conv (Stride 1) + MaxPool
    x = Conv2D(64, 3, strides=1, padding='same', kernel_initializer=initializer, use_bias=False)(x)
    x = LeakyReLU(alpha=0.2)(x)
    x = MaxPooling2D(pool_size=(2, 2))(x) # Clean downsampling

    # 3. Block 2: 3x3 Conv (Stride 1) + MaxPool
    x = Conv2D(128, 3, strides=1, padding='same', kernel_initializer=initializer, use_bias=False)(x)
    x = BatchNormalization()(x)
    x = LeakyReLU(alpha=0.2)(x)
    x = MaxPooling2D(pool_size=(2, 2))(x)

    # 4. Block 3: 3x3 Conv (Stride 1) + MaxPool
    x = Conv2D(256, 3, strides=1, padding='same', kernel_initializer=initializer, use_bias=False)(x)
    x = BatchNormalization()(x)
    x = LeakyReLU(alpha=0.2)(x)
    x = MaxPooling2D(pool_size=(2, 2))(x)

    # 5. Block 4: The "Patch" Feature Layer (No pooling here)
    x = Conv2D(512, 3, strides=1, padding='same', kernel_initializer=initializer, use_bias=False)(x)
    x = BatchNormalization()(x)
    x = LeakyReLU(alpha=0.2)(x)

    # 6. Output Layer: Final 3x3 to create the 2D Probability Grid
    last = Conv2D(1, 3, strides=1, padding='same', kernel_initializer=initializer)(x)

    return Model(inputs=[inp, tar], outputs=last, name="Proper_3x3_Pool_PatchGAN")

# Initialize it
discriminator = build_proper_3x3_patchgan(image_shape=(240, 240, 2), mask_channels=4)
discriminator.summary()

import tensorflow as tf

discriminator = build_proper_3x3_patchgan(image_shape=(IMG_SIZE, IMG_SIZE, 2), mask_channels=4)

# 2. Optimizers
generator_optimizer = tf.keras.optimizers.Adam(learning_rate=0.0002, beta_1=0.5)
discriminator_optimizer = tf.keras.optimizers.Adam(learning_rate=0.0002, beta_1=0.5)

# 3. Loss functions
bce_loss = tf.keras.losses.BinaryCrossentropy()
cce_loss = tf.keras.losses.CategoricalCrossentropy()

# 4. Weight for the segmentation loss vs adversarial loss
LAMBDA_SEG = 100

@tf.function
def train_step(mri_images, real_masks):
    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
        
        # 1. Generator creates a fake mask
        fake_masks = generator(mri_images, training=True)

        # 2. Discriminator evaluates both real and fake masks
        real_output = discriminator([mri_images, real_masks], training=True)
        fake_output = discriminator([mri_images, fake_masks], training=True)

        # 3. Calculate Discriminator Loss
        disc_loss_real = bce_loss(tf.ones_like(real_output), real_output)
        disc_loss_fake = bce_loss(tf.zeros_like(fake_output), fake_output)
        total_disc_loss = disc_loss_real + disc_loss_fake

        # 4. Calculate Generator Loss
        # Adversarial loss: generator wants discriminator to guess 1 (Real)
        gen_adv_loss = bce_loss(tf.ones_like(fake_output), fake_output)
        # Segmentation loss: generated mask must match the real mask shape/classes
        gen_seg_loss = cce_loss(real_masks, fake_masks)
        
        total_gen_loss = gen_adv_loss + (LAMBDA_SEG * gen_seg_loss)

    # 5. Calculate gradients
    generator_gradients = gen_tape.gradient(total_gen_loss, generator.trainable_variables)
    discriminator_gradients = disc_tape.gradient(total_disc_loss, discriminator.trainable_variables)

    # 6. Apply gradients
    generator_optimizer.apply_gradients(zip(generator_gradients, generator.trainable_variables))
    discriminator_optimizer.apply_gradients(zip(discriminator_gradients, discriminator.trainable_variables))

    return total_gen_loss, total_disc_loss, gen_seg_loss

import time

EPOCHS = 35

print("Starting cGAN Training...")
train_start_time = time.time()

for epoch in range(EPOCHS):
    epoch_start_time = time.time()
    print(f"\n--- Epoch {epoch + 1}/{EPOCHS} ---")
    
    gen_loss_list = []
    disc_loss_list = []
    seg_loss_list = []
    
    # Iterate over the training generator
    for batch_idx in range(len(training_generator)):
        mri_batch, mask_batch = training_generator[batch_idx]
        
        g_loss, d_loss, seg_loss = train_step(mri_batch, mask_batch)
        
        gen_loss_list.append(g_loss)
        disc_loss_list.append(d_loss)
        seg_loss_list.append(seg_loss)
        
        # Clean logging
        if batch_idx == 0:
            print(f"  ↳ Batch 1 compiled & completed!")
        elif (batch_idx + 1) % 50 == 0:
            print(f"  ↳ Processing batch {batch_idx + 1}/{len(training_generator)}...")

    # Calculate metrics for the epoch
    avg_g_loss = tf.reduce_mean(gen_loss_list)
    avg_d_loss = tf.reduce_mean(disc_loss_list)
    avg_seg_loss = tf.reduce_mean(seg_loss_list)
    
    epoch_duration = time.time() - epoch_start_time
    print(f" Epoch {epoch + 1} Finished in {epoch_duration/60:.2f} mins")
    print(f" Gen Loss: {avg_g_loss:.4f} | Disc Loss: {avg_d_loss:.4f} | Seg Loss: {avg_seg_loss:.4f}")
    
    # Save the generator weights periodically to prevent data loss
    if (epoch + 1) % 5 == 0:
        generator.save_weights(f'cgan_generator_epoch_{epoch+1}.weights.h5')

total_duration = time.time() - train_start_time
print(f"\nTraining Complete! Total Time: {total_duration/60:.2f} minutes")

generator.save('cgan_unet_braTS_3x3_IIcGAN_generator.keras')

generator.save_weights('cgan_unet_braTS_3x3_IIcGAN.weights.h5')

print("✅ Model secured! Check the '/kaggle/working/' directory in your output tab to download it.")

import tensorflow as tf
import numpy as np

# 1. Load the fully trained 3x3 cGAN Generator
print("📥 Loading saved 3x3 II cGAN model...")
# Update this filename if you saved it as something else!
model = tf.keras.models.load_model('/kaggle/working/cgan_unet_braTS_3x3_IIcGAN_generator.keras', compile=False)

print("📊 Evaluating 3x3 II cGAN (Strictly on slices containing tumors)...")
total_psnr = 0.0
total_ssim = 0.0
tumor_slice_count = 0
num_batches = len(valid_generator)

# 2. Run the evaluation loop
for i in range(num_batches):
    mri_batch, mask_batch = valid_generator[i]
    predictions = model.predict(mri_batch, verbose=0)
    
    real_argmax = np.argmax(mask_batch, axis=-1)
    pred_argmax = np.argmax(predictions, axis=-1)
    
    # Loop through each image in the batch individually
    for j in range(len(real_argmax)):
        # ONLY evaluate slices that actually have tumor pixels in the ground truth
        if np.max(real_argmax[j]) > 0:
            r_tensor = tf.convert_to_tensor(real_argmax[j], dtype=tf.float32)
            p_tensor = tf.convert_to_tensor(pred_argmax[j], dtype=tf.float32)
            
            # Add batch and channel dimensions back: (1, Height, Width, 1)
            r_tensor = tf.expand_dims(tf.expand_dims(r_tensor, axis=0), axis=-1)
            p_tensor = tf.expand_dims(tf.expand_dims(p_tensor, axis=0), axis=-1)
            
            # Calculate metrics for this specific tumor slice
            # (max_val=3.0 assumes your classes are 0, 1, 2, 3)
            slice_psnr = tf.image.psnr(r_tensor, p_tensor, max_val=3.0)[0].numpy()
            slice_ssim = tf.image.ssim(r_tensor, p_tensor, max_val=3.0)[0].numpy()
            
            total_psnr += slice_psnr
            total_ssim += slice_ssim
            tumor_slice_count += 1
            
    if (i + 1) % 50 == 0:
        print(f"  ↳ Processed {i + 1}/{num_batches} batches...")

# Calculate true final averages safely
avg_psnr = total_psnr / tumor_slice_count if tumor_slice_count > 0 else 0
avg_ssim = total_ssim / tumor_slice_count if tumor_slice_count > 0 else 0

print("\n🏆 Final 3x3 cGAN Structural Metrics 🏆")
print(f"Evaluated strictly on {tumor_slice_count} tumor slices.")
print(f"Average PSNR: {avg_psnr:.4f} dB")
print(f"Average SSIM: {avg_ssim:.4f}")

import numpy as np
import matplotlib.pyplot as plt
import random

def plot_random_tumor_results(model, data_generator):
    found_tumor = False
    max_attempts = 50 # Prevent infinite loops
    
    for _ in range(max_attempts):
        # 1. Pick a RANDOM batch index
        idx = random.randint(0, len(data_generator) - 1)
        mri_batch, mask_batch = data_generator[idx]
        
        # 2. Check images in this random batch
        indices = list(range(len(mri_batch)))
        random.shuffle(indices) # Look through the batch in random order
        
        for j in indices:
            real_mask_2d = np.argmax(mask_batch[j], axis=-1)
            
            if np.max(real_mask_2d) > 0: # Found a slice with a tumor!
                # Run Inference
                input_tensor = np.expand_dims(mri_batch[j], axis=0)
                prediction = model.predict(input_tensor, verbose=0)
                pred_mask_2d = np.argmax(prediction[0], axis=-1)
                
                # Plotting
                plt.figure(figsize=(15, 5))
                
                plt.subplot(1, 3, 1)
                plt.title(f"MRI (Batch {idx}, Slice {j})")
                plt.imshow(mri_batch[j][:, :, 0], cmap='gray')
                plt.axis('off')
                
                plt.subplot(1, 3, 2)
                plt.title("Ground Truth Mask")
                plt.imshow(real_mask_2d, cmap='magma', vmin=0, vmax=3)
                plt.axis('off')
                
                plt.subplot(1, 3, 3)
                plt.title("3x3 cGAN Prediction")
                plt.imshow(pred_mask_2d, cmap='magma', vmin=0, vmax=3)
                plt.axis('off')
                
                plt.show()
                found_tumor = True
                break
        if found_tumor: break

# Run it multiple times to see different results!
plot_random_tumor_results(model, valid_generator)