import numpy as np
import cv2

class ThreatLabeler:
    def __init__(self):
        """
        Initialize the threat labeling system for deepfake detection
        """
        # Define threat level thresholds
        self.threat_levels = {
            'low': (0, 30),      # Low threat: 0-30% difference
            'medium': (30, 60),  # Medium threat: 30-60% difference
            'high': (60, 100)    # High threat: 60-100% difference
        }
        
        # Define colors for each threat level (BGR format for OpenCV)
        self.threat_colors = {
            'low': (0, 255, 0),      # Green
            'medium': (0, 165, 255), # Orange
            'high': (0, 0, 255)      # Red
        }
    
    def calculate_threat_level(self, diff_value):
        """
        Calculate threat level based on difference value
        
        Args:
            diff_value: Normalized difference value (0-100)
            
        Returns:
            Threat level string ('low', 'medium', or 'high')
        """
        for level, (min_val, max_val) in self.threat_levels.items():
            if min_val <= diff_value < max_val:
                return level
        return 'high'  # Default to high if outside ranges
    
    def calculate_region_threat(self, diff_image, bbox):
        """
        Calculate the threat level for a specific region
        
        Args:
            diff_image: Difference image (normalized 0-255)
            bbox: Bounding box (x, y, w, h)
            
        Returns:
            Threat level string and average difference value
        """
        x, y, w, h = bbox
        region = diff_image[y:y+h, x:x+w]
        
        # Calculate average difference in the region (normalized to 0-100)
        avg_diff = np.mean(region) / 255 * 100
        
        # Determine threat level
        threat_level = self.calculate_threat_level(avg_diff)
        
        return threat_level, avg_diff
    
    def label_regions(self, image, diff_image, bounding_boxes):
        """
        Label regions with threat levels and colors
        
        Args:
            image: Original image to label
            diff_image: Difference image
            bounding_boxes: List of bounding boxes (x, y, w, h)
            
        Returns:
            Labeled image and list of regions with threat levels
        """
        # Make a copy of the image to avoid modifying the original
        labeled_image = image.copy()
        labeled_regions = []
        
        # Process each bounding box
        for bbox in bounding_boxes:
            x, y, w, h = bbox
            
            # Calculate threat level for this region
            threat_level, avg_diff = self.calculate_region_threat(diff_image, bbox)
            
            # Get color for this threat level
            color = self.threat_colors[threat_level]
            
            # Draw colored rectangle based on threat level
            cv2.rectangle(labeled_image, (x, y), (x + w, y + h), color, 2)
            
            # Add label text with threat level and percentage
            label_text = f"{threat_level.upper()}: {avg_diff:.1f}%"
            cv2.putText(labeled_image, label_text, (x, y - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
            
            # Store labeled region info
            labeled_regions.append({
                'bbox': bbox,
                'threat_level': threat_level,
                'difference_percentage': avg_diff
            })
        
        return labeled_image, labeled_regions
    
    def get_threat_summary(self, labeled_regions):
        """
        Generate a summary of threat levels in the image
        
        Args:
            labeled_regions: List of labeled regions with threat levels
            
        Returns:
            Dictionary with threat summary statistics
        """
        if not labeled_regions:
            return {
                'total_regions': 0,
                'threat_counts': {'low': 0, 'medium': 0, 'high': 0},
                'max_threat': None,
                'average_difference': 0
            }
        
        # Count regions by threat level
        threat_counts = {'low': 0, 'medium': 0, 'high': 0}
        total_diff = 0
        max_threat = {'level': 'low', 'percentage': 0}
        
        for region in labeled_regions:
            level = region['threat_level']
            diff = region['difference_percentage']
            
            # Update counts
            threat_counts[level] += 1
            total_diff += diff
            
            # Track maximum threat
            if diff > max_threat['percentage']:
                max_threat = {'level': level, 'percentage': diff}
        
        # Calculate average difference
        avg_diff = total_diff / len(labeled_regions) if labeled_regions else 0
        
        return {
            'total_regions': len(labeled_regions),
            'threat_counts': threat_counts,
            'max_threat': max_threat,
            'average_difference': avg_diff
        }