# ----------------------------------------------------------
# Enhanced Embedding Network
# ----------------------------------------------------------
import keras
import tensorflow as tf
from keras_facenet import FaceNet
# from app import eucledian_distance

def euclidean_distance(vectors):
    x, y = vectors
    return tf.sqrt(tf.reduce_sum(tf.square(x - y), axis=1, keepdims=True))

def build_embedding_network(input_shape):
    # Load the FaceNet model
    facenet = FaceNet()
    base_model = facenet.model

    # Freeze all layers except the last 10 for fine-tuning
    for layer in base_model.layers[:-10]:
        layer.trainable = False

    # Define input and pass it through the base model
    inputs = keras.Input(input_shape)
    x = base_model(inputs)  # Pass input through FaceNet

    # Final Embedding Layer
    x = keras.layers.Dense(256, activation=None)(x)  # No activation before L2 norm
    x = keras.layers.Lambda(lambda x: tf.math.l2_normalize(x, axis=1))(x)  # L2-normalized embeddings

    return keras.Model(inputs, x)

# ----------------------------------------------------------
# Siamese Network
# ----------------------------------------------------------
def build_siamese():
    
    input_shape = (256, 256, 3)  # FaceNet expects 160x160 input
    embedding_network = build_embedding_network(input_shape)

# Define inputs for the Siamese network
    input_1 = keras.layers.Input(input_shape)
    input_2 = keras.layers.Input(input_shape)

# Pass inputs through the embedding network
    tower_1 = embedding_network(input_1)
    tower_2 = embedding_network(input_2)

# Compute Euclidean distance between embeddings
    distance = keras.layers.Lambda(euclidean_distance, name='distance_layer')([tower_1, tower_2])

# Custom Contrastive Output Layer
    output = keras.layers.Dense(1, activation='sigmoid',
                      kernel_regularizer=keras.regularizers.l2(1e-4))(distance)

# Build the Siamese model
    siamese = keras.Model(inputs=[input_1, input_2], outputs=output)
    return siamese