#!/usr/bin/env python3
"""
OpenCV Face Recognition Demo
Advanced face recognition that can learn and identify specific people
"""

import cv2
import numpy as np
import os
import pickle
from datetime import datetime

class FaceRecognitionDemo:
    def __init__(self):
        self.face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
        self.face_recognizer = cv2.face.LBPHFaceRecognizer_create()
        self.known_faces = []
        self.known_names = []
        self.face_data_file = 'face_data.pkl'
        self.load_face_data()
    
    def load_face_data(self):
        """Load previously saved face data"""
        if os.path.exists(self.face_data_file):
            try:
                with open(self.face_data_file, 'rb') as f:
                    data = pickle.load(f)
                    self.known_faces = data['faces']
                    self.known_names = data['names']
                
                if self.known_faces:
                    labels = list(range(len(self.known_faces)))
                    self.face_recognizer.train(self.known_faces, np.array(labels))
                    print(f"Loaded {len(self.known_faces)} known faces")
            except:
                print("Could not load face data")
    
    def save_face_data(self):
        """Save face data to file"""
        data = {
            'faces': self.known_faces,
            'names': self.known_names
        }
        with open(self.face_data_file, 'wb') as f:
            pickle.dump(data, f)
        print("Face data saved!")
    
    def register_face_from_webcam(self, name):
        """Register a new face from webcam"""
        print(f"Registering face for: {name}")
        print("Look at the camera and press SPACE when ready, ESC to cancel")
        
        cap = cv2.VideoCapture(0)
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            faces = self.face_cascade.detectMultiScale(gray, 1.3, 5)
            
            for (x, y, w, h) in faces:
                cv2.rectangle(frame, (x, y), (x+w, y+h), (0, 255, 0), 2)
                cv2.putText(frame, f'Registering: {name}', (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
            
            cv2.putText(frame, 'Press SPACE to capture, ESC to cancel', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2)
            cv2.imshow('Register Face', frame)
            
            key = cv2.waitKey(1) & 0xFF
            if key == 32:  # Space key
                if len(faces) > 0:
                    # Use the first detected face
                    (x, y, w, h) = faces[0]
                    face_roi = gray[y:y+h, x:x+w]
                    face_roi = cv2.resize(face_roi, (100, 100))
                    
                    # Add to known faces
                    self.known_faces.append(face_roi)
                    self.known_names.append(name)
                    
                    # Retrain recognizer
                    labels = list(range(len(self.known_faces)))
                    self.face_recognizer.train(self.known_faces, np.array(labels))
                    
                    self.save_face_data()
                    print(f"✅ Face registered for {name}!")
                    break
                else:
                    print("❌ No face detected. Try again.")
            elif key == 27:  # ESC key
                print("Registration cancelled")
                break
        
        cap.release()
        cv2.destroyAllWindows()
    
    def recognize_faces_webcam(self):
        """Real-time face recognition from webcam"""
        print("Starting face recognition...")
        print("Press 'q' to quit")
        
        cap = cv2.VideoCapture(0)
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            faces = self.face_cascade.detectMultiScale(gray, 1.3, 5)
            
            for (x, y, w, h) in faces:
                face_roi = gray[y:y+h, x:x+w]
                face_roi = cv2.resize(face_roi, (100, 100))
                
                name = "Unknown"
                confidence = 0
                
                if len(self.known_faces) > 0:
                    try:
                        label, conf = self.face_recognizer.predict(face_roi)
                        if conf < 100:  # Confidence threshold
                            name = self.known_names[label]
                            confidence = 100 - conf
                    except:
                        pass
                
                # Draw rectangle and label
                color = (0, 255, 0) if name != "Unknown" else (0, 0, 255)
                cv2.rectangle(frame, (x, y), (x+w, y+h), color, 2)
                
                label_text = f'{name}'
                if name != "Unknown":
                    label_text += f' ({confidence:.1f}%)'
                
                cv2.putText(frame, label_text, (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)
            
            # Display info
            cv2.putText(frame, f'Known faces: {len(self.known_faces)}', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2)
            cv2.putText(frame, 'Press Q to quit', (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2)
            
            cv2.imshow('Face Recognition', frame)
            
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
        
        cap.release()
        cv2.destroyAllWindows()
    
    def list_known_faces(self):
        """List all registered faces"""
        if not self.known_names:
            print("No faces registered yet.")
        else:
            print(f"\n📋 Known faces ({len(self.known_names)}):")
            for i, name in enumerate(self.known_names):
                print(f"  {i+1}. {name}")
    
    def clear_all_faces(self):
        """Clear all registered faces"""
        confirm = input("Are you sure you want to clear all faces? (yes/no): ")
        if confirm.lower() == 'yes':
            self.known_faces = []
            self.known_names = []
            if os.path.exists(self.face_data_file):
                os.remove(self.face_data_file)
            print("✅ All faces cleared!")
        else:
            print("Operation cancelled.")

def main():
    """Main demo function"""
    print("🎯 OpenCV Face Recognition Demo")
    print("=" * 40)
    
    demo = FaceRecognitionDemo()
    
    while True:
        print("\n🔍 Face Recognition Menu:")
        print("1. Register new face")
        print("2. Start face recognition")
        print("3. List known faces")
        print("4. Clear all faces")
        print("5. Exit")
        
        choice = input("\nEnter your choice (1-5): ").strip()
        
        if choice == '1':
            name = input("Enter person's name: ").strip()
            if name:
                demo.register_face_from_webcam(name)
            else:
                print("Please enter a valid name.")
        
        elif choice == '2':
            if demo.known_faces:
                demo.recognize_faces_webcam()
            else:
                print("❌ No faces registered yet. Please register some faces first.")
        
        elif choice == '3':
            demo.list_known_faces()
        
        elif choice == '4':
            demo.clear_all_faces()
        
        elif choice == '5':
            print("Goodbye! 👋")
            break
        
        else:
            print("Invalid choice. Please try again.")

if __name__ == "__main__":
    main()