import gradio as gr
import os
import numpy as np
from PIL import Image
from sklearn.metrics.pairwise import cosine_similarity
import insightface
from insightface.app import FaceAnalysis

# -----------------------------
# Setup directories
# -----------------------------
os.makedirs("registered_faces", exist_ok=True)
os.makedirs("embeddings", exist_ok=True)

# -----------------------------
# Load ArcFace Model (InsightFace)
# -----------------------------
app = FaceAnalysis(name="buffalo_l")  # ArcFace best model set
app.prepare(ctx_id=0, det_size=(640, 640))  # ctx_id=0 uses GPU if available, else CPU

# -----------------------------
# Helper: Generate Embedding
# -----------------------------
def get_embedding(image):
    img = np.array(image)
    faces = app.get(img)

    if len(faces) == 0:
        return None, "❌ No face detected. Try another image."
    
    # Take first detected face
    embedding = faces[0].embedding
    return embedding, None

# -----------------------------
# Register New Face
# -----------------------------
def register_face(name, image):
    if not name:
        return "⚠️ Please enter a name."

    embedding, error = get_embedding(image)
    if embedding is None:
        return error

    # Save image & embedding
    image.save(f"registered_faces/{name}.jpg")
    np.save(f"embeddings/{name}.npy", embedding)

    return f"✅ Registered {name} successfully!"

# -----------------------------
# Recognize Face
# -----------------------------
def recognize_face(image):
    embedding, error = get_embedding(image)
    if embedding is None:
        return error

    best_match = None
    highest_score = 0

    for file in os.listdir("embeddings"):
        if file.endswith(".npy"):
            saved_emb = np.load(os.path.join("embeddings", file))
            score = cosine_similarity([embedding], [saved_emb])[0][0]

            if score > highest_score:
                highest_score = score
                best_match = file.replace(".npy", "")

    # Threshold decision
    if best_match and highest_score > 0.60:
        return f"🟢 Match Found: **{best_match}** (Similarity: {highest_score:.2f})"
    
    return f"🔴 No match found. Best score = {highest_score:.2f}"

# -----------------------------
# Gradio UI
# -----------------------------
with gr.Blocks(title="Face Recognition Attendance System") as demo:
    
    gr.Markdown("## 🧠 Facial Recognition System (ArcFace Based)")
    gr.Markdown("Upload a face to register or recognize.")

    with gr.Tab("📌 Register Employee"):
        name_input = gr.Textbox(label="Employee Name")
        reg_image = gr.Image(label="Upload Face", type="pil")
        reg_button = gr.Button("Register")
        reg_output = gr.Textbox(label="Status")
        reg_button.click(register_face, inputs=[name_input, reg_image], outputs=reg_output)

    with gr.Tab("🔍 Recognize Face"):
        recog_image = gr.Image(label="Upload Face", type="pil")
        recog_button = gr.Button("Recognize")
        recog_output = gr.Textbox(label="Result")
        recog_button.click(recognize_face, inputs=recog_image, outputs=recog_output)

demo.launch(server_name="0.0.0.0", server_port=7860)