""" Import statements and check for GPU """

import os

import re

import glob

import math

import cv2

import csv 

import pandas as pd

import matplotlib.pyplot as plt

import numpy as np

from sklearn.model_selection import train_test_split

from transunet import TransUNet

import tensorflow as tf

from tensorflow.keras.models import Model

from tensorflow.keras.optimizers import Adam

from tensorflow.keras.preprocessing.image import ImageDataGenerator

from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping, CSVLogger

from tensorflow import keras

from tensorflow.keras import layers

# List available GPUs

gpus = tf.config.list_physical_devices('GPU')

print("GPUs: ", gpus)

if gpus:

    print("TensorFlow is using the GPU.")

else:

    print("TensorFlow is not using the GPU.")

""" Function Definitions """

def rle_to_binary(rle, shape):

    """

    Decodes run length encoded masks into a binary image

    Parameters:

        rle (list): list containing the starts and lengths that make up each RLE mask

        shape (tuple): the original shape of the associated image

    """

    # Initialize a flat mask with zeros

    mask = np.zeros(shape[0] * shape[1], dtype=np.uint8)

    if rle == '' or rle == '0':  # Handle empty RLE

        return mask.reshape(shape, order='C')

    # Decode RLE into mask

    rle_numbers = list(map(int, rle.split()))

    for i in range(0, len(rle_numbers), 2):

        start = rle_numbers[i] - 1  # Convert to zero-indexed

        length = rle_numbers[i + 1]

        mask[start:start + length] = 1

    # Reshape flat mask into 2D

    return mask.reshape(shape, order='C')

def custom_generator(gdf, dir, batch_size, target_size=(224, 224), test_mode=False):

    """

    Custom data generator that dynamically aligns images and masks using RLE decoding.

    Parameters:

        gdf (GroupBy): Grouped dataframe containing image IDs and RLEs.

        dir (str): Root directory of the dataset.

        batch_size (int): Number of samples per batch.

        target_size (tuple): Target size for resizing (default=(224, 224)).

        test_mode (bool): If True, yields one image and mask at a time.

    """

    ids = list(gdf.groups.keys())

    dir2 = 'train'

    while True:

        sample_ids = np.random.choice(ids, size=batch_size, replace=False)

        images, masks = [], []

        for id_num in sample_ids:

            # Get the dataframe rows for the current image

            img_rows = gdf.get_group(id_num)

            rle_list = img_rows['segmentation'].tolist()

            # Construct the file path for the image

            sections = id_num.split('_')

            case = sections[0]

            day = sections[0] + '_' + sections[1]

            slice_id = sections[2] + '_' + sections[3]

            pattern = os.path.join(dir, dir2, case, day, "scans", f"{slice_id}*.png")

            filelist = glob.glob(pattern)

            if filelist:

                file = filelist[0]

                image = cv2.imread(file, cv2.IMREAD_COLOR)

                if image is None:

                    print(f"Image not found: {file}")

                    continue  # Skip if the image is missing

                # Original shape of the image

                original_shape = image.shape[:2]

                # Resize the image

                resized_image = cv2.resize(image, target_size, interpolation=cv2.INTER_LINEAR)

                # Decode and resize the masks

                mask = np.zeros((target_size[0], target_size[1], len(rle_list)), dtype=np.uint8)

                for i, rle in enumerate(rle_list):

                    if rle != '0':  # Check if the RLE is valid

                        decoded_mask = rle_to_binary(rle, original_shape)

                        resized_mask = cv2.resize(decoded_mask, target_size, interpolation=cv2.INTER_NEAREST)

                        mask[:, :, i] = resized_mask

                if test_mode:

                    # Return individual samples in test mode

                    yield resized_image[np.newaxis], mask[np.newaxis], pattern

                else:

                    images.append(resized_image)

                    masks.append(mask)

        if not test_mode:

            x = np.array(images)

            y = np.array(masks)

            yield x, y, None

""" Loss function: dice loss ignores negative class thus negating class imbalance issues """     

