import random
import numpy as np
import os
import pandas as pd
import cv2
import warnings
from sklearn.model_selection import KFold
import tensorflow as tf
from sklearn.metrics import accuracy_score, confusion_matrix, ConfusionMatrixDisplay
import matplotlib.pyplot as plt
from sklearn.metrics import precision_score, recall_score, f1_score, classification_report

# Ensure h5py is installed
import h5py

# Set the random seed for numpy, tensorflow, and python built-in random module
np.random.seed(42)
tf.random.set_seed(42)
random.seed(42)

warnings.filterwarnings("ignore", category=UserWarning, message=".*iCCP:.*")

# Define data paths
train_real_folder = 'datasets/training_set/real/'
train_fake_folder = 'datasets/training_set/fake/'
test_real_folder = 'datasets/test_set/real/'
test_fake_folder = 'datasets/test_set/fake/'

# Load train image paths and labels
train_image_paths = []
train_labels = []

# Load train_real image paths and labels
for filename in os.listdir(train_real_folder):
    image_path = os.path.join(train_real_folder, filename)
    label = 0  # Real images have label 0
    train_image_paths.append(image_path)
    train_labels.append(label)

# Load train_fake image paths and labels
for filename in os.listdir(train_fake_folder):
    image_path = os.path.join(train_fake_folder, filename)
    label = 1  # Fake images have label 1
    train_image_paths.append(image_path)
    train_labels.append(label)

# Load test image paths and labels
test_image_paths = []
test_labels = []

# Load test_real image paths and labels
for filename in os.listdir(test_real_folder):
    image_path = os.path.join(test_real_folder, filename)
    label = 0  # Assuming test real images are all real (label 0)
    test_image_paths.append(image_path)
    test_labels.append(label)

# Load test_fake image paths and labels
for filename in os.listdir(test_fake_folder):
    image_path = os.path.join(test_fake_folder, filename)
    label = 1  # Assuming test fake images are all fake (label 1)
    test_image_paths.append(image_path)
    test_labels.append(label)

# Create DataFrames
train_dataset = pd.DataFrame({'image_path': train_image_paths, 'label': train_labels})
test_dataset = pd.DataFrame({'image_path': test_image_paths, 'label': test_labels})

# Function to preprocess images
def preprocess_image(image_path):
    """Loads, resizes, and normalizes an image."""
    image = cv2.imread(image_path)
    resized_image = cv2.resize(image, (224, 224))  # Target size defined here
    normalized_image = resized_image.astype(np.float32) / 255.0
    return normalized_image

# Preprocess all images and convert labels to numpy arrays
X = np.array([preprocess_image(path) for path in train_image_paths])
Y = np.array(train_labels)

# Define ResNet50 model
resnet_model = tf.keras.applications.ResNet50(weights='imagenet', include_top=False, input_shape=(224, 224, 3))

# Freeze all layers except the last one
for layer in resnet_model.layers[:-1]:
    layer.trainable = False

# Modify final layer
x = resnet_model.output
x = tf.keras.layers.Flatten()(x)
predictions = tf.keras.layers.Dense(1, activation='sigmoid')(x)  # Binary classification

# Create new model with modified top
new_model = tf.keras.models.Model(inputs=resnet_model.input, outputs=predictions)

# Compile the new model
new_model.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])

# Set parameters for cross-validation
kf = KFold(n_splits=4, shuffle=True, random_state=42)
batch_size = 32
epochs = 5
weights_file = 'model_2.weights.h5'
# Lists to store accuracy and loss for each fold
accuracy_per_fold = []
loss_per_fold = []

