import traceback
import numpy as np
import torch
import gradio as gr
import timm
import torch.nn.functional as F


MODEL_NAME = "hf_hub:gaunernst/vit_tiny_patch8_112.arcface_ms1mv3"

model = timm.create_model(
    MODEL_NAME,
    pretrained=True
)

model.eval()


def image_to_tensor(image):
    image = image.convert("RGB")
    image = image.resize((112, 112))

    img = np.array(image).astype(np.float32) / 255.0

    img = (img - 0.5) / 0.5

    img = np.transpose(img, (2, 0, 1))

    tensor = torch.tensor(img).unsqueeze(0)

    return tensor


def get_embeddings(image):
    if image is None:
        raise ValueError("Please upload an image.")

    tensor = image_to_tensor(image)

    with torch.no_grad():
        embedding = model(tensor)

    embedding = F.normalize(embedding, p=2, dim=1)

    return embedding


def compare_faces(image1, image2):
    try:
        face1 = get_embeddings(image1)
        face2 = get_embeddings(image2)

        similarity = F.cosine_similarity(face1, face2).item()
        similarity = round(similarity, 4)

        if similarity >= 0.6:
            result = "THIS IS THE SAME PERSON"
        else:
            result = "FACE DOES NOT MATCH"

        return result, similarity

    except Exception:
        return traceback.format_exc(), 0


demo = gr.Interface(
    fn=compare_faces,
    inputs=[
        gr.Image(type="pil", label="Upload your first face"),
        gr.Image(type="pil", label="Upload your second face")
    ],
    outputs=[
        gr.Textbox(label="Result"),
        gr.Number(label="Similarity")
    ],
    title="Face Recognition App",
    description="Upload two clear cropped face images and compare them."
)

demo.launch()