# *************************************************************************** #
#                                                                              #
#    style_function.py                                                         #
#                                                                              #
#    By: Widium <ebennace@student.42lausanne.ch>                               #
#    Github : https://github.com/widium                                        #
#                                                                              #
#    Created: 2022/11/15 13:38:04 by  ebennace                                 #
#    Updated: 2023/05/03 16:05:48 by Widium                                    #
#                                                                              #
# **************************************************************************** #

import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt 

from tensorflow.keras.applications import VGG19
from tensorflow.keras.optimizers import Optimizer
from tensorflow import Tensor
from tensorflow import Variable
from tensorflow import GradientTape 
from keras import Model

from .compute import optimize_gradients
from .extract import get_features_map
from .extract import extract_style

# ===================================================== # 

def compute_style_loss(
    style_generated : Tensor, 
    style_target : Tensor
)->float:
    """
    Compute the Style Loss of the Generated Image with the Generated and the Original Style as Tensor.

    1. Iterate through the generated and target style tensors.
    2. Compute the mean squared error for each style layer.
    3. Append the style layer loss to the list of style losses.
    4. Compute the average style loss across all style layers.

    Args:
        style_generated (Tensor): List of image tensors representing the generated style.
        style_target (Tensor): List of image tensors representing the target style.

    Returns:
        float: Mean squared error of style differences.
    """
    all_style_loss = list()

    for generated, target in zip(style_generated, style_target):

        style_layer_loss = tf.reduce_mean((generated - target)**2)
        all_style_loss.append(style_layer_loss)

    num_style_layers = len(all_style_loss)
    style_loss = tf.add_n(all_style_loss) / num_style_layers

    return (style_loss)

# ===================================================== # 

@tf.function
def update_style(
    model : Model,
    style_target : Tensor, 
    generated_img : Variable,
    optimizer : Optimizer
):
    """
    Updates the generated image to minimize the style loss.

    1. Extracts the features map from the model for the generated image.
    2. Extracts the style from the features map.
    3. Computes the style loss based on the style and style target.
    4. Calculates the gradients of the style loss with respect to the generated image.
    5. Updates the generated image using the optimizer and gradients.

    Args:
        model (Model): The pre-trained CNN model (e.g., VGG19) for feature extraction.
        style_target (Tensor): The target style features as a tensor.
        generated_img (Variable): The generated image as a TensorFlow Variable.
        optimizer (Optimizer): The optimizer used for updating the generated image.
    """

    with GradientTape() as tape :

        features_map = get_features_map(model, generated_img)
        style_generated = extract_style(features_map)
        
        style_loss = compute_style_loss(style_generated, style_target)

    gradients = tape.gradient(style_loss, generated_img)
    
    optimize_gradients(
        gradients=gradients,
        optimizer=optimizer, 
        generated_img=generated_img,
    )

# ===================================================== #