# *************************************************************************** #
#                                                                              #
#    vgg.py                                                                    #
#                                                                              #
#    By: Widium <ebennace@student.42lausanne.ch>                               #
#    Github : https://github.com/widium                                        #
#                                                                              #
#    Created: 2022/11/15 15:25:02 by  ebennace                                 #
#    Updated: 2023/05/03 16:05:48 by Widium                                    #
#                                                                              #
# **************************************************************************** ## =============== Import =================== #
import tensorflow as tf
import numpy as np

from tensorflow.keras.applications import VGG19
from keras import Model

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

def create_list_of_vgg_layer():
    """
    Create a list of VGG19 layer names that are important for style transfer.
    
    Returns:
        list: A list of VGG19 layer names used for style transfer.
    """
    style_layer_names   = [
        'block1_conv1',
        'block2_conv1',
        'block3_conv1',
        'block4_conv1',
        'block5_conv1'
    ]

    return (style_layer_names)

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

def load_vgg19()-> Model:
    """
    Load the pre-trained VGG19 model from Keras with ImageNet weights.
    
    Returns:
        Model: The VGG19 model without the top classification layers.
    """
    vgg = VGG19(include_top=False, weights='imagenet')
    return vgg

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

def create_multi_output_model(style_layers : list)-> Model:
    """
    Create a multi-output model using VGG19 for style transfer.
    
    Args:
        style_layers (list): A list of style layer names from VGG19 model.
        
    Returns:
        Model: A model with multiple outputs for the specified style layers.
    """
    vgg19 = load_vgg19()
    
    layers_name = style_layers
    layers_output = list()
    
    for name in layers_name:
        layer = vgg19.get_layer(name)
        output = layer.output
        layers_output.append(output)

    multi_output_model = Model([vgg19.input], layers_output)
    multi_output_model.trainable = False

    return (multi_output_model)