import gradio as gr

import numpy as np
from deepface import DeepFace
from pymongo.mongo_client import MongoClient
import cv2

credentials = "jamshaid:jamshaid19gh"

uri = f"mongodb+srv://{credentials}@cluster0.uimyui3.mongodb.net/?retryWrites=true&w=majority"
client = MongoClient(uri)
db = client["Face_identification"]
identities_collection = db["face_identities"]
model_name="Facenet"
debug=False

def save_identity(image , name):
  try:
    embeddings = DeepFace.represent(image , model_name=model_name , detector_backend = "retinaface")
    embeddings = embeddings[0]

    identity = {"embeddings":embeddings["embedding"] , "name" : name }

    result = identities_collection.insert_one(identity)

    return f"{name} stored in database successfully.It is recommended to add 2 or 3 high quality images for one person"
  except Exception as error:
    return str(error)

def findCosineDistance(source_representation, test_representation):
    a = np.matmul(np.transpose(source_representation), test_representation)
    b = np.sum(np.multiply(source_representation, source_representation))
    c = np.sum(np.multiply(test_representation, test_representation))
    return 1 - (a / (np.sqrt(b) * np.sqrt(c)))

def findThreshold(model_name, distance_metric):

    base_threshold = {"cosine": 0.40, "euclidean": 0.55, "euclidean_l2": 0.75}

    thresholds = {
        "VGG-Face": {"cosine": 0.40, "euclidean": 0.60, "euclidean_l2": 0.86},
        "Facenet": {"cosine": 0.40, "euclidean": 10, "euclidean_l2": 0.80},
        "Facenet512": {"cosine": 0.30, "euclidean": 23.56, "euclidean_l2": 1.04},
        "ArcFace": {"cosine": 0.68, "euclidean": 4.15, "euclidean_l2": 1.13},
        "Dlib": {"cosine": 0.07, "euclidean": 0.6, "euclidean_l2": 0.4},
        "SFace": {"cosine": 0.593, "euclidean": 10.734, "euclidean_l2": 1.055},
        "OpenFace": {"cosine": 0.10, "euclidean": 0.55, "euclidean_l2": 0.55},
        "DeepFace": {"cosine": 0.23, "euclidean": 64, "euclidean_l2": 0.64},
        "DeepID": {"cosine": 0.015, "euclidean": 45, "euclidean_l2": 0.17},
    }

    threshold = thresholds.get(model_name, base_threshold).get(distance_metric, 0.4)

    return threshold

threshold = findThreshold(model_name , "cosine")

def predict_image(image):
    original_image = np.copy(image)
    if debug: 
      print("1")
    # getting face embeddings from database
    results = identities_collection.find()
    faces = [dict(result) for result in results]

    if debug:
      print("2")

    # generate face embeddings for all detected faces in image
    target_embedding_array = DeepFace.represent(
        img_path=image,
        model_name=model_name,
        detector_backend = "retinaface"
    )
    identities = []
    # for each face compare its embeddings with all face embeddings in database
    for target_embedding_obj in target_embedding_array:
        target_embedding = target_embedding_obj["embedding"]

        if debug:
          print("4")
        # compare the face embedding with all other faces

        name = "Unknown"
        for face in faces:
          distance = findCosineDistance(face["embeddings"], target_embedding)
          if distance <= threshold:
            name = face["name"]
            break

        if debug:
          print("5")

        identities.append({"name":name , "facial_area":target_embedding_obj["facial_area"]})  

    output_img = np.copy(original_image)
    for identity in identities:    
        # Draw the rectangle on the image
        x = identity["facial_area"]["x"]
        y = identity["facial_area"]["y"]
        w = identity["facial_area"]["w"]
        h = identity["facial_area"]["h"]
        cv2.rectangle(output_img, (x,y), (x+w,y+h), (0, 0, 255), 1)

        # Define the text position
        text_position = (x, y+h+5)

        # Add the text to the image
        cv2.putText(output_img ,identity["name"], text_position, cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255,0 ), 1)
    return output_img


# image_input = gr.inputs.Image(shape=(160,160))


# Create the Gradio interface
# gr.Interface(fn=predict_image, inputs=image_input, outputs=label_output).launch()


# Create Gradio interfaces for input and output
image_input = gr.inputs.Image(shape=(160, 160))
label_input = gr.inputs.Textbox(label="Enter Name")
label_output = gr.outputs.Textbox()


# Create the Gradio interface
interface1 = gr.Interface(
    fn=save_identity,
    inputs=[image_input, label_input],
    outputs=label_output,
    title="Face Identification",
    description="Upload an image, enter the person name and store the person in database",
)


# Create Gradio interfaces for image input and output
image_input2 = gr.inputs.Image(shape=(160,160))
output_image = gr.outputs.Image(type="numpy")
# output_image = gr.outputs.Textbox()

# Create the Gradio interface for image input and output
interface2 = gr.Interface(
    fn=predict_image,
    inputs=image_input2,
    outputs=output_image,
    title="Face Identification",
    description="Upload an image and get the identity of person",
)

# Create the Gradio interface with two tabs
# interface = gr.Interface(title="Face identification App")
# interface.add_view(interface1, "Save", "Add new person")
# interface.add_view(interface2, "Predict", "Get identity of person")

gr.TabbedInterface(
    [interface2 , interface1],
    tab_names=["Predict Persons","Add new Person"]
).queue().launch()

# Launch the Gradio interface
# interface.launch()