import cv2
import os
import numpy as np

# Initialize the face detector
face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')

# Font for displaying the user ID on the image
font = cv2.FONT_HERSHEY_SIMPLEX

# Function to recognize the user from the registered faces
def recognize_user():
    # Load all registered user folders
    registered_users = os.listdir("registered_faces")
    
    # Initialize webcam for live detection
    cap = cv2.VideoCapture(0)

    # Prompt the user to ensure they can see themselves
    print("Starting face recognition. Please make sure you can see your face on the screen.")
    
    # Capture a few frames to allow the user to see themselves
    for _ in range(30):  # Show a few frames (about 1 second) to the user
        ret, frame = cap.read()
        if not ret:
            print("Failed to grab frame!")
            break
        cv2.putText(frame, "Look at the camera", (50, 50), font, 1, (255, 255, 255), 2, cv2.LINE_AA)
        cv2.imshow("Face Recognition", frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):  # Allow user to exit if they press 'q'
            break

    print("Proceeding with face recognition...")

    while True:
        ret, frame = cap.read()
        if not ret:
            print("Failed to grab frame!")
            break

        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))

        for (x, y, w, h) in faces:
            face = gray[y:y + h, x:x + w]

            # Compare the live face with each registered user face
            recognized_user_id = None
            for user_id in registered_users:
                user_folder = f"registered_faces/{user_id}"
                if not os.path.exists(user_folder):
                    continue

                user_faces = []
                for filename in os.listdir(user_folder):
                    if filename.endswith(".jpg"):
                        img = cv2.imread(f"{user_folder}/{filename}")
                        user_faces.append(cv2.cvtColor(img, cv2.COLOR_BGR2GRAY))

                for stored_face in user_faces:
                    res = cv2.matchTemplate(face, stored_face, cv2.TM_CCOEFF_NORMED)
                    threshold = 0.7
                    if np.max(res) > threshold:
                        recognized_user_id = user_id
                        break

                if recognized_user_id:
                    break

            if recognized_user_id:
                cv2.putText(frame, f"User ID: {recognized_user_id} Recognized", (x, y - 10), font, 1, (0, 255, 0), 2, cv2.LINE_AA)
                cv2.rectangle(frame, (x, y), (x + w, y + h), (255, 0, 0), 2)
            else:
                cv2.putText(frame, "Unknown User", (x, y - 10), font, 1, (0, 0, 255), 2, cv2.LINE_AA)

        cv2.imshow("Face Recognition", frame)

        # Exit if 'q' is pressed
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()

# Run face recognition and user ID detection
recognize_user()