import numpy as np
import os
import torch
from facenet_pytorch import MTCNN, InceptionResnetV1
import logging

logger = logging.getLogger(__name__)


class FacialProcessing:
    def __init__(self):
        # Set the cache directory to a writable location
        os.environ['TORCH_HOME'] = '/tmp/.cache/torch'

        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        self.mtcnn = MTCNN(keep_all=True, device=device)
        self.resnet = InceptionResnetV1(pretrained='vggface2').eval().to(device)


    def extract_embeddings_vgg(self, image):
        try:
            # Preprocess the image
            preprocessed_image = self.mtcnn(image)

            if preprocessed_image is None:
                logger.warning(f"No face detected in image")
                return None
            
            # Extract the face embeddings
            embeddings = self.resnet(preprocessed_image.unsqueeze(0)).detach().cpu().numpy().tolist()
            if embeddings:
                return embeddings[0]
            
        except Exception as e:
            logger.error(f"An error occurred while extracting embeddings: {e}")
            return None