import numpy as np
from typing import Tuple

from deepface import DeepFace
from deepface.modules import modeling, preprocessing

class ArcFaceRecognition:
    def __init__(self):
        """Initialize ArcFace recognition system."""
        pass
    
    def _preprocess_face_batch(self, faces: np.ndarray, target_size: Tuple[int, int], normalization: str) -> np.ndarray:
        """Preprocess a batch of face images for model inference"""
        batch_size = faces.shape[0]
        processed_faces = []
        
        for i in range(batch_size):
            face = faces[i]
            # Convert RGB to BGR (DeepFace expects BGR)
            face = face[:, :, ::-1]
            
            # Resize to model input size
            resized = preprocessing.resize_image(face, target_size)
            
            # Normalize
            normalized = preprocessing.normalize_input(resized, normalization)
            
            processed_faces.append(normalized)
        
        # Stack into batch and remove the extra dimension added by resize_image
        batch = np.vstack(processed_faces)
        return batch
    
    def get_face_embeddings_batch(self, faces: np.ndarray) -> np.ndarray:
        """Get face embeddings for a batch of images efficiently
        
        Args:
            faces: np.ndarray of shape (batch_size, height, width, channels)
            
        Returns:
            ArcFace embeddings as np.ndarray
        """
        # Load ArcFace model (cached by DeepFace)
        arcface_model = modeling.build_model(task="facial_recognition", model_name="ArcFace")
        
        # Preprocess faces
        arcface_batch = self._preprocess_face_batch(faces, arcface_model.input_shape, "ArcFace")
        
        # Get embeddings using direct model inference
        arcface_embeddings = arcface_model.model(arcface_batch, training=False).numpy()
        
        return arcface_embeddings

def extract_faces(image):
    """Extract faces from an image using DeepFace"""
    return DeepFace.extract_faces(image, detector_backend="yolov8")