import tensorflow as tf
import os
import cv2
import numpy as np
import random
from tqdm import tqdm

SEED = 7
np.random.seed(SEED)
tf.random.set_seed(SEED)
random.seed(SEED)

def load_data(dir_path, img_size=(96, 96)):
    X = []
    y = []
    i = 0
    labels = dict()
    for path in tqdm(sorted(os.listdir(dir_path))):
        if not path.startswith('.'):
            labels[i] = path
            for file in os.listdir(dir_path + path):
                if not file.startswith('.'):
                    img = cv2.imread(os.path.join(dir_path, path, file))
                    img = cv2.resize(img, img_size)
                    img = img / 255.0
                    X.append(img)
                    y.append(i)
            i += 1
    X = np.array(X)
    y = np.array(y)
    print(f'{len(X)} images loaded from {dir_path} directory.')
    return X, y, labels

 
TRAIN_DIR = 'data/train/'
TEST_DIR = 'data/test/'
VAL_DIR = 'data/val/'
IMG_SIZE = (96, 96)

X_train, y_train, labels = load_data(TRAIN_DIR, IMG_SIZE)
X_test, y_test, _ = load_data(TEST_DIR, IMG_SIZE)
X_val, y_val, _ = load_data(VAL_DIR, IMG_SIZE)


 # Low contrast
"""
def add_low_contrast(image, factor=0.5):
    image = np.float32(image)
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    avg_intensity = np.mean(gray)
    low_contrast_image = image * factor + avg_intensity * (1 - factor)
    low_contrast_image = np.clip(low_contrast_image, 0., 1.)  
    return low_contrast_image

def add_low_contrast_to_dataset(X):
    X_low_contrast = []
    for img in X:
        low_contrast_img = add_low_contrast(img)
        X_low_contrast.append(low_contrast_img)
    return np.array(X_low_contrast)

X_train_low_contrast = add_low_contrast_to_dataset(X_train)
X_test_low_contrast = add_low_contrast_to_dataset(X_test)
X_val_low_contrast = add_low_contrast_to_dataset(X_val)
"""
 # Gaussian Noise
"""
def add_gaussian_noise(image, mean=0, sigma=0.04):
    noise = np.random.normal(mean, sigma, image.shape)
    noisy_image = np.clip(image + noise, 0., 1.)  # Ensure the pixel values stay within [0, 1]
    return noisy_image

def add_noise_to_dataset(X):
    X_noisy = []
    for img in X:
        noisy_img = add_gaussian_noise(img)
        X_noisy.append(noisy_img)
    return np.array(X_noisy)

X_train_noisy = add_noise_to_dataset(X_train)
X_test_noisy = add_noise_to_dataset(X_test)
X_val_noisy = add_noise_to_dataset(X_val)
"""

from tensorflow.keras.layers import *
from tensorflow.keras.models import Model

