import gradio as gr

from PIL import Image as im_lib
from PIL import ImageFilter
from random import randint
import numpy as np
from keras import models
from keras import layers
from keras.layers import Layer, Input, concatenate, MaxPooling2D, Conv2D, Dense, Dropout, Flatten
from keras import Model
from sklearn.preprocessing import MinMaxScaler
from scipy.fftpack import dct



def crop_image(image): #Gradio takes image, and automatically converts into an ndarray.
    #This function will return a value which indicates whether the crop succeeded or not
    #and why it failed, if it did. It will also return the cropped image if succeeded, and the original one otherwise.
    #-1 indicates failure because of grayscale or non-RGB.
    #-2 indicates failure because image is too small to crop
    # 1 indicates success.
    if len(np.asarray(image).shape)!=3: #This is a grayscale image.
        return -1, image
    x_dim, y_dim = image.size
    if x_dim < 256 or y_dim < 256:
        return -2, image
    left, upper = randint(0, x_dim-256), randint(0, y_dim-256)
    right, lower = 256+left, 256+upper
    image = image.crop((left,upper,right,lower))
    return 1, image

def HPF_filter(image):
    return im_lib.fromarray(np.asarray(image)-np.asarray(image.filter(ImageFilter.GaussianBlur)))

def final_image_array_single(image):
    cropped_key, image = crop_image(image)
    if cropped_key == 1:
        image_array = np.asarray(HPF_filter(image))
    return image_array

def create_model_single(model_weights_file = "single_channel_model_best_val_real_precision_weights"):
    try:
        recalled_model1
    except NameError:
        pass
    else:
        del recalled_model1
    
    recalled_model1 = models.Sequential()
    recalled_model1.add(layers.Conv2D( 32, (3,3), activation='relu', input_shape=(256,256,3,)))
    recalled_model1.add( layers.MaxPooling2D( (2,2), strides = 2 ) )
    recalled_model1.add( layers.Conv2D(64, (3,3), activation='relu'))
    recalled_model1.add( layers.MaxPooling2D( (2,2), strides=2) )
    recalled_model1.add( layers.Flatten() )
    recalled_model1.add(layers.Dense(64, activation='relu'))
    recalled_model1.add(layers.Dense(14, activation='softmax'))
    recalled_model1.load_weights(model_weights_file)
    recalled_model1.compile(optimizer='adam',
                  loss='categorical_crossentropy')
    return recalled_model1

model_test_single = create_model_single()

def real_or_not_single(image, model = model_test_single):
    cropped_key, cropped_image = crop_image(image)
    if cropped_key == -1:
        return "This image cannot be processed as it is not RGB."
    elif cropped_key == -2:
        return "This image is too small to be processed."
    else:
        image_array = final_image_array_single(cropped_image)
        prediction_list = model.predict(image_array.reshape(1,256,256,3))
        if np.argmax(prediction_list) == len(prediction_list[0]) - 1:
            return "This image is probably real."
        else:
            keywords = ["biggan", "crn", "cyclegan","deepfake","gaugan","imle","progan","san","seeingdark","stargan", "stylegan2", "stylegan",
           "whichfaceisreal"]
           

            return f"This image is probably fake and generated by {keywords[np.argmax(prediction_list)]}."


def normalize_image(image,normalizing_factor=255):
    if image.mode == 'RGB':
        return np.asarray(image).reshape(image.size[0],image.size[1],3)/normalizing_factor
    if image.mode == 'L':
        return np.assarray(image)/normalizing_factor

#combines the filters for each color channel as one image
def highpassrgb(image):
    red, green, blue = image.split()
    return normalize_image(im_lib.merge(mode='RGB',bands=(red.filter(ImageFilter.Kernel((3,3),(0,-1,0,-1,4,-1,0,-1,0),1,0)), green.filter(ImageFilter.Kernel((3,3),(0,-1,0,-1,4,-1,0,-1,0),1,0)), blue.filter(ImageFilter.Kernel((3,3),(0,-1,0,-1,4,-1,0,-1,0),1,0)))))

#grayscale discrete cosine transform
def gdct(image):
    a = np.array(image.convert('L'))
    return dct(dct(a.T, norm='ortho').T, norm='ortho')