def dice_coef(y_true, y_pred, smooth=1e-6):

    # Ensure consistent data types

    y_true = tf.cast(y_true, tf.float32)

    y_pred = tf.cast(y_pred, tf.float32)

    y_true_f = tf.keras.backend.flatten(y_true)

    y_pred_f = tf.keras.backend.flatten(y_pred)

    intersection = tf.reduce_sum(y_true_f * y_pred_f)

    return (2. * intersection + smooth) / (tf.reduce_sum(y_true_f) + tf.reduce_sum(y_pred_f) + smooth)

def dice_loss(y_true, y_pred):

    y_true = tf.cast(y_true, tf.float32)

    y_pred = tf.cast(y_pred, tf.float32)

    return 1 - dice_coef(y_true, y_pred)

""" Construct pipeline """

# dir = '../path/Dataset'

dir = './Dataset'

target_size = 224

batch_size = 24

epochs = 124

# read the csv file into a dataframe. os.path.join makes code executable across operating systes

df = pd.read_csv(os.path.join('.', dir, 'train.csv'))

df['segmentation'] = df['segmentation'].fillna('0')

# split into training, testing and validation sets

train_ids, temp_ids = train_test_split(df.id.unique(), test_size=0.25, random_state=42)

val_ids, test_ids = train_test_split(temp_ids, test_size=0.5, random_state=42)

# convert dfs into groupby objects to make sure rows are grouped by id

train_grouped_df = df[df.id.isin(train_ids)].groupby('id')

val_grouped_df = df[df.id.isin(val_ids)].groupby('id')

test_grouped_df = df[df.id.isin(test_ids)].groupby('id')

# steps per epoch is typically train length / batch size to use all training examples

train_steps_per_epoch = math.ceil(len(train_ids) / batch_size)

val_steps_per_epoch = math.ceil(len(val_ids) / batch_size)

test_steps_per_epoch = math.ceil(len(test_ids) / batch_size)

# create the training and validation datagens

train_generator = custom_generator(train_grouped_df, dir, batch_size, (target_size, target_size))

val_generator = custom_generator(val_grouped_df, dir, batch_size, (target_size, target_size))

test_generator = custom_generator(test_grouped_df, dir, batch_size, (target_size, target_size), test_mode=True)

""" Build the model or load the trained model """

loading = True

if loading:

    weights_path = './impmodels/model_weights.h5'

    model = TransUNet(image_size=224, pretrain=False)

    model.load_weights(weights_path)

    model.compile(optimizer='adam', loss=dice_loss, metrics=['accuracy'])

else:

    # create the optimizer and learning rate scheduler

    lr_schedule = tf.keras.optimizers.schedules.CosineDecay(

        initial_learning_rate = 1e-3,

        # decay_steps=train_steps_per_epoch * epochs,

        decay_steps=epochs+2,

        alpha=1e-2 

    )

    optimizer = tf.keras.optimizers.Adam(learning_rate=lr_schedule)

    # create the U-net neural network

    model = TransUNet(image_size=target_size, freeze_enc_cnn=False, pretrain=True)

    model.compile(optimizer=optimizer, loss=dice_loss, metrics=['accuracy'])

    # set up model checkpoints and early stopping

    checkpoints_path = os.path.join('Checkpoints', 'model_weights.h5')

    model_checkpoint = ModelCheckpoint(filepath=checkpoints_path, save_best_only=True, monitor='val_loss')

    early_stopping = EarlyStopping(monitor='val_loss', mode='min', patience=8)

    # log the training to a .csv for reference

    csv_logger = CSVLogger('training_log.csv', append=True)

    history = model.fit(train_generator, validation_data=val_generator, steps_per_epoch=train_steps_per_epoch, validation_steps=val_steps_per_epoch, epochs=epochs, callbacks=[model_checkpoint, early_stopping, csv_logger])

""" Display some predictions """

preds = []

ground_truths = []

num_samples = 50

# Generate predictions and ground truths

for i in range(num_samples):

    # Fetch a batch from the test generator

    batch = next(test_generator)

    image, mask = batch  

    preds.append(model.predict(image))  # Predict using the model

    ground_truths.append(mask)

best_threshold = 0.99

# Apply the best threshold to all predictions

final_preds = [(pred >= best_threshold).astype(int) for pred in preds]

