import cv2
import numpy as np
from typing import List, Dict, Tuple
import yaml
import math

class DamageVisualizer:
    """Visualize detection and comparison results"""

    def __init__(self, config_path: str = "config.yaml"):
        """Initialize visualizer with configuration"""
        with open(config_path, 'r') as f:
            self.config = yaml.safe_load(f)

        self.line_thickness = self.config['visualization']['line_thickness']
        self.font_scale = self.config['visualization']['font_scale']
        self.colors = self.config['visualization']['colors']

    def draw_detections(self, image: np.ndarray, detections: Dict,
                       color_type: str = 'new_damage') -> np.ndarray:
        """
        Draw bounding boxes and labels on image

        Args:
            image: Input image
            detections: Detection results
            color_type: Type of color to use ('new_damage', 'existing_damage', 'matched_damage')

        Returns:
            Image with drawn detections
        """
        img_copy = image.copy()
        color = self.colors.get(color_type, [255, 0, 0])

        for i, box in enumerate(detections['boxes']):
            x1, y1, x2, y2 = box
            label = f"{detections['classes'][i]} ({detections['confidences'][i]:.2f})"

            # Draw rectangle
            cv2.rectangle(img_copy, (x1, y1), (x2, y2), color, self.line_thickness)

            # Draw label background
            label_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX,
                                           self.font_scale, self.line_thickness)
            cv2.rectangle(img_copy, (x1, y1 - label_size[1] - 5),
                         (x1 + label_size[0], y1), color, -1)

            # Draw label text
            cv2.putText(img_copy, label, (x1, y1 - 5),
                       cv2.FONT_HERSHEY_SIMPLEX, self.font_scale,
                       [255, 255, 255], self.line_thickness)

        return img_copy

    def create_comparison_visualization(self, before_img: np.ndarray, after_img: np.ndarray,
                                      before_detections: Dict, after_detections: Dict,
                                      comparison_result: Dict) -> np.ndarray:
        """
        Create side-by-side comparison visualization

        Args:
            before_img, after_img: Input images
            before_detections, after_detections: Detection results
            comparison_result: Comparison analysis results

        Returns:
            Combined visualization image
        """
        # Draw matched damages in yellow
        before_vis = before_img.copy()
        after_vis = after_img.copy()

        # Draw matched damages
        for match in comparison_result['matched_damages']:
            # Draw on before image
            x1, y1, x2, y2 = match['box_before']
            cv2.rectangle(before_vis, (x1, y1), (x2, y2),
                         self.colors['matched_damage'], self.line_thickness)

            # Draw on after image
            x1, y1, x2, y2 = match['box_after']
            cv2.rectangle(after_vis, (x1, y1), (x2, y2),
                         self.colors['matched_damage'], self.line_thickness)

        # Draw repaired damages (only on before) in green
        for damage in comparison_result['repaired_damages']:
            x1, y1, x2, y2 = damage['box']
            cv2.rectangle(before_vis, (x1, y1), (x2, y2),
                         self.colors['existing_damage'], self.line_thickness)
            cv2.putText(before_vis, "REPAIRED", (x1, y1 - 5),
                       cv2.FONT_HERSHEY_SIMPLEX, self.font_scale,
                       self.colors['existing_damage'], self.line_thickness)

        # Draw new damages (only on after) in red
        for damage in comparison_result['new_damages']:
            x1, y1, x2, y2 = damage['box']
            cv2.rectangle(after_vis, (x1, y1), (x2, y2),
                         self.colors['new_damage'], self.line_thickness + 1)
            cv2.putText(after_vis, "NEW!", (x1, y1 - 5),
                       cv2.FONT_HERSHEY_SIMPLEX, self.font_scale * 1.5,
                       self.colors['new_damage'], self.line_thickness)

        # Combine images side by side
        h1, w1 = before_vis.shape[:2]
        h2, w2 = after_vis.shape[:2]
        max_height = max(h1, h2)

        # Resize if needed
        if h1 != max_height:
            before_vis = cv2.resize(before_vis, (int(w1 * max_height / h1), max_height))
        if h2 != max_height:
            after_vis = cv2.resize(after_vis, (int(w2 * max_height / h2), max_height))

        # Create combined image
        combined = np.hstack([before_vis, after_vis])

        # Add status text
        status_height = 100
        status_img = np.ones((status_height, combined.shape[1], 3), dtype=np.uint8) * 255

        # Add case message
        case_color = (0, 128, 0) if 'SUCCESS' in comparison_result['case'] else (0, 0, 255)
        cv2.putText(status_img, comparison_result['message'], (20, 40),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.8, case_color, 2)

        # Add statistics
        stats_text = f"Before: {comparison_result['statistics']['total_before']} | " \
                    f"After: {comparison_result['statistics']['total_after']} | " \
                    f"Matched: {comparison_result['statistics']['matched']} | " \
                    f"New: {comparison_result['statistics']['new']}"
        cv2.putText(status_img, stats_text, (20, 70),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 1)

        # Combine with status
        final_image = np.vstack([status_img, combined])

        return final_image

    def create_summary_grid(self, comparison_results: List[Dict],
                           image_pairs: List[Tuple[np.ndarray, np.ndarray]], 
                           num_pairs: int = None) -> np.ndarray:
        """
        Create a flexible grid visualization for 1-6 position comparisons

        Args:
            comparison_results: List of comparison results for each position
            image_pairs: List of (before, after) image pairs
            num_pairs: Number of image pairs (if None, inferred from length)

        Returns:
            Grid visualization of all positions
        """
        if num_pairs is None:
            num_pairs = len(comparison_results)
        
        # Ensure we don't exceed available data
        num_pairs = min(num_pairs, len(comparison_results), len(image_pairs))
        
        if num_pairs == 0:
            # Return empty image if no pairs
            return np.ones((400, 600, 3), dtype=np.uint8) * 255
        
        grid_images = []
        target_size = (300, 200)  # Size for each position in grid

        for i in range(num_pairs):
            result = comparison_results[i]
            before_img, after_img = image_pairs[i]
            
            # Create mini comparison for each position
            before_small = cv2.resize(before_img, target_size)
            after_small = cv2.resize(after_img, target_size)

            # Add position label
            position_label = f"Position {i+1}"
            cv2.putText(before_small, position_label, (10, 20),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
            cv2.putText(after_small, position_label, (10, 20),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)

            # Add case indicator
            case_color = (0, 255, 0) if 'SUCCESS' in result['case'] else (0, 0, 255)
            if 'NEW_DAMAGE' in result['case']:
                case_color = (0, 0, 255)

            cv2.rectangle(after_small, (0, 0), (target_size[0], 30), case_color, -1)
            cv2.putText(after_small, result['case'][:10], (10, 20),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1)

            pair_img = np.hstack([before_small, after_small])
            grid_images.append(pair_img)

        # Calculate optimal grid layout based on number of pairs
        grid_layout = self._calculate_grid_layout(num_pairs)
        rows, cols = grid_layout['rows'], grid_layout['cols']
        
        # Create grid based on calculated layout
        grid_rows = []
        
        for row in range(rows):
            row_images = []
            for col in range(cols):
                idx = row * cols + col
                if idx < len(grid_images):
                    row_images.append(grid_images[idx])
                else:
                    # Create empty placeholder for missing positions
                    empty_img = np.ones((target_size[1], target_size[0] * 2, 3), dtype=np.uint8) * 240
                    cv2.putText(empty_img, "Empty", (target_size[0] - 30, target_size[1] // 2),
                               cv2.FONT_HERSHEY_SIMPLEX, 0.6, (150, 150, 150), 2)
                    row_images.append(empty_img)
            
            if row_images:
                grid_rows.append(np.hstack(row_images))
        
        # Combine all rows
        if grid_rows:
            grid = np.vstack(grid_rows)
        else:
            grid = np.ones((400, 600, 3), dtype=np.uint8) * 255
        
        # Add header with summary information
        header_height = 80
        header_img = np.ones((header_height, grid.shape[1], 3), dtype=np.uint8) * 250
        
        # Add title
        title_text = f"Damage Comparison Summary - {num_pairs} Position{'s' if num_pairs != 1 else ''}"
        cv2.putText(header_img, title_text, (20, 30),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 0), 2)
        
        # Add overall statistics
        total_new = sum(len(result.get('new_damages', [])) for result in comparison_results[:num_pairs])
        total_matched = sum(len(result.get('matched_damages', [])) for result in comparison_results[:num_pairs])
        
        stats_text = f"Total New Damages: {total_new} | Total Matched: {total_matched}"
        cv2.putText(header_img, stats_text, (20, 60),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 1)
        
        # Combine header with grid
        final_grid = np.vstack([header_img, grid])
        
        return final_grid
    
    def _calculate_grid_layout(self, num_pairs: int) -> Dict[str, int]:
        """
        Calculate optimal grid layout for given number of image pairs
        
        Args:
            num_pairs: Number of image pairs to display
            
        Returns:
            Dictionary with 'rows' and 'cols' keys
        """
        if num_pairs <= 0:
            return {'rows': 1, 'cols': 1}
        elif num_pairs == 1:
            return {'rows': 1, 'cols': 1}
        elif num_pairs == 2:
            return {'rows': 1, 'cols': 2}
        elif num_pairs == 3:
            return {'rows': 1, 'cols': 3}
        elif num_pairs == 4:
            return {'rows': 2, 'cols': 2}
        elif num_pairs == 5:
            return {'rows': 2, 'cols': 3}  # 2x3 grid with one empty slot
        elif num_pairs == 6:
            return {'rows': 2, 'cols': 3}
        else:
            # For more than 6, calculate optimal layout
            cols = math.ceil(math.sqrt(num_pairs))
            rows = math.ceil(num_pairs / cols)
            return {'rows': rows, 'cols': cols}