class InvolutionLayer(tf.keras.layers.Layer):
    def __init__(self, Channel, Group_Number, Kernel_Size, Stride, Reduction_Ratio, **kwargs):
        super().__init__(**kwargs)
        self.channel = Channel
        self.group_number = Group_Number
        self.kernel_size = Kernel_Size
        self.stride = Stride
        self.red_ratio = Reduction_Ratio

    def build(self, input_shape):
        (_, height, width, num_channels) = input_shape
        height = height // self.stride
        width = width // self.stride

        self.stride_layer = (tf.keras.layers.AveragePooling2D(
            pool_size=self.stride, strides=self.stride, padding="same"
        )
            if self.stride > 1
            else tf.identity
        )

        self.kernel_gen = tf.keras.Sequential([
            tf.keras.layers.Conv2D(filters=self.channel // self.red_ratio, kernel_size=1),
            tf.keras.layers.BatchNormalization(),
            tf.keras.layers.ReLU(),
            tf.keras.layers.Conv2D(filters=self.kernel_size * self.kernel_size * self.group_number, kernel_size=1)
        ])

        self.kernel_reshape = tf.keras.layers.Reshape(
            target_shape=(
                height,
                width,
                self.kernel_size * self.kernel_size,
                1,
                self.group_number,
            )
        )

        self.input_patches_reshape = tf.keras.layers.Reshape(
            target_shape=(
                height,
                width,
                self.kernel_size * self.kernel_size,
                num_channels // self.group_number,
                self.group_number,
            )
        )

        self.output_reshape = tf.keras.layers.Reshape(
            target_shape=(height, width, num_channels)
        )

    def call(self, x):
        kernel_input = self.stride_layer(x)
        kernel = self.kernel_gen(kernel_input)
        kernel = self.kernel_reshape(kernel)

        input_patches = tf.image.extract_patches(
            images=x,
            sizes=[1, self.kernel_size, self.kernel_size, 1],
            strides=[1, self.stride, self.stride, 1],
            rates=[1, 1, 1, 1],
            padding='SAME',
        )

        input_patches = self.input_patches_reshape(input_patches)

        output = tf.multiply(kernel, input_patches)
        output = tf.reduce_sum(output, axis=3)
        output = self.output_reshape(output)

        return output

    def get_config(self):
        config = super().get_config().copy()
        config.update({
            "channel": self.channel,
            "group_number": self.group_number,
            "kernel_size": self.kernel_size,
            "stride": self.stride,
            "reduction_ratio": self.red_ratio,
        })
        return config

def DepthwiseConvModule(filters, kernel_size, activation='relu', depth_multiplier=1):
    def layer(x):
        x = DepthwiseConv2D(kernel_size, depth_multiplier=depth_multiplier, padding='same')(x)
        x = Conv2D(filters, 1, activation=activation)(x)
        return x
    return layer

def DilatedConvModule(filters, kernel_size, dilation_rate, activation='relu'):
    def layer(x):
        x = Conv2D(filters, kernel_size, dilation_rate=dilation_rate, padding='same', activation=activation)(x)
        return x
    return layer

def InvolutionModule(channel, group_number, kernel_size, stride, reduction_ratio, activation='relu'):
    def layer(x):
        x = InvolutionLayer(Channel=channel, Group_Number=group_number, Kernel_Size=kernel_size, Stride=stride, Reduction_Ratio=reduction_ratio)(x)
        x = ReLU()(x)
        return x
    return layer

def MCICModule(filters, involution_params, depth_multiplier=1, dilation_rate=2, pool_size=(2, 2)):
    def layer(x):
        residual = x  # For the skip connection
        residual = Conv2D(filters, (1, 1), padding='same')(residual)

        x = DepthwiseConvModule(filters=filters, kernel_size=(3, 3), depth_multiplier=depth_multiplier)(x)
        x = DilatedConvModule(filters=filters, kernel_size=(3, 3), dilation_rate=dilation_rate)(x)
        x = InvolutionModule(**involution_params)(x)
        x = Add()([x, residual])
        x = MaxPooling2D(pool_size=pool_size)(x)
        return x
    return layer

def MultiScaleFeatureFusionBlock(filters, activation='relu'):
    def layer(x):
        conv1x1 = Conv2D(filters, (1, 1), padding='same', activation=activation)(x)
        conv3x3 = Conv2D(filters, (3, 3), padding='same', activation=activation)(x)
        conv5x5 = Conv2D(filters, (5, 5), padding='same', activation=activation)(x)
        concatenated = Concatenate()([conv1x1, conv3x3, conv5x5])
        fused = Conv2D(filters, (1, 1), padding='same', activation=activation)(concatenated)
        fused = ReLU()(fused)
        residual = Conv2D(filters, (1, 1), padding='same')(x)
        fused = Add()([residual, fused])
        return fused
    return layer

def build_MCIC_model():
    input_img = Input(shape=(96, 96, 3), name='ImageInput')

    involution_params = {
        'channel': 16,
        'group_number': 4,
        'kernel_size': 3,
        'stride': 1,
        'reduction_ratio': 2
    }

    x = MultiScaleFeatureFusionBlock(16, activation='relu')(input_img)
    x = MCICModule(32, involution_params)(x)
    involution_params['channel'] = 64
    x = MCICModule(64, involution_params)(x)
    involution_params['channel'] = 128
    x = MCICModule(128, involution_params, depth_multiplier=2, dilation_rate=4)(x)
    involution_params['channel'] = 256
    x = MCICModule(256, involution_params, depth_multiplier=2, dilation_rate=4)(x)

    model = Model(inputs=input_img, outputs=x)
    return model

feature_extractor = build_MCIC_model()

import lightgbm as lgb
from sklearn.metrics import accuracy_score, log_loss, roc_auc_score, roc_curve, confusion_matrix, classification_report, matthews_corrcoef
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd

def extract_features(model, data):
    features = model.predict(data, batch_size=32)
    return features.reshape(features.shape[0], -1)

train_features_flat = extract_features(feature_extractor, X_train)
val_features_flat = extract_features(feature_extractor, X_val)
test_features_flat = extract_features(feature_extractor, X_test)

train_data = lgb.Dataset(train_features_flat, label=y_train)
val_data = lgb.Dataset(val_features_flat, label=y_val, reference=train_data)

training_metrics = []
validation_metrics = []

def record_eval_results(env):
    def compute_metrics(y_true, y_pred_proba):
        y_pred = np.round(y_pred_proba)
        accuracy = accuracy_score(y_true, y_pred)
        loss = log_loss(y_true, y_pred_proba)
        return {'accuracy': accuracy, 'loss': loss}

    train_pred_proba = env.model.predict(train_features_flat)
    val_pred_proba = env.model.predict(val_features_flat)
    training_metrics.append(compute_metrics(y_train, train_pred_proba))
    validation_metrics.append(compute_metrics(y_val, val_pred_proba))

callbacks = [lgb.record_evaluation(eval_result={}), record_eval_results]

params = {
    'objective': 'binary',
    'metric': ['binary_error', 'binary_logloss'],
    'force_col_wise': True,
    'n_estimators': 200,
    'seed': SEED,
}

lgbm_classifier = lgb.train(params, train_data, valid_sets=[train_data, val_data], callbacks=callbacks)

def plot_metrics(metrics_history, metric_name):
    plt.figure(figsize=(6, 3.5), dpi=120)
    plt.plot([metrics[metric_name] for metrics in metrics_history['train']], label=f'Training {metric_name.capitalize()}', color='#0f044c')
    plt.plot([metrics[metric_name] for metrics in metrics_history['val']], label=f'Validation {metric_name.capitalize()}', color='#00909E')
    plt.title(f'{metric_name.capitalize()} over Iterations')
    plt.xlabel('Iteration')
    plt.ylabel(metric_name.capitalize())
    plt.legend()
    plt.savefig(f'{metric_name}_plot.png', dpi=400)
    plt.show()

plot_metrics({'train': training_metrics, 'val': validation_metrics}, 'accuracy')
plot_metrics({'train': training_metrics, 'val': validation_metrics}, 'loss')

def save_metrics_to_csv(metrics_history, filename):
    df = pd.DataFrame({
        'Iteration': range(1, len(metrics_history['train']) + 1),
        'Train Accuracy': [m['accuracy'] for m in metrics_history['train']],
        'Validation Accuracy': [m['accuracy'] for m in metrics_history['val']],
        'Train Loss': [m['loss'] for m in metrics_history['train']],
        'Validation Loss': [m['loss'] for m in metrics_history['val']]
    })
    df.to_csv(filename, index=False)

save_metrics_to_csv({'train': training_metrics, 'val': validation_metrics}, 'MCICNet_Inv_LGBM_metrics_history.csv')

def specificity_score(y_true, y_pred):
    cm = confusion_matrix(y_true, y_pred)
    tn, fp, fn, tp = cm.ravel()
    return tn / (tn + fp)

def evaluate_model(y_true, y_pred_proba, y_pred, set_name):
    accuracy = accuracy_score(y_true, y_pred)
    roc_auc = roc_auc_score(y_true, y_pred_proba)
    mcc = matthews_corrcoef(y_true, y_pred)
    spec = specificity_score(y_true, y_pred)

    print(f'{set_name} Set Evaluation:')
    print(f'Accuracy: {accuracy}')
    print(f'ROC AUC: {roc_auc:.4f}')
    print(f'MCC: {mcc:.4f}')
    print(f'Specificity: {spec:.4f}')
    print('\nClassification Report:\n', classification_report(y_true, y_pred, digits=4))

    plt.figure(figsize=(6, 4), dpi=120)
    heatmap = sns.heatmap(confusion_matrix(y_true, y_pred),
                          annot=True,
                          fmt='d',
                          cmap='cividis',
                          cbar=True,
                          xticklabels=['glaucoma', 'normal'],
                          yticklabels=['glaucoma', 'normal'])
    plt.title(f'{set_name} Confusion Matrix')
    plt.xlabel('Predicted', fontsize=12)
    plt.ylabel('True', fontsize=12)
    plt.savefig(f'{set_name}_CM.png', dpi=400)
    plt.show()

    fpr, tpr, _ = roc_curve(y_true, y_pred_proba)
    plt.figure(figsize=(6, 4), dpi=120)
    plt.plot(fpr, tpr, color='#0f044c', lw=1.5, label=f'AUC = {roc_auc:.4f}')
    plt.plot([0, 1], [0, 1], color='gray', linestyle='--', lw=1)
    plt.xlim([-0.02, 1.0])
    plt.ylim([0.0, 1.02])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title(f'{set_name} ROC Curve')
    plt.legend(loc='lower right')
    plt.savefig(f'{set_name}_ROC.png', dpi=400)
    plt.show()

for set_name, X, y in [('Train', train_features_flat, y_train), ('Validation', val_features_flat, y_val), ('Test', test_features_flat, y_test)]:
    y_pred_proba = lgbm_classifier.predict(X)
    y_pred = np.round(y_pred_proba)
    evaluate_model(y, y_pred_proba, y_pred, set_name)

import time

start_inference_time = time.time()
test_pred = lgbm_classifier.predict(test_features_flat)
end_inference_time = time.time()
inference_time = end_inference_time - start_inference_time
print(f"Total Inference time on test data: {inference_time:.4f} seconds")
print(f"Inference time per sample: {inference_time / len(test_features_flat):.6f} seconds")

def save_metrics_to_csv(metrics, file_name):
    with open(file_name, 'w', newline='') as csvfile:
        fieldnames = ['False Positive Rate', 'True Positive Rate', 'AUC Score']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()
        for metric in metrics:
            writer.writerow(metric)

validation_metrics_csv = [{'False Positive Rate': fpr, 'True Positive Rate': tpr, 'AUC Score': roc_auc} for fpr, tpr, roc_auc in zip(*roc_curve(y_val, lgbm_classifier.predict(val_features_flat)))]
save_metrics_to_csv(validation_metrics_csv, 'validation_roc.csv')

test_metrics_csv = [{'False Positive Rate': fpr, 'True Positive Rate': tpr, 'AUC Score': roc_auc} for fpr, tpr, roc_auc in zip(*roc_curve(y_test, lgbm_classifier.predict(test_features_flat)))]
save_metrics_to_csv(test_metrics_csv, 'test_roc.csv')


#Model Predictions (visual outcomes of glaucoma detection)
def plot_samples_with_labels(X_test, y_test, y_pred, labels, num_samples=20, cols=4):
    indices = np.random.choice(range(len(X_test)), num_samples, replace=False)
    rows = int(np.ceil(num_samples / cols))
    plt.figure(figsize=(20, 5 * rows))
    for i, idx in enumerate(indices):
        plt.subplot(rows, cols, i + 1)
        img = np.array(X_test[idx], dtype=np.float32)
        if img.max() <= 1.0:
            img = (img * 255).astype(np.uint8)
        elif img.max() > 1.0 and img.max() <= 255:  
            img = img.astype(np.uint8)

        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        
        plt.imshow(img_rgb)
        plt.title(f'Actual: {labels[y_test[idx]]}\nPredicted: {labels[int(y_pred[idx])]}', fontsize=16, fontweight='bold')
        plt.axis('off')  
    plt.tight_layout(pad=2.0)  
    plt.subplots_adjust(top=0.9) 
    plt.savefig('test_samples_with_labels.png', dpi=400)
    plt.show()

y_test_pred_proba = lgbm_classifier.predict(test_features_flat)
y_test_pred = np.round(y_test_pred_proba)

plot_samples_with_labels(X_test, y_test, y_test_pred, labels, num_samples=20, cols=4)

#Feature Map Analysis
feature_extractor = build_MCIC_model()


def get_intermediate_model(model, layer_name):
    intermediate_model = Model(inputs=model.input, outputs=model.get_layer(layer_name).output)
    return intermediate_model

model = get_intermediate_model(feature_extractor, 'add_5')

input_image = X_test[0]  
input_image = np.expand_dims(input_image, axis=0)

feature_map = model.predict(input_image)
feature_map = feature_map[0]

num_filters = feature_map.shape[-1]
size = feature_map.shape[0]

num_cols = 16
num_rows = (num_filters + num_cols - 1) // num_cols

fig, axes = plt.subplots(num_rows, num_cols, figsize=(20, 2 * num_rows))
axes = axes.flatten()

# Plot each feature map
for i in range(num_filters):
    im = axes[i].imshow(feature_map[:, :, i], cmap='jet')
    axes[i].axis('off')

for j in range(num_filters, len(axes)):
    axes[j].axis('off')

plt.subplots_adjust(hspace=0, wspace=0)
plt.suptitle('Feature Maps', fontsize=16, y=0.9)
plt.show()