# *************************************************************************** #
#                                                                              #
#    compute.py                                                                #
#                                                                              #
#    By: Widium <ebennace@student.42lausanne.ch>                               #
#    Github : https://github.com/widium                                        #
#                                                                              #
#    Created: 2022/11/10 09:05:01 by  ebennace                                 #
#    Updated: 2023/05/03 16:05:48 by Widium                                    #
#                                                                              #
# **************************************************************************** ## =============== Import =================== #

import tensorflow as tf

from tensorflow.keras.optimizers import Optimizer
from tensorflow import Variable

from .image import clip_pixel

# ======================================== #
# def gram_matrix(input_tensor : Tensor):

#   Gram = tf.linalg.einsum('bijc,bijd->bcd', input_tensor, input_tensor)
#   input_shape   = tf.shape(input_tensor)
#   num_locations = tf.cast(input_shape[1]*input_shape[2], tf.float32) #
#   Gram_Normalized = Gram/num_locations
#   return (Gram_Normalized)

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

def flatten_filters(Feature_Maps):
    """
    Flatten the feature maps into a matrix of pixels.
    
    Args:
        Feature_Maps: The feature maps tensor of shape (batch, height, width, filters).
    
    Returns:
        Tensor: The flattened matrix of pixels with shape (batch, nbr_pixels, nbr_filter).
    """
    batch = int(Feature_Maps.shape[0])
    nbr_pixels = int(Feature_Maps.shape[1] * Feature_Maps.shape[2])
    nbr_filter = int(Feature_Maps.shape[3])
    
    matrix_pixels = tf.reshape(Feature_Maps, (batch, nbr_pixels, nbr_filter))
    return (matrix_pixels)

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

def normalize_matrix(G, Feature_Maps):
    """
    Normalize a matrix by dividing it by the number of pixels in the feature maps.
    
    Args:
        G: The matrix to be normalized.
        Feature_Maps: The feature maps tensor used to compute the number of pixels.
    
    Returns:
        Tensor: The normalized matrix.
    """
    height =  tf.cast(Feature_Maps.shape[1], tf.float32)
    width =  tf.cast(Feature_Maps.shape[2], tf.float32)
    number_pixels = height * width
    G = G / number_pixels
    return (G)

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

def gram_matrix(Feature_Maps):
    """
    Compute the Gram matrix of the given feature maps.
    
    Args:
        Feature_Maps: The feature maps tensor of shape (batch, height, width, filters).
    
    Returns:
        Tensor: The computed Gram matrix.
    """
    F = flatten_filters(Feature_Maps)
    Gram = tf.matmul(F, F, transpose_a=True)
    Gram = normalize_matrix(Gram, Feature_Maps)
    return Gram

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

def optimize_gradients(
    gradients, 
    optimizer : Optimizer, 
    generated_img : Variable, 
):
    """
    Optimize gradients, apply them to the generated image, and clip its pixel values.
    
    Args:
        gradients: The gradients to be optimized.
        optimizer: The optimizer used for optimizing the gradients.
        generated_img: The generated image variable that will be updated.
        loss: The loss value used for optimization.
    
    Returns:
        Variable: The updated generated image variable.
    """
    optimizer.apply_gradients(grads_and_vars=[(gradients, generated_img)])
    generated_img.assign(clip_pixel(generated_img))

    return (generated_img)

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