import cv2
import os
import insightface
import numpy as np
from collections import deque

# Initialize face detector and recognizer
model = insightface.app.FaceAnalysis(allowed_modules=['detection', 'recognition'])
model.prepare(ctx_id=0, det_size=(640, 640))  # Use ctx_id=-1 for CPU

def normalize(v):
    return v / np.linalg.norm(v)

def cosine_similarity(a, b):
    return np.dot(a, b)

# Load known faces
known_embs = []
names = []
for fname in os.listdir("images"):
    if fname.lower().endswith(('.jpg', '.png')):
        img = cv2.imread(os.path.join("images", fname))
        faces = model.get(img)
        if faces:
            emb = normalize(faces[0].embedding)
            known_embs.append(emb)
            names.append(os.path.splitext(fname)[0])
            print(f"Loaded {fname}")
        else:
            print(f"No face in {fname}")

# Smoothing buffers for each face
face_buffers = {}  # key: face_id, value: deque of embeddings

# Start video
cap = cv2.VideoCapture(0)
face_id_counter = 0

while True:
    ret, frame = cap.read()
    if not ret:
        break

    faces = model.get(frame)
    current_buffers = {}

    for i, face in enumerate(faces):
        x1, y1, x2, y2 = face.bbox.astype(int)
        emb = normalize(face.embedding)

        # Use a temporary ID for each face based on bbox location
        face_id = f"{x1}-{y1}-{x2}-{y2}"
        if face_id not in face_buffers:
            face_buffers[face_id] = deque(maxlen=5)

        face_buffers[face_id].append(emb)
        current_buffers[face_id] = face_buffers[face_id]  # mark as active this frame

        # Smooth embedding
        avg_emb = normalize(np.mean(face_buffers[face_id], axis=0))

        # Find best match
        sims = [cosine_similarity(avg_emb, known) for known in known_embs]
        max_idx = np.argmax(sims)
        name = "Unknown"
        if sims[max_idx] > 0.5:
            name = names[max_idx]

        # Draw result
        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 200, 0), 2)
        cv2.putText(frame, name, (x1, y1 - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 200, 0), 2)

    # Remove stale buffers
    face_buffers = {fid: buf for fid, buf in face_buffers.items() if fid in current_buffers}

    # Show frame
    cv2.imshow("InsightFace Multi-Face Recognition", frame)
    if cv2.waitKey(1) == 27:  # ESC key
        break

cap.release()
cv2.destroyAllWindows()