# drive_paddy/detection/strategies/cnn_model.py
from src.detection.base_processor import BaseProcessor
import numpy as np
import torch
import torchvision.transforms as transforms
from torchvision.models import efficientnet_b7
import cv2
import dlib
from PIL import Image
import os

class CnnProcessor(BaseProcessor):
    """
    Drowsiness detection using a pre-trained EfficientNet-B7 model.
    """
    def __init__(self, config):
        self.settings = config['cnn_model_settings']
        self.model_path = self.settings['model_path']
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # Initialize dlib for face detection
        self.face_detector = dlib.get_frontal_face_detector()
        
        # Load the model
        self.model = self._load_model()
        
        # Define image transformations
        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])

    def _load_model(self):
        """Loads the EfficientNet-B7 model and custom weights."""
        if not os.path.exists(self.model_path):
            print(f"Error: Model file not found at {self.model_path}")
            print("Please run 'python download_model.py' first.")
            return None
            
        try:
            # Initialize the model structure
            model = efficientnet_b7()
            # Modify the final classifier layer to match the number of output classes (e.g., 2: drowsy, not_drowsy)
            num_ftrs = model.classifier[1].in_features
            model.classifier[1] = torch.nn.Linear(num_ftrs, 2) # Assuming 2 output classes

            # Load the saved weights
            model.load_state_dict(torch.load(self.model_path, map_location=self.device))
            model.to(self.device)
            model.eval() # Set the model to evaluation mode
            print(f"CNN Model '{self.model_path}' loaded successfully on {self.device}.")
            return model
        except Exception as e:
            print(f"Error loading CNN model: {e}")
            return None

    def process_frame(self, frame):
        """
        Processes a frame to detect drowsiness using the CNN model.
        """
        if self.model is None:
            return frame, {"cnn_prediction": False}

        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = self.face_detector(gray)
        is_drowsy_prediction = False

        for face in faces:
            x1, y1, x2, y2 = face.left(), face.top(), face.right(), face.bottom()
            
            # Crop the face from the frame
            face_crop = frame[y1:y2, x1:x2]
            
            # Ensure the crop is valid before processing
            if face_crop.size == 0:
                continue
                
            # Convert to PIL Image and apply transformations
            pil_image = Image.fromarray(cv2.cvtColor(face_crop, cv2.COLOR_BGR2RGB))
            image_tensor = self.transform(pil_image).unsqueeze(0).to(self.device)
            
            # Perform inference
            with torch.no_grad():
                outputs = self.model(image_tensor)
                _, preds = torch.max(outputs, 1)
                # Assuming class 1 is 'drowsy' and class 0 is 'not_drowsy'
                print(preds)
                if preds.item() == 1:
                    is_drowsy_prediction = True

            # Draw bounding box for visualization
            cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 255, 0), 2)
            label = "Drowsy" if is_drowsy_prediction else "Awake"
            cv2.putText(frame, f"CNN: {label}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2)
            
            # Process only the first detected face
            break
            
        return frame, {"cnn_prediction": is_drowsy_prediction}