import cv2
import numpy as np
import os
import math
import pickle # Used to save and load the AI's memory instantly

try:
    import face_recognition
except ImportError:
    print("[!] Fatal Error: 'face_recognition' library not found.")

class FaceRecognizer:
    """
    Module 3: The Identity & Security Engine (Enterprise Edition)
    Features HUD UI, Confidence Matrix, Mass Folder Training, Anti-Lag, and Memory Persistence.
    """
    def __init__(self, strictness=0.5, fast_mode_skip=2):
        print("[Aelion Vision] Booting Advanced Identity Engine...")
        self.known_encodings = []
        self.known_names = []
        self.strictness = strictness 
        
        # Anti-Lag Variables
        self.fast_mode_skip = fast_mode_skip
        self.frame_count = 0
        self.last_locations = []
        self.last_names = []
        self.last_confidences = []

    def train(self, name, image_path):
        """Teaches the AI a single face."""
        try:
            image = face_recognition.load_image_file(image_path)
            encodings = face_recognition.face_encodings(image)
            
            if len(encodings) > 0:
                self.known_encodings.append(encodings[0])
                self.known_names.append(name)
                print(f"    [+] Successfully initialized identity: {name}")
            else:
                print(f"    [!] Failed: No face detected in {image_path}")
        except Exception as e:
            print(f"    [!] Failed to read {image_path}. Error: {e}")

    def train_folder(self, folder_path):
        """MASS TRAINING: Scans a folder and learns every face inside it."""
        print(f"[Aelion Vision] Scanning secure folder: '{folder_path}'...")
        if not os.path.exists(folder_path):
            print(f"    [!] Folder '{folder_path}' does not exist.")
            return

        for filename in os.listdir(folder_path):
            if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
                name = os.path.splitext(filename)[0]
                filepath = os.path.join(folder_path, filename)
                self.train(name, filepath)

    # --- ENTERPRISE UPGRADE: DATA PERSISTENCE ---
    def save_brain(self, filepath="aelion_faces.pkl"):
        """Saves the trained neural encodings to the hard drive for instant booting later."""
        with open(filepath, 'wb') as f:
            pickle.dump({'encodings': self.known_encodings, 'names': self.known_names}, f)
        print(f"[Aelion Vision] Neural Brain saved successfully to {filepath}")

    def load_brain(self, filepath="aelion_faces.pkl"):
        """Instantly loads a pre-trained brain without needing the original photos."""
        if os.path.exists(filepath):
            with open(filepath, 'rb') as f:
                data = pickle.load(f)
                self.known_encodings = data['encodings']
                self.known_names = data['names']
            print(f"[Aelion Vision] Brain loaded! Knows {len(self.known_names)} identities.")
            return True
        else:
            print(f"[!] Brain file {filepath} not found. You need to train it first.")
            return False

    def _face_distance_to_conf(self, face_distance, face_match_threshold=0.6):
        """Translates neural distance into a readable confidence percentage."""
        if face_distance > face_match_threshold:
            range_val = (1.0 - face_match_threshold)
            linear_val = (1.0 - face_distance) / (range_val * 2.0)
            return linear_val
        else:
            range_val = face_match_threshold
            linear_val = 1.0 - (face_distance / (range_val * 2.0))
            return linear_val + ((1.0 - linear_val) * math.pow((linear_val - 0.5) * 2, 0.2))

    def _draw_hud(self, frame, left, top, right, bottom, name, color, confidence=None):
        """Draws a high-tech Sci-Fi targeting bracket and crosshair."""
        length = 20
        thickness = 4
        
        # 1. Corner Brackets
        cv2.line(frame, (left, top), (left + length, top), color, thickness)
        cv2.line(frame, (left, top), (left, top + length), color, thickness)
        cv2.line(frame, (right, top), (right - length, top), color, thickness)
        cv2.line(frame, (right, top), (right, top + length), color, thickness)
        cv2.line(frame, (left, bottom), (left + length, bottom), color, thickness)
        cv2.line(frame, (left, bottom), (left, bottom - length), color, thickness)
        cv2.line(frame, (right, bottom), (right - length, bottom), color, thickness)
        cv2.line(frame, (right, bottom), (right, bottom - length), color, thickness)

        # 2. Inner Crosshair
        cx, cy = (left + right) // 2, (top + bottom) // 2
        cv2.line(frame, (cx - 10, cy), (cx + 10, cy), color, 1)
        cv2.line(frame, (cx, cy - 10), (cx, cy + 10), color, 1)

        # 3. ID Tag
        label = f"{name} ({confidence}%)" if confidence else name
        cv2.rectangle(frame, (left, bottom + 10), (right, bottom + 40), color, cv2.FILLED)
        cv2.putText(frame, label, (left + 5, bottom + 32), cv2.FONT_HERSHEY_DUPLEX, 0.6, (255, 255, 255), 1)

    def recognize(self, frame, draw=True):
        """Scans the frame using Anti-Lag Frame Skipping logic."""
        small_frame = cv2.resize(frame, (0, 0), fx=0.25, fy=0.25)
        rgb_small_frame = cv2.cvtColor(small_frame, cv2.COLOR_BGR2RGB)
        
        self.frame_count += 1

        # ENTERPRISE UPGRADE: Only run the heavy math every N frames
        if self.frame_count % self.fast_mode_skip == 0:
            self.last_locations = face_recognition.face_locations(rgb_small_frame)
            face_encodings = face_recognition.face_encodings(rgb_small_frame, self.last_locations)
            
            self.last_names = []
            self.last_confidences = []

            for face_encoding in face_encodings:
                matches = face_recognition.compare_faces(self.known_encodings, face_encoding, tolerance=self.strictness)
                name = "UNKNOWN INTRUDER"
                confidence_pct = None

                if True in matches:
                    face_distances = face_recognition.face_distance(self.known_encodings, face_encoding)
                    best_match_index = np.argmin(face_distances)
                    
                    if matches[best_match_index]:
                        name = self.known_names[best_match_index]
                        conf_float = self._face_distance_to_conf(face_distances[best_match_index])
                        confidence_pct = int(round(conf_float * 100))

                self.last_names.append(name)
                self.last_confidences.append(confidence_pct)

        # Draw the HUD continuously using the cached data
        if draw:
            for (top, right, bottom, left), name, conf in zip(self.last_locations, self.last_names, self.last_confidences):
                top *= 4; right *= 4; bottom *= 4; left *= 4
                color = (0, 255, 0) if name != "UNKNOWN INTRUDER" else (0, 0, 255)
                self._draw_hud(frame, left, top, right, bottom, name, color, conf)

        return self.last_names