# Install gradio
# Import necessary libs
import tensorflow as tf
import keras
import numpy as np

# Load the VGG19 model
from keras.applications.vgg19 import VGG19
from keras.applications.vgg19 import preprocess_input, decode_predictions

model = VGG19(include_top=False, weights='imagenet')

# Define style layers and content layers for VGG19
content_layers = 'block5_conv2'

style_layers = ['block1_conv1',
                'block2_conv1',
                'block3_conv1',
                'block4_conv1',
                'block5_conv1']

num_content_layers = len(content_layers)
num_style_layers = len(style_layers)

# Model for extracting style feature
style_model = keras.Model(inputs=model.input, outputs=[model.get_layer(name).output for name in style_layers], name='Style_model')

# Model for extracting content feature
content_model = keras.Model(inputs=model.input, outputs=model.get_layer(content_layers).output, name='Content_model')

# Style model summary
print(f'Style model: {style_model.name}')
style_model.summary()

# Content model summary
print(f'Content model: {content_model.name}')
content_model.summary()

# Compute content loss
def content_loss(gen, content_img):
    return tf.reduce_sum(tf.square(content_img - gen))

# Gram matrix
def gram_matrix(image):
    # Reshape image to 2-D array
    channel = image.shape[-1] # Number of channels
    a = tf.reshape(image, [-1, channel])
    # Compute Gram matrix corresponding to the input image
    dimension = a.shape[0] # Dimension M*N of an MxNxc image
    return tf.matmul(a, a, transpose_a=True) / tf.cast(dimension, tf.float32)

# Style loss
def style_loss(gen, style_img):
    total_loss = 0.
    for i in range(len(gen)):
        gen_gram = gram_matrix(gen[i])
        style_gram = gram_matrix(style_img[i])
        # Compute style loss
        total_loss += tf.reduce_sum(tf.square(gen_gram-style_gram))/4
    return total_loss

# Training parameter
lr = 10.
optimizer = keras.optimizers.Adam(lr)
alpha = 1e-7
beta = 1e-10
style_weight = len(style_layers)

# Train step
@tf.function(reduce_retracing=True)
def train_step(image, process_content_image, process_style_image):
    with tf.GradientTape() as tape:
        # Calculate content loss
        gen_content_feature = content_model(image)
        content_feature  = content_model(process_content_image)
        loss_content = content_loss(gen_content_feature, content_feature)
        # Calculate style loss
        gen_style_feature = style_model(image)
        style_feature = style_model(process_style_image)
        loss_style = style_loss(gen_style_feature, style_feature)
        # Calculate total loss
        total_loss = alpha*loss_content + beta*loss_style/style_weight
        # Calculate gradient
        grad = tape.gradient(total_loss, image)
        # Apply gradient
        optimizer.apply_gradients([(grad, image)])
    # Return loss for visualize
    return total_loss

# Deprocess image
def deprocess_image(x):
    # Util function to convert a tensor into a valid image
    x = x.reshape((256,256,3))
    # Remove zero-center by mean pixel
    x[:, :, 0] += 103.939
    x[:, :, 1] += 116.779
    x[:, :, 2] += 123.68
    # 'BGR'->'RGB'
    x = x[:, :, ::-1]
    x = np.clip(x, 0, 255).astype("uint8")
    return x

def load_image(image, RESO=256):
  image = tf.image.convert_image_dtype(image, tf.float32) * 255
  image = tf.image.resize(image,[RESO, RESO],method=tf.image.ResizeMethod.BICUBIC)
  image = tf.expand_dims(preprocess_input(image), 0)
  return image

def style_transfer(content_image, style_image, step=100):
  process_content_image = load_image(content_image)
  process_style_image = load_image(style_image)
  generate_image = tf.Variable(process_content_image, trainable=True)
  for step in range(1, step+1):
      # Train step
      loss = train_step(generate_image, process_content_image, process_style_image)
  return deprocess_image(generate_image[0].numpy())

import gradio as gr

demo = gr.Interface(
    fn=style_transfer,
    inputs=[gr.Image(label='Input Image'), gr.Image(label='Style Image'), gr.Slider(10, 200, 50, step=10, label='Step', show_label=True)],
    outputs=gr.Image(label='Style Transfer Image'),
)

demo.launch(debug=True, share=True)