import os
import json

os.environ["DEEPFACE_HOME"] = "."

import pyzipper
import numpy as np
import gradio as gr
from annoy import AnnoyIndex
from deepface.commons import functions
from deepface.basemodels import Facenet512


# load the face model
model = Facenet512.loadModel()

input_shape_x, input_shape_y = functions.find_input_shape(model)

index = AnnoyIndex(512, "euclidean")
index.load(f"face.db")

ANNOY_INDEX = json.load(open(f"face.json"))

with pyzipper.AESZipFile('persons.zip') as zf:
    zf.setpassword(b"4321ecafhsats"[::-1])
    PERFORMER_DB = json.loads(zf.read('performers.json'))


def predict(image, threshold=20.0, results=3):
    image_array = np.array(image)

    img = functions.preprocess_face(
        img=image_array,
        target_size=(input_shape_x, input_shape_y),
        enforce_detection=True,
        detector_backend="retinaface",
        align=True,
    )

    img = functions.normalize_input(img, normalization="Facenet2018")
    face = model.predict(img)[0].tolist()
    ids, distances = index.get_nns_by_vector(
        face, 50, search_k=10000, include_distances=True
    )

    persons = {}
    for p, distance in zip(ids, distances):
        id = ANNOY_INDEX[p]
        if id in persons:
            persons[id]["hits"] += 1
            persons[id]["distance"] -= 0.5
            continue

        persons[id] = {
            "id": id,
            "distance": round(distance, 2),
            "hits": 1,
        }

        if id in PERFORMER_DB:
            persons[id].update(PERFORMER_DB.get(id))

    persons = sorted(persons.values(), key=lambda x: x["distance"])
    persons = [p for p in persons if p["distance"] < threshold]
    return persons[:results]


gr.Interface(
    fn=predict,
    inputs=[
        gr.components.Image(),
        gr.components.Slider(label="threshold",minimum=0.0, maximum=30.0, value=20.0),
        gr.components.Slider(label="results", minimum=0, maximum=10, value=3),
    ],
    outputs=gr.outputs.JSON(label=""),
    title="Who is in the photo?",
    description="Upload an image of a person and we'll tell you who it is.",
).launch(enable_queue=True)