import gradio as gr
from model import NeuralStyleTransfer
import tensorflow as tf
from keras import backend as K
import numpy as np


def change_dtype_inputs(
    n_style_layers,
    n_content_layers,
    epochs,
    learning_rate,
    steps_per_epoch,
    style_weight,
    content_weight,
    var_weight,
):
    return (
        int(n_style_layers),
        int(n_content_layers),
        int(epochs),
        float(learning_rate),
        int(steps_per_epoch),
        float(style_weight),
        float(content_weight),
        float(var_weight),
    )


def fit_style_transfer(
    style_image,
    content_image,
    extractor="inception_v3",
    n_style_layers=2,
    n_content_layers=3,
    epochs=4,
    learning_rate=60.0,
    steps_per_epoch=100,
    style_weight=1e-2,
    content_weight=1e-4,
    var_weight=1e-12,
):
    """
    Fit the style transfer model to the content and style images.

    Parameters
    ----------

    style_image: str
        The path to the style image.

    content_image: str
        The path to the content image.

    extractor: str
        The name of the feature extractor to use. Options are
        "inception_v3", "vgg19", "resnet50", and "mobilenet_v2".

    n_style_layers: int
        The number of layers to use for the style loss.

    n_content_layers: int
        The number of layers to use for the content loss.

    epochs: int
        The number of epochs to train the model for.

    learning_rate: float
        The learning rate to use for the Adam optimizer.

    steps_per_epoch: int
        The number of steps to take per epoch.

    style_weight: float
        The weight to use for the style loss.

    content_weight: float
        The weight to use for the content loss.

    var_weight: float
        The weight to use for the total variation loss.

    Returns
    -------
    display_image: np.array
    """

    (
        n_style_layers,
        n_content_layers,
        epochs,
        learning_rate,
        steps_per_epoch,
        style_weight,
        content_weight,
        var_weight,
    ) = change_dtype_inputs(
        n_style_layers,
        n_content_layers,
        epochs,
        learning_rate,
        steps_per_epoch,
        style_weight,
        content_weight,
        var_weight,
    )

    model = NeuralStyleTransfer(
        style_image=style_image,
        content_image=content_image,
        extractor=extractor,
        n_style_layers=n_style_layers,
        n_content_layers=n_content_layers,
    )

    style_image = model.style_image
    content_image = model.content_image

    content_and_style_layers = model.get_output_layers()

    # build the model with the layers we need to extract the features from
    K.clear_session()
    model.build(content_and_style_layers)

    style_features = model.get_features(style_image, type="style")
    content_features = model.get_features(content_image, type="content")

    optimizer = tf.optimizers.Adam(
        tf.keras.optimizers.schedules.ExponentialDecay(
            initial_learning_rate=learning_rate, decay_steps=100, decay_rate=0.80
        )
    )

    generated_image = tf.cast(content_image, tf.float32)
    generated_image = tf.Variable(generated_image)

    step = 0

    for epoch in range(epochs):
        for step in range(steps_per_epoch):
            losses = model._update_image_with_style(
                generated_image,
                style_features,
                content_features,
                style_weight,
                content_weight,
                optimizer,
                var_weight,
            )

            display_image = model.tensor_to_image(generated_image)

            step += 1

            style_loss, content_loss, var_loss = losses

            yield np.array(display_image), style_loss, content_loss, var_loss, epoch, step



def main():
    content_image = gr.Image(type="filepath", label="Content Image", shape=(512, 512))
    style_image = gr.Image(type="filepath", label="Style Image", shape=(512, 512))

    extractor = gr.Dropdown(
        ["inception_v3", "vgg19", "resnet50", "mobilenet_v2"],
        label="Feature Extractor",
        value="inception_v3",

    )

    n_content_layers = gr.Slider(
        1,
        5,
        value=3,
        step=1,
        label="Content Layers",
    )

    n_style_layers = gr.Slider(
        1,
        5,
        value=2,
        step=1,
        label="Style Layers",
    )

    epochs = gr.Slider(2, 20, value=4, step=1, label="Epochs")

    learning_rate = gr.Slider(1, 100, value=60, step=1, label="Learning Rate")

    steps_per_epoch = gr.Slider(
        1,
        100,
        value=80,
        step=1,
        label="Steps Per Epoch",
    )

    style_weight = gr.Slider(
        1e-4,
        0.5,
        value=1e-1,
        step=1e-4,
        label="Style Weight",
    )

    content_weight = gr.Slider(
        1e-3,
        0.5,
        value=0.3,
        step=1e-4,
        label="Content Weight",
    )

    var_weight = gr.Slider(
        0,
        1e-1,
        value=1e-5,
        step=1e-12,
        label="Total Variation Weight",
    )

    inputs = [
        style_image,
        content_image,
        extractor,
        n_style_layers,
        n_content_layers,
        epochs,
        learning_rate,
        steps_per_epoch,
        style_weight,
        content_weight,
        var_weight,
    ]

    examples = [
        [
            "examples/style_1.jpg",
            "examples/content_1.jpg",
            "inception_v3",
            3,
            2,
            4,
            60,
            100,
            1e-2,
            1e-2,
            1e-11,
        ],
        [
            "examples/style_2.jpg",
            "examples/content_2.jpg",
            "inception_v3",
            3,
            2,
            4,
            60,
            100,
            1e-2,
            1e-2,
            1e-11,
        ],
        [
            "examples/style_3.jpg",
            "examples/content_3.jpg",
            "inception_v3",
            3,
            2,
            4,
            60,
            100,
            1e-2,
            1e-2,
            1e-11,
        ]

    ]

    output_image = gr.Image(type="numpy", label="Output Image", shape=(512, 512))

    style_loss = gr.Number(label="Current Style Loss")

    content_loss = gr.Number(label="Current Content Loss")

    var_loss = gr.Number(label="Current Total Variation Loss")

    curr_epoch = gr.Number(label="Current Epoch")

    curr_step = gr.Number(label="Current Step")



    outputs = [output_image, style_loss, content_loss, var_loss, curr_epoch, curr_step]

    interface = gr.Interface(
        fn=fit_style_transfer,
        inputs=inputs,
        outputs=outputs,
        examples=examples
    )

    interface.queue().launch(debug=True)


main()