#log-scaled and normalized gdct
def normalized_gdct(image):
    array = gdct(image)
    sgn = np.sign(array)
    logscale = sgn*np.log(abs(array)+0.000000001) #symmetric log scale, shifted slightly to be defined at 0
    scaler = MinMaxScaler()
    scaler.fit(array)
    return scaler.transform(logscale)
    

def final_image_array_dual(image): 
    cropped_key, image= crop_image(image)
    if cropped_key == 1:
        two_images = [normalized_gdct(image), highpassrgb(image)]
    return two_images  

def create_model_dual(model_weights_file = "dual_channel_model_best_val_real_precision_weights"):
    try:
        recalled_model2
    except NameError:
        pass
    else:
        del recalled_model2
  
    gdct_input = Input(shape=(256,256,1))
    gdct_1 = Conv2D(filters=32,kernel_size=(3,3),activation='relu')(gdct_input)
    gdct_2 = MaxPooling2D(pool_size=(2,2),strides=2)(gdct_1)
    gdct_3 = Conv2D(filters=64,kernel_size=(3,3),activation='relu')(gdct_2)
    gdct_4 = MaxPooling2D(pool_size=(2,2),strides=2)(gdct_3)

    #highpass filter
    highpass_input = Input(shape=(256,256,3))
    highpass_1 = Conv2D(filters=32,kernel_size=(3,3),activation='relu')(highpass_input)
    highpass_2 = MaxPooling2D(pool_size=(2,2),strides=2)(highpass_1)
    highpass_3 = Conv2D(filters=64,kernel_size=(3,3),activation='relu')(highpass_2)
    highpass_4 = MaxPooling2D(pool_size=(2,2),strides=2)(highpass_3)


    merged_1 = concatenate([gdct_4, highpass_4])
    merged_2 = Flatten()(merged_1)
    merged_3 = Dense(units=64, activation='relu')(merged_2)
    multiclass= Dense(units=14, activation = 'softmax')(merged_3)
    #binary = Dense(units=1,activation='softmax')(merged_3)

    recalled_model2 = Model(inputs = [gdct_input,highpass_input], outputs = multiclass)

    recalled_model2.load_weights(model_weights_file)
    recalled_model2.compile(optimizer='adam',
                  loss='categorical_crossentropy')
    return recalled_model2


model_test_dual = create_model_dual()


def real_or_not_dual(image, model = model_test_dual): #must take an array
    cropped_key, cropped_image = crop_image(image)
    if cropped_key == -1:
        return "This image cannot be processed as it is not RGB."
    elif cropped_key == -2:
        return "This image is too small to be processed."
    else:
        image1, image2 = final_image_array_dual(cropped_image)
        prediction_list = model.predict([image1.reshape(1,256,256,1), image2.reshape(1,256,256,3)])
        if np.argmax(prediction_list) == len(prediction_list[0]) - 1:
            return "This image is probably real."
        else:
            keywords = ["biggan", "crn", "cyclegan","deepfake","gaugan","imle","progan","san","seeingdark","stargan", "stylegan2", "stylegan",
           "whichfaceisreal"]
            

            return f"This image is probably fake and generated by {keywords[np.argmax(prediction_list)]}."


interface1 = gr.Interface(fn = real_or_not_single,
                         title = "AI or Real?", 
                         inputs = gr.Image(show_label = False, type = "pil"), 
                         outputs = "text", 
                         description = "<center>Upload your image and we will determine whether it's <strong> real </strong> or <strong> AI-generated </strong> using a Single Channel Neural Network. <br> Please flag any erroneous output.</center>",
                         allow_flagging = "manual"
                        )


interface2 = gr.Interface(fn = real_or_not_dual,
                         title = "AI or Real?", 
                         inputs = gr.Image(show_label = False, type = "pil"), 
                         outputs = "text", 
                         description = "<center>Upload your image and we will determine whether it's <strong> real </strong> or <strong> AI-generated </strong> using a Dual Channel Neural Network. <br> Please flag any erroneous output.</center>",
                         allow_flagging = "manual"
                        )


demo = gr.TabbedInterface([interface1, interface2], ["Single Channel", "Dual Channel"])


demo.launch()