"""
AI-powered smart cropping and composition.

This module provides intelligent cropping and composition features using
deep learning models for content-aware video processing.
"""

import torch
import numpy as np
from typing import List, Dict, Tuple, Optional
import logging
import cv2
from dataclasses import dataclass

logger = logging.getLogger(__name__)

@dataclass
class CropRegion:
    """Represents a recommended crop region."""
    x: int
    y: int
    width: int
    height: int
    score: float
    label: str

class SmartCropper:
    """Content-aware video cropping."""

    def __init__(self, model_name: str = "yolov5s", device: str = "auto"):
        """
        Initialize smart cropper.

        Args:
            model_name: Name of detection model
            device: Device to run inference on ("cpu", "cuda", or "auto")
        """
        self.model_name = model_name

        # Determine device
        if device == "auto":
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
        else:
            self.device = device

        # Load model
        try:
            if model_name.startswith("yolov5"):
                self.model = torch.hub.load("ultralytics/yolov5", model_name)
                self.model.to(self.device)
            else:
                raise ValueError(f"Unsupported model: {model_name}")
        except Exception as e:
            logger.error(f"Failed to load model {model_name}: {e}")
            raise

    def find_regions(self, frame: np.ndarray,
                    target_ratio: float = 16/9) -> List[CropRegion]:
        """
        Find optimal crop regions.

        Args:
            frame: Input frame (RGB)
            target_ratio: Target aspect ratio (width/height)

        Returns:
            List of recommended crop regions
        """
        # Run object detection
        img = torch.from_numpy(frame).to(self.device)
        img = img.permute(2, 0, 1)  # HWC to CHW
        results = self.model([img])

        regions = []
        height, width = frame.shape[:2]

        # Process detections
        for pred in results.pred[0]:
            if pred[4] >= 0.5:  # Confidence threshold
                # Get box coordinates
                x1, y1, x2, y2 = pred[:4].cpu().numpy()
                box_width = x2 - x1
                box_height = y2 - y1

                # Get class info
                cls_id = int(pred[5])
                label = self.model.names[cls_id]
                confidence = float(pred[4])

                # Calculate crop region around object
                target_width = box_width * 1.5
                target_height = target_width / target_ratio

                # Center crop around object
                cx = (x1 + x2) / 2
                cy = (y1 + y2) / 2

                x = max(0, int(cx - target_width/2))
                y = max(0, int(cy - target_height/2))
                w = min(width - x, int(target_width))
                h = min(height - y, int(target_height))

                regions.append(CropRegion(x, y, w, h, confidence, label))

        return regions

    def get_best_region(self, frame: np.ndarray,
                       target_ratio: float = 16/9) -> Optional[CropRegion]:
        """
        Get single best crop region.

        Args:
            frame: Input frame (RGB)
            target_ratio: Target aspect ratio (width/height)

        Returns:
            Best crop region or None if no good regions found
        """
        regions = self.find_regions(frame, target_ratio)
        if not regions:
            return None

        # Return region with highest confidence
        return max(regions, key=lambda r: r.score)

    def crop_frame(self, frame: np.ndarray, region: CropRegion) -> np.ndarray:
        """
        Crop frame to specified region.

        Args:
            frame: Input frame
            region: Crop region

        Returns:
            Cropped frame
        """
        return frame[region.y:region.y+region.height,
                    region.x:region.x+region.width]