# Perform K-Fold Cross-Validation
for train_index, val_index in kf.split(X):
    X_train, X_val = X[train_index], X[val_index]
    Y_train, Y_val = Y[train_index], Y[val_index]
    
    # Load weights if they exist
    if os.path.exists(weights_file):
        new_model.load_weights(weights_file)
        print(f"Loaded weights from {weights_file}")
    
    # Train only the last layer
    history = new_model.fit(X_train, Y_train, epochs=epochs, batch_size=batch_size, verbose=1, validation_data=(X_val, Y_val))
    
    # Save weights after training
    new_model.save_weights(weights_file)
    print(f"Saved weights to {weights_file}")
    
    # Evaluate the model on the validation data
    val_loss, val_accuracy = new_model.evaluate(X_val, Y_val)
    
    # Store the accuracy score for this fold
    accuracy_per_fold.append(val_accuracy)
    loss_per_fold.append(val_loss)
    print(f'Fold accuracy: {val_accuracy*100:.2f}%')
    print(f'Fold loss: {val_loss:.4f}')
    
# Print average accuracy and loss across all folds
print(f'\nAverage accuracy across all folds: {np.mean(accuracy_per_fold)*100:.2f}%')
print(f'Average loss across all folds: {np.mean(loss_per_fold):.4f}')

# Evaluate the preprocessed test images using the final model
test_loss, test_accuracy = new_model.evaluate(np.array([preprocess_image(path) for path in test_image_paths]), np.array(test_labels))
print(f"\nTest Loss: {test_loss}, Test Accuracy: {test_accuracy}")

# Predict labels for the test set
predictions = new_model.predict(np.array([preprocess_image(path) for path in test_image_paths]))
predicted_labels = (predictions > 0.5).astype(int).flatten()

# Summarize the classification results
true_real_correct = np.sum((np.array(test_labels) == 0) & (predicted_labels == 0))
true_real_incorrect = np.sum((np.array(test_labels) == 0) & (predicted_labels == 1))
true_fake_correct = np.sum((np.array(test_labels) == 1) & (predicted_labels == 1))
true_fake_incorrect = np.sum((np.array(test_labels) == 1) & (predicted_labels == 0))

print("\nClassification Summary:")
print(f"Real images correctly classified: {true_real_correct}")
print(f"Real images incorrectly classified: {true_real_incorrect}")
print(f"Fake images correctly classified: {true_fake_correct}")
print(f"Fake images incorrectly classified: {true_fake_incorrect}")

# Print detailed classification report
print("\nClassification Report:")
print(classification_report(test_labels, predicted_labels, target_names=['Real', 'Fake']))

# Plot confusion matrix
cm = confusion_matrix(test_labels, predicted_labels)
disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=['Real', 'Fake'])
disp.plot(cmap=plt.cm.Blues)
plt.title("Confusion Matrix")
plt.show()

# Plot training & validation accuracy values
plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(history.history['accuracy'])
plt.plot(history.history['val_accuracy'])
plt.title('Model accuracy')
plt.ylabel('Accuracy')
plt.xlabel('Epoch')
plt.legend(['Train', 'Validation'], loc='upper left')
plt.xticks(np.arange(0, len(history.history['accuracy']), step=1), np.arange(1, len(history.history['accuracy']) + 1, step=1))

# Plot training & validation loss values
plt.subplot(1, 2, 2)
plt.plot(history.history['loss'])
plt.plot(history.history['val_loss'])
plt.title('Model loss')
plt.ylabel('Loss')
plt.xlabel('Epoch')
plt.legend(['Train', 'Validation'], loc='upper left')
plt.xticks(np.arange(0, len(history.history['loss']), step=1), np.arange(1, len(history.history['loss']) + 1, step=1))

plt.tight_layout()
plt.show()

# Plot validation accuracy and loss per fold
plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(range(1, kf.get_n_splits() + 1), accuracy_per_fold, marker='o')
plt.title('Validation Accuracy per Fold')
plt.xlabel('Fold')
plt.ylabel('Accuracy')

plt.subplot(1, 2, 2)
plt.plot(range(1, kf.get_n_splits() + 1), loss_per_fold, marker='o')
plt.title('Validation Loss per Fold')
plt.xlabel('Fold')
plt.ylabel('Loss')

plt.tight_layout()
plt.show()