# 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
from PIL import Image
import os

class CnnProcessor(BaseProcessor):
    """
    Drowsiness detection using a pre-trained EfficientNet-B7 model.
    This version receives face landmarks from another processor instead of using dlib.
    """
    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")
        
        # dlib is no longer needed.
        # self.face_detector = dlib.get_frontal_face_detector()
        
        self.model = self._load_model()
        
        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}")
            return None
            
        try:
            model = efficientnet_b7()
            num_ftrs = model.classifier[1].in_features
            model.classifier[1] = torch.nn.Linear(num_ftrs, 2)
            model.load_state_dict(torch.load(self.model_path, map_location=self.device))
            model.to(self.device)
            model.eval()
            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, face_landmarks=None):
        """
        Processes a frame using the CNN model with pre-supplied landmarks.
        """
        if self.model is None or face_landmarks is None:
            return frame, {"cnn_prediction": False}

        is_drowsy_prediction = False
        h, w, _ = frame.shape
        
        landmarks = face_landmarks[0].landmark
        
        # Calculate bounding box from landmarks
        x_coords = [lm.x * w for lm in landmarks]
        y_coords = [lm.y * h for lm in landmarks]
        x1, y1 = int(min(x_coords)), int(min(y_coords))
        x2, y2 = int(max(x_coords)), int(max(y_coords))

        # Add some padding to the bounding box
        padding = 10
        x1 = max(0, x1 - padding)
        y1 = max(0, y1 - padding)
        x2 = min(w, x2 + padding)
        y2 = min(h, y2 + padding)

        # Crop the face
        face_crop = frame[y1:y2, x1:x2]
        
        if face_crop.size > 0:
            pil_image = Image.fromarray(cv2.cvtColor(face_crop, cv2.COLOR_BGR2RGB))
            image_tensor = self.transform(pil_image).unsqueeze(0).to(self.device)
            
            with torch.no_grad():
                outputs = self.model(image_tensor)
                _, preds = torch.max(outputs, 1)
                if preds.item() == 1: # Assuming class 1 is 'drowsy'
                    is_drowsy_prediction = True

            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)
            
        return frame, {"cnn_prediction": is_drowsy_prediction}