# Compute Dice loss for each prediction

for i in range(len(final_preds)):

    loss = dice_loss(ground_truths[i], final_preds[i])  

    print(f"Image {i + 1}: Dice Loss = {loss:.4f}")

def visualize_predictions(generator, model, num_samples=8, target_size=(224, 224)):

    """

    Visualize predictions vs. ground truths overlaid on original images.

    Parameters:

        generator (generator): Data generator

        model (Model): Trained segmentation model

        num_samples (int): Number of samples to visualize

        target_size (tuple): Target size for resizing (default=(224, 224)).

    """

    fig, axes = plt.subplots(num_samples, 3, figsize=(15, 5 * num_samples))

    for i in range(num_samples):

        # Fetch one image and mask from the generator

        image_batch, mask_batch = next(generator)

        image = image_batch[0]  # Single image

        ground_truth = mask_batch[0]  # Corresponding ground truth mask

        # Ensure image is RGB

        if len(image.shape) == 2:

            image = np.stack([image] * 3, axis=-1)  # Convert grayscale to RGB

        # Ensure ground truth is a single-channel binary mask

        if ground_truth.ndim == 3 and ground_truth.shape[-1] == 3:

            ground_truth = ground_truth[:, :, 0]  # Extract the first channel

        # Generate prediction

        raw_prediction = model.predict(image[np.newaxis])[0]  # Add batch dimension for prediction

        # Ensure prediction is single-channel

        if raw_prediction.ndim == 3 and raw_prediction.shape[-1] == 3:

            prediction = raw_prediction[:, :, 0]  # Extract the first channel

        else:

            prediction = raw_prediction

        prediction = (prediction >= 0.99).astype(np.uint8)  # Threshold prediction

        # Create overlays

        gt_overlay = image.copy()

        pred_overlay = image.copy()

        # Overlay ground truth in red

        gt_overlay[ground_truth == 1] = [255, 0, 0]

        # Overlay prediction in green

        pred_overlay[prediction == 1] = [0, 255, 0]

        # Plot original image, ground truth overlay, and prediction overlay

        axes[i, 0].imshow(image)

        axes[i, 0].set_title(f"Image {i + 1}")

        axes[i, 0].axis('off')

        axes[i, 1].imshow(gt_overlay)

        axes[i, 1].set_title(f"Ground Truth Overlay {i + 1}")

        axes[i, 1].axis('off')

        axes[i, 2].imshow(pred_overlay)

        axes[i, 2].set_title(f"Prediction Overlay {i + 1}")

        axes[i, 2].axis('off')

    plt.tight_layout()

    plt.show()

# Call the function with your test generator and trained model

visualize_predictions(test_generator, model, num_samples=24)

def binary_to_rle(binary_mask):

    """

    Converts a binary mask to RLE (Run-Length Encoding).

    """

    # Flatten mask in column-major order

    flat_mask = binary_mask.T.flatten()

    rle = []

    start = -1

    for i, val in enumerate(flat_mask):

        if val == 1 and start == -1:

            start = i

        elif val == 0 and start != -1:

            rle.extend([start + 1, i - start])

            start = -1

    if start != -1:

        rle.extend([start + 1, len(flat_mask) - start])

    return ' '.join(map(str, rle))

def save_predictions_to_csv(test_generator, model, output_csv_path):

    """

    Generates predictions using the trained model and writes them to a CSV file in RLE format.

    Parameters:

        test_generator: The data generator for the test set.

        model: The trained segmentation model.

        output_csv_path: Path to save the CSV file.

    """

    with open(output_csv_path, mode='w', newline='') as csvfile:

        csv_writer = csv.writer(csvfile)

        csv_writer.writerow(['id', 'segmentation'])  # Header row

        for image, masks, ids in test_generator:

            predictions = model.predict(image)

            predictions = (predictions > 0.99).astype(int)  

            for pred_mask, mask_id in zip(predictions, ids):

                rle = binary_to_rle(pred_mask.squeeze())

                csv_writer.writerow([mask_id, rle])

            print(f"Processed {len(ids)} predictions...")

save_predictions_to_csv(test_generator, model, 'model_output.csv')