import os
#os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'

#Imports
import tensorflow as tf
from tensorflow import keras
import matplotlib.pyplot as plt
import tensorflow_hub as hub

from PIL import Image
import gradio as gr

from helper_functions import *

def style_transfer(input_image, artist):

    style_path_van_gogh = keras.utils.get_file('Starry-Night-canvas-Vincent-van-Gogh-New-1889.jpg', 
                                           'https://cdn.britannica.com/78/43678-050-F4DC8D93/Starry-Night-canvas-Vincent-van-Gogh-New-1889.jpg')
    style_path_davinci = keras.utils.get_file('Leonardo_da_Vinci_-_Mona_Lisa_%28La_Gioconda%29_-_WGA12711.jpg', 
                                          'https://upload.wikimedia.org/wikipedia/commons/f/f2/Leonardo_da_Vinci_-_Mona_Lisa_%28La_Gioconda%29_-_WGA12711.jpg')
    style_path_dali = keras.utils.get_file('The_Persistence_of_Memory.jpg', 
                                       'https://upload.wikimedia.org/wikipedia/en/d/dd/The_Persistence_of_Memory.jpg')
    style_path_monet = keras.utils.get_file('Claude_Monet_-_Water_Lilies_-_Google_Art_Project_%28462013%29.jpg',
                                        'https://upload.wikimedia.org/wikipedia/commons/a/af/Claude_Monet_-_Water_Lilies_-_Google_Art_Project_%28462013%29.jpg')
    style_path_picasso = keras.utils.get_file('Picasso_The_Weeping_Woman_Tate_identifier_T05010_10.jpg', 
                                          'https://upload.wikimedia.org/wikipedia/en/1/14/Picasso_The_Weeping_Woman_Tate_identifier_T05010_10.jpg')
    style_path_rembrandt = keras.utils.get_file('1259px-The_Nightwatch_by_Rembrandt_-_Rijksmuseum.jpg', 
                                            'https://upload.wikimedia.org/wikipedia/commons/thumb/9/94/The_Nightwatch_by_Rembrandt_-_Rijksmuseum.jpg/1259px-The_Nightwatch_by_Rembrandt_-_Rijksmuseum.jpg')
    
    #set dimensions of input image                                        
    oc_max_dim = 1080

    #set parameters for each choice of artist
    if artist == "Vincent van Gogh":
        style_max_dim = 442
        style_path = style_path_van_gogh
    elif artist == "Claude Monet":
        style_max_dim = 256
        style_path = style_path_monet
    elif artist == "Leonardo da Vinci":
        style_max_dim = 442
        style_path = style_path_davinci
    elif artist == "Rembrandt":
        style_max_dim = 256
        style_path = style_path_rembrandt
    elif artist == "Pablo Picasso":
        style_max_dim = 256
        style_path = style_path_picasso
    elif artist == "Salvador Dali":
        style_max_dim = 512
        style_path = style_path_dali
    
    #load content and style images    
    content_image = load_img(input_image, content=True, max_dim=oc_max_dim)
    style_image = load_img(style_path, content=False, max_dim=style_max_dim)
    
    #Load Magenta Arbitrary Image Stylization network
    hub_module = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/1')
    
    #Pass content and style images as arguments in TensorFlow Constant object format
    stylized_image = hub_module(tf.constant(content_image), tf.constant(style_image))[0]
    
    print("stylized_image:")
    print(stylized_image[0])
    print(stylized_image)
    return tf.keras.preprocessing.image.img_to_array(stylized_image[0])
    
app = gr.Interface(
    style_transfer,
    [gr.Image(type='pil'), gr.Radio(["Vincent van Gogh", "Claude Monet", "Leonardo da Vinci", "Rembrandt", "Pablo Picasso", "Salvador Dali"])],
    gr.Image(type='pil'),
    title="Artist Style Transfer Tool",
    description="""
Fast style transfer using the Magenta model lets you make your own art in the style of six famous artists using a pretrained neural network and deep learning! Simply upload an image and select an artist's style to have transferred to your picture. Each artist's styles are based on a single one of their most famous paintings, shown below for reference: Starry Night (van Gogh), Water Lilies (Monet), Mona Lisa (da Vinci), The Night Watch (Rembrandt), The Weeping Woman (Picasso), and The Persistence of Memory (Dali). 
Note that some input images may be rotated 90 degrees in the output to facilitate the style transfer.  
""",
    article = """
<table>
  <tr>
     <td style="text-align:center">van Gogh</td>
     <td style="text-align:center">Monet</td>
     <td style="text-align:center">da Vinci</td>
     <td style="text-align:center">Rembrandt</td>
     <td style="text-align:center">Picasso</td>
     <td style="text-align:center">Dali</td>
  </tr>
  <tr>
    <td><img src="https://cdn.britannica.com/78/43678-050-F4DC8D93/Starry-Night-canvas-Vincent-van-Gogh-New-1889.jpg" width="200"></td>
    <td><img src="https://upload.wikimedia.org/wikipedia/commons/a/af/Claude_Monet_-_Water_Lilies_-_Google_Art_Project_%28462013%29.jpg" width="200"></td>
    <td><img src="https://upload.wikimedia.org/wikipedia/commons/f/f2/Leonardo_da_Vinci_-_Mona_Lisa_%28La_Gioconda%29_-_WGA12711.jpg" width="200"></td>
    <td><img src="https://upload.wikimedia.org/wikipedia/commons/thumb/9/94/The_Nightwatch_by_Rembrandt_-_Rijksmuseum.jpg/1259px-The_Nightwatch_by_Rembrandt_-_Rijksmuseum.jpg" width="200"></td>
    <td><img src="https://upload.wikimedia.org/wikipedia/en/1/14/Picasso_The_Weeping_Woman_Tate_identifier_T05010_10.jpg" width="200"></td>
    <td><img src="https://upload.wikimedia.org/wikipedia/en/d/dd/The_Persistence_of_Memory.jpg" width="200"></td>
  </tr>
</table>

This app uses [Arbitrary Style Transfer with Magenta](https://arxiv.org/abs/1705.06830).    
"""
)

app.launch()