# -*- coding: utf-8 -*-
"""
Created on Fri May 24 14:31:20 2024

@author: beni
"""

from keras.models import Sequential
from keras.layers import Conv2D, MaxPool2D, Dropout, Flatten, Dense
from cnn_ela_training import convert_to_ela_image, shuffle_and_split_data, labeling
import pandas as pd
import numpy as np
from PIL import Image
import os
from pylab import *
import re
from PIL import Image, ImageChops, ImageEnhance
import tensorflow as tf
import itertools
from tensorflow.keras.utils import to_categorical 
from tensorflow.keras.optimizers.legacy import RMSprop
from sklearn.metrics import confusion_matrix
import seaborn as sns
import matplotlib.pyplot as plt
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, confusion_matrix
from copy import deepcopy



##########################################
#define the Cnn network as it was in the training file


model = Sequential()

model.add(Conv2D(filters = 32, kernel_size = (5,5),padding = 'valid', 
                     activation ='relu', input_shape = (128,128,3)))
print("Input: ", model.input_shape)
print("Output: ", model.output_shape)

model.add(Conv2D(filters = 32, kernel_size = (5,5),padding = 'valid', 
                     activation ='relu'))
print("Input: ", model.input_shape)
print("Output: ", model.output_shape)

model.add(MaxPool2D(pool_size=(2,2)))

model.add(Dropout(0.25))
print("Input: ", model.input_shape)
print("Output: ", model.output_shape)

model.add(Flatten())
model.add(Dense(256, activation = "relu"))
model.add(Dropout(0.5))
model.add(Dense(2, activation = "softmax"))

model.summary()


# Load saved weights
model.load_weights("ELA_CNN_ART_V2.h5")

#set the optimizer
optimizer = RMSprop(lr=0.0005, rho=0.9, epsilon=1e-08, decay=0.0)
model.compile(optimizer = optimizer , loss = "categorical_crossentropy", metrics=["accuracy"])

###########################################
#label and order the data

test_real_folder = 'datasets/test_set/real/'
test_fake_folder = 'datasets/test_set/fake/'

test_ela_output = 'datasets/training_set/ela_output/'

test_set = labeling(test_real_folder, test_fake_folder)
X_test = []
Y_test = []


# Preprocess test set using ELA and then store the output
for index, row in test_set.iterrows():
    X_test.append(array(convert_to_ela_image(row[0], 90, test_ela_output).resize((128, 128))).flatten() / 255.0)
    Y_test.append(row[1])

# Convert to numpy arrays
X_test = np.array(X_test)
Y_test = to_categorical(Y_test, 2)

# Reshape images
X_test = X_test.reshape(-1, 128, 128, 3)

# Evaluate the model on test set
test_loss, test_accuracy = model.evaluate(X_test, Y_test)
print()
print("~~~~~art Dataset~~~~")
print()
print("Test Loss:", test_loss)
print("Test Accuracy:", test_accuracy)

#######################################################

def calculate_acc(y_true, y_pred):
    

    # Calculate precision
    precision = precision_score(y_true, y_pred)

    # Calculate recall
    recall = recall_score(y_true, y_pred)

    # Calculate F1 score
    f1 = f1_score(y_true, y_pred)
    
    # Calculate confusion matrix
    conf_matrix = confusion_matrix(y_true, y_pred)


    print("Precision:", precision)
    print("Recall:", recall)
    print("F1 Score:", f1)
    print("Confusion Matrix:")

    # Plot confusion matrix
    plt.figure(figsize=(8, 6))
    sns.heatmap(conf_matrix, annot=True, fmt='d', cmap='Blues', cbar=False)
    plt.xlabel('Predicted Label')
    plt.ylabel('True Label')
    plt.title('Confusion Matrix')
    plt.show()
    
############################################
# model preformance metrics and plots

# Get predicted probabilities
Y_pred_prob = model.predict(X_test)

# Convert predicted probabilities to class labels
Y_pred = np.argmax(Y_pred_prob, axis=1)

Y_true = np.argmax(Y_test, axis=1)

# Calculate accuracies
calculate_acc(Y_true, Y_pred)

model.summary()