import gradio as gr
import numpy as np
import cv2
import plotly.graph_objects as go
from sklearn.metrics.pairwise import cosine_distances

# =========================
# FACE SETUP
# =========================

face_cascade = cv2.CascadeClassifier(
    cv2.data.haarcascades + "haarcascade_frontalface_default.xml"
)

def extract_face(image):
    if image is None:
        return None, None

    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    faces = face_cascade.detectMultiScale(gray, 1.3, 5)

    if len(faces) == 0:
        return image, None

    x, y, w, h = faces[0]
    cv2.rectangle(image, (x,y), (x+w,y+h), (0,255,0), 2)
    face = cv2.resize(gray[y:y+h, x:x+w], (64,64))
    return image, face.flatten()

def embed(face):
    vec = face @ np.random.randn(face.shape[0], 128)
    return vec / np.linalg.norm(vec)

# =========================
# STEP FUNCTIONS
# =========================

def step_detect(img):
    img, _ = extract_face(img)
    return img

def step_enroll(img):
    img, face = extract_face(img)
    if face is None:
        return img, None, "❌ No face detected"

    emb = embed(face)
    return img, emb, "✅ Face enrolled"

def step_verify(img, stored):
    img, face = extract_face(img)
    if face is None or stored is None:
        return img, "❌ Missing data", None

    live = embed(face)
    dist = cosine_distances([stored], [live])[0][0]
    status = "🔓 UNLOCKED" if dist < 0.35 else "🔒 DENIED"

    fig = go.Figure(go.Indicator(
        mode="gauge+number",
        value=dist,
        title={"text": status},
        gauge={"axis": {"range": [0, 1]}}
    ))

    return img, f"Distance: {dist:.3f}", fig

# =========================
# UI
# =========================

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    page = gr.State(0)
    stored_embedding = gr.State()

    gr.Markdown("# 🔐 Face Unlock — How Mobile Face ID Works")

    with gr.Row():
        back = gr.Button("⬅ Back")
        nextb = gr.Button("Next ➡")

    # ---------- PAGE 0 ----------
    page0 = gr.Column(visible=True)
    with page0:
        gr.Markdown("## 📸 Face Detection")
        cam0 = gr.Image(sources=["webcam"], type="numpy")
        out0 = gr.Image()
        gr.Button("Detect Face").click(step_detect, cam0, out0)

    # ---------- PAGE 1 ----------
    page1 = gr.Column(visible=False)
    with page1:
        gr.Markdown("## 🧠 Face Enrollment")
        cam1 = gr.Image(sources=["webcam"], type="numpy")
        out1 = gr.Image()
        msg1 = gr.Markdown()
        gr.Button("Enroll Face").click(
            step_enroll,
            cam1,
            [out1, stored_embedding, msg1]
        )

    # ---------- PAGE 2 ----------
    page2 = gr.Column(visible=False)
    with page2:
        gr.Markdown("## 🔓 Face Verification")
        cam2 = gr.Image(sources=["webcam"], type="numpy")
        out2 = gr.Image()
        msg2 = gr.Markdown()
        gauge = gr.Plot()
        gr.Button("Verify").click(
            step_verify,
            [cam2, stored_embedding],
            [out2, msg2, gauge]
        )

    # =========================
    # NAVIGATION
    # =========================

    def navigate(p, step):
        p = min(2, max(0, p + step))
        return (
            p,
            gr.update(visible=p == 0),
            gr.update(visible=p == 1),
            gr.update(visible=p == 2)
        )

    back.click(navigate, [page, gr.State(-1)], [page, page0, page1, page2])
    nextb.click(navigate, [page, gr.State(1)], [page, page0, page1, page2])

demo.launch()