import numpy as np
import random
import glob
import sys
import os
from PIL import Image
import matplotlib.pyplot as plt
from tensorflow.keras.losses import mean_squared_error
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.models import load_model
from pathlib import Path

'''
Evaluates the trained model on a new dataset
Uses Mean Square Error as evaluation metric
'''

# Path for evaluation dataset
EVAL_PATH = Path(sys.argv[1]).resolve()
BATCH_SIZE = 1
TEST_NUM = len(glob.glob(str(EVAL_PATH)))


# Normalize input for evaluation
def normalize_batch(imgs):
    return (imgs - np.array([0.485, 0.456, 0.406])) / np.array([0.229, 0.224, 0.225])


# Denormalize output for prediction
def denormalize_batch(imgs, should_clip=True):
    imgs = (imgs * np.array([0.229, 0.224, 0.225])) + np.array([0.485, 0.456, 0.406])

    if should_clip:
        imgs = np.clip(imgs, 0, 1)
    return imgs


# Create a data generator for evaluation
test_imgen = ImageDataGenerator(rescale=1. / 255)


# Compute L2 loss(MSE) for secret image
def custom_loss_1(secret, secret_pred):
    secret_mse = mean_squared_error(secret, secret_pred)
    return secret_mse


# Compute L2 loss(MSE) for cover image
def custom_loss_2(cover, cover_pred):
    cover_mse = mean_squared_error(cover, cover_pred)
    return cover_mse


# Load the trained model
model = load_model(sys.argv[2], custom_objects={'custom_loss_1': custom_loss_1, 'custom_loss_2': custom_loss_2})


# Custom data generator
def generate_generator_multiple(generator, direc):
    genX1 = generator.flow_from_directory(direc, target_size=(224, 224), batch_size=12, shuffle=True, seed=3,
                                          class_mode=None)
    genX2 = generator.flow_from_directory(direc, target_size=(224, 224), batch_size=12, shuffle=True, seed=8,
                                          class_mode=None)

    while True:
        X1i = normalize_batch(genX1.next())
        X2i = normalize_batch(genX2.next())

        yield ({'secret': X1i, 'cover': X2i},
               {'hide_conv_f': X2i, 'revl_conv_f': X1i})  # Yield both images and their mutual label


# Load data using generator
testgenerator = generate_generator_multiple(test_imgen, direc=str(EVAL_PATH.parent))

# Evaluates the model using data generator
score = model.evaluate(testgenerator, steps=TEST_NUM / BATCH_SIZE, verbose=0)

# Print mean square error
print('Mean square error:', score[0])

'''
Test the model on a  random pair of images (test)
Plots the input and output for verification
'''


# Perform prediction for single input
def predict(source, cover):
    # Normalize inputs
    secret = np.array(source / 255.0)
    cover = np.array(cover / 255.0)

    # Predict output
    coverout, secretout = model.predict(
        [normalize_batch(np.reshape(secret, (1, 224, 224, 3))), normalize_batch(np.reshape(cover, (1, 224, 224, 3)))])

    # Postprocess output cover image
    coverout = denormalize_batch(coverout)
    coverout = np.squeeze(coverout) * 255.0
    coverout = np.uint8(coverout)

    # Postprocess output secret image
    secretout = denormalize_batch(secretout)
    secretout = np.squeeze(secretout) * 255.0
    secretout = np.uint8(secretout)

    # Plot output images
    fig_out, ax_out = plt.subplots(1, 2, figsize=(10, 10))
    fig_out.suptitle('Outputs')
    ax_out[0].title.set_text("Secret output")
    ax_out[0].imshow(secretout)
    ax_out[1].title.set_text("Cover output")
    ax_out[1].imshow(coverout)
    print("Prediction Completed!")


# Load random test image pairs
sec, cov = random.sample(os.listdir(str(EVAL_PATH)), k=2)

source = np.array(Image.open(str(EVAL_PATH) + '/' + sec))
cover = np.array(Image.open(str(EVAL_PATH) + '/' + cov))

# Plot input images
fig_in, ax_in = plt.subplots(1, 2, figsize=(10, 10))
fig_in.suptitle('Inputs')
ax_in[0].title.set_text("Secret input")
ax_in[0].imshow(source)
ax_in[1].title.set_text("Cover input")
ax_in[1].imshow(cover)

# Perform prediction
predict(source, cover)

# Sample run: python eval.py dataset/eval_data checkpoints/steg_model-06-0.03.hdf5