import spaces
import gradio as gr
import numpy as np
import tensorflow as tf
from tensorflow import keras
from model import build_siamese
import cv2
import os

# Suppress oneDNN warnings (optional)
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'

# Define custom functions
def euclidean_distance(vectors):
    x, y = vectors
    sum_square = tf.reduce_sum(tf.square(x - y), axis=1, keepdims=True)
    return tf.sqrt(tf.maximum(sum_square, tf.keras.backend.epsilon()))

def contrastive_loss(y_true, y_pred, margin=1.0):
    square_pred = tf.square(y_pred)
    margin_square = tf.square(tf.maximum(margin - y_pred, 0))
    return tf.reduce_mean(y_true * square_pred + (1 - y_true) * margin_square)

# Load model with custom objects

siamese = build_siamese()

siamese.load_weights("my_siamese.keras", skip_mismatch=True)

#try:
#    siamese = keras.models.load_model(
#        "my_siamese.keras",
#        custom_objects={
#            "euclidean_distance": euclidean_distance,
#            "contrastive_loss": contrastive_loss
#        },
#        compile=False
#    )
#except Exception as e:
#    print(f"Error loading model: {e}")
#    raise

# Load the model
# try:
#     siamese = keras.models.load_model("my_siamese.keras", compile=False)
# except Exception as e:
#     print(f"Error loading model: {e}")
#     raise

# Load stored images
def load_stored_images():
    try:
        stored_images = [np.zeros((256, 256, 3)) for _ in range(3)]  # Mock images
        return np.array([np.expand_dims(img, axis=0) for img in stored_images])
    except Exception as e:
        print(f"Error loading stored images: {e}")
        raise

stored_imgs = load_stored_images()

# Inference function
@spaces.GPU
def check_membership(uploaded_image):
    try:
        cv_img = cv2.cvtColor(np.array(uploaded_image), cv2.COLOR_RGB2BGR)
        image = cv2.resize(cv_img, (256, 256))

        uploaded_image = np.expand_dims(image, axis=0).astype("float32")

        predictions = [siamese.predict([uploaded_image, img])[0][0] for img in stored_imgs]

        if predictions[0] < 0.5:
            return "You are welcome Munzali"
        elif predictions[1] < 0.5:
            return "You are welcome Ahmad"
        elif predictions[2] < 0.5:
            return "You are welcome Usama"
        else:
            return "You are not a member"
    except Exception as e:
        return f"Error: {e}"

# Gradio interface

def app():
    with gr.Blocks():
        with gr.Row():
            with gr.Column():
                image = gr.Image(type="pil", label="Image")
                siamese_infer = gr.Button(value="Search the image")

                # Get files in the current working directory ending with .jpg or .jpeg
                example_list = [
                    [filename] 
                    for filename in os.listdir(".") 
                    if filename.lower().endswith(('.jpg', '.jpeg'))
                ]
                
                
                # Gradio Examples component
                gr.Examples(
                    examples=example_list,
                    inputs=[image]  # assuming `image` is your gr.Image input component
                )

    
            with gr.Column():
                output_image = gr.Textbox( label="Resulted Search", lines=1)
        siamese_infer.click(
            fn=check_membership,
            inputs=image,
            outputs=output_image
        )
#[gr.Textbox(label="greeting", lines=1)], gr.Image(image_mode="RGB")
gradio_app = gr.Blocks()

with gradio_app:
    gr.HTML(
        """
    <h1 style='text-align: center'>
    SIAMESE: Real-Time End-to-End Face Security System
    </h1>
    """)
    gr.HTML(
        """
        <h3 style='text-align: center'>
        Follow me for more!
        <a href='munzalialhassan.gumel@gmail.co' target='_blank'>Email</a> |
        <a href='https://github.com/ExplorerGumel' target='_blank'>Github</a> | 
        <a href='https://linkedin.com/in/munzali-alhassan' target='_blank'>Linkedin</a>  | 
<!--        <a href='https://www.huggingface.co/kadirnar/' target='_blank'>HuggingFace</a> -->
        </h3>
        """
    )
 
    with gr.Row():
        with gr.Column():
            app()
gradio_app.launch(debug=True, share=True)
# iface = gr.Interface(
#     fn=check_membership,
#     inputs=gr.Image(image_mode="RGB"),
#     outputs="text",
#     title="Siamese Network Membership Check",
#     description="A security check model used to give access to the assign individual and blocks the rest"
# )

# iface.launch()
#share=True, auth=("username", "password")