import os
import time
import cv2
import numpy as np
import torch
from ultralytics import YOLO
from PIL import Image as PILImage
from keras_facenet import FaceNet
from transformers import pipeline
import gradio as gr
from datetime import datetime, timedelta
import gc

# -----------------------------
# Device Setup
# -----------------------------
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {DEVICE}")

# -----------------------------
# Load YOLOv8 Face Model
# -----------------------------
MODEL_PATH = "yolov8n-face.pt"  # put this file in your Space repository
face_model = YOLO(MODEL_PATH).to(DEVICE)

# -----------------------------
# Load FaceNet Embedder
# -----------------------------
embedder = FaceNet()  # CPU mode

# -----------------------------
# Load HuggingFace Age & Gender Models
# -----------------------------
age_model = pipeline(
    "image-classification",
    model="prithivMLmods/Age-Classification-SigLIP2",
    device=-1
)
gender_model = pipeline(
    "image-classification",
    model="dima806/fairface_gender_image_detection",
    device=-1
)

# -----------------------------
# Face DB
# -----------------------------
FACE_DB = []
NEXT_ID = 1

def clean_gpu():
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    gc.collect()

def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

# -----------------------------
# Main Inference Function
# -----------------------------
def process_image(image):
    global NEXT_ID, FACE_DB
    start_time = time.time()

    rgb_img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # Detect faces
    results = face_model(rgb_img, verbose=False)
    boxes = results[0].boxes.xyxy.cpu().numpy().astype(int)

    now = datetime.now()
    FACE_DB = [f for f in FACE_DB if now - f["time"] <= timedelta(hours=1)]

    faces = []
    for (x1, y1, x2, y2) in boxes:
        face_crop = rgb_img[y1:y2, x1:x2]
        if face_crop.size == 0:
            continue

        face_embedding = embedder.embeddings([face_crop])[0]

        assigned_id = None
        age_pred, gender_pred = "unknown", "unknown"

        # Compare with DB
        if FACE_DB:
            similarities = [cosine_similarity(face_embedding, entry["embedding"]) for entry in FACE_DB]
            best_match_index = int(np.argmax(similarities))
            best_score = similarities[best_match_index]

            if best_score > 0.6:
                assigned_id = FACE_DB[best_match_index]["id"]
                FACE_DB[best_match_index]["time"] = now
                FACE_DB[best_match_index]["seen_count"] += 1
                age_pred = FACE_DB[best_match_index]["age"]
                gender_pred = FACE_DB[best_match_index]["gender"]

        if assigned_id is None:
            assigned_id = NEXT_ID
            face_pil = PILImage.fromarray(face_crop)
            try:
                age_pred = age_model(face_pil)[0]["label"]
                gender_pred = gender_model(face_pil)[0]["label"]
            except Exception:
                age_pred, gender_pred = "unknown", "unknown"

            FACE_DB.append({
                "id": assigned_id,
                "embedding": face_embedding,
                "time": now,
                "seen_count": 1,
                "age": age_pred,
                "gender": gender_pred
            })
            NEXT_ID += 1

        faces.append({
            "id": assigned_id,
            "age": age_pred,
            "gender": gender_pred,
            "box": [int(x1), int(y1), int(x2), int(y2)]
        })

    total_time = round(time.time() - start_time, 3)
    clean_gpu()

    return {"faces": faces, "processing_time_sec": total_time}

# -----------------------------
# Gradio UI
# -----------------------------
iface = gr.Interface(
    fn=process_image,
    inputs=gr.Image(type="numpy"),
    outputs="json",
    title="Face Detection + Age/Gender"
)

if __name__ == "__main__":
    iface.launch()