"""
ID Card Detector Module
Sử dụng YOLO để detect và cắt ID cards từ ảnh lớn, kết hợp với data augmentation
Tích hợp với YOLOv8 French ID Card Detection model
"""
import cv2
import numpy as np
from pathlib import Path
from typing import List, Tuple, Optional, Dict, Any, Union
import torch
import torch.nn as nn
from ultralytics import YOLO
import logging
from data_augmentation import DataAugmentation
from utils import load_image, save_image, create_augmented_filename, print_progress
import os
import json
import yaml

class IDCardDetector:
    """Class để detect và cắt ID cards từ ảnh lớn sử dụng YOLO"""
    
    def __init__(self, model_path: str = None, config: Dict[str, Any] = None):
        """
        Initialize ID Card Detector
        
        Args:
            model_path: Đường dẫn đến model YOLO đã train
            config: Configuration dictionary
        """
        self.config = config or {}
        self.model_path = model_path
        self.model = None
        self.data_augmentation = DataAugmentation(config)
        self.logger = self._setup_logger()
        
        # Default model path nếu không được cung cấp
        if not model_path:
            default_model_path = "data/weights/id_cards_yolov8n.pt"
            if os.path.exists(default_model_path):
                model_path = default_model_path
                self.model_path = model_path
        
        # Load YOLO model nếu có
        if model_path and os.path.exists(model_path):
            self.load_model(model_path)
    
    def _setup_logger(self) -> logging.Logger:
        """Setup logger cho module"""
        logger = logging.getLogger(__name__)
        logger.setLevel(logging.INFO)
        
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
            )
            handler.setFormatter(formatter)
            logger.addHandler(handler)
        
        return logger
    
    def load_model(self, model_path: str) -> bool:
        """
        Load YOLO model từ file
        
        Args:
            model_path: Đường dẫn đến model file
            
        Returns:
            True nếu load thành công, False nếu thất bại
        """
        try:
            self.model = YOLO(model_path)
            self.logger.info(f"Loaded YOLO model from: {model_path}")
            return True
        except Exception as e:
            self.logger.error(f"Failed to load model: {e}")
            return False
    
    def detect_id_cards(self, image: np.ndarray, confidence: float = 0.5, iou_threshold: float = 0.45) -> List[Dict[str, Any]]:
        """
        Detect ID cards trong ảnh sử dụng YOLO
        
        Args:
            image: Input image
            confidence: Confidence threshold
            iou_threshold: IoU threshold cho NMS
            
        Returns:
            List các detection results với format:
            {
                'bbox': [x1, y1, x2, y2],
                'confidence': float,
                'class_id': int,
                'class_name': str
            }
        """
        if self.model is None:
            self.logger.error("Model chưa được load!")
            return []
        
        try:
            # Run inference
            results = self.model(image, conf=confidence, iou=float(iou_threshold), verbose=False)
            
            detections = []
            for result in results:
                boxes = result.boxes
                if boxes is not None:
                    for box in boxes:
                        # Get bbox coordinates
                        x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
                        
                        # Get confidence and class
                        confidence_score = float(box.conf[0].cpu().numpy())
                        class_id = int(box.cls[0].cpu().numpy())
                        class_name = self.model.names[class_id] if hasattr(self.model, 'names') else f"class_{class_id}"
                        
                        detection = {
                            'bbox': [int(x1), int(y1), int(x2), int(y2)],
                            'confidence': confidence_score,
                            'class_id': class_id,
                            'class_name': class_name
                        }
                        detections.append(detection)
            
            self.logger.info(f"Detected {len(detections)} ID cards")
            return detections
            
        except Exception as e:
            self.logger.error(f"Error during detection: {e}")
            return []
    
    def crop_id_card(self, image: np.ndarray, bbox: List[int], padding: int = 10, 
                     crop_mode: str = "bbox", target_size: Tuple[int, int] = None) -> np.ndarray:
        """
        Cắt ID card từ ảnh gốc dựa trên bbox với nhiều options
        
        Args:
            image: Input image
            bbox: Bounding box [x1, y1, x2, y2]
            padding: Padding thêm xung quanh bbox
            crop_mode: Mode cắt ("bbox", "square", "aspect_ratio")
            target_size: Kích thước target (width, height) nếu muốn resize
            
        Returns:
            Cropped ID card image
        """
        x1, y1, x2, y2 = bbox
        
        # Thêm padding
        height, width = image.shape[:2]
        x1 = max(0, x1 - padding)
        y1 = max(0, y1 - padding)
        x2 = min(width, x2 + padding)
        y2 = min(height, y2 + padding)
        
        # Cắt ảnh theo mode
        if crop_mode == "bbox":
            # Cắt theo bbox gốc
            cropped = image[y1:y2, x1:x2]
        elif crop_mode == "square":
            # Cắt thành hình vuông
            center_x = (x1 + x2) // 2
            center_y = (y1 + y2) // 2
            size = max(x2 - x1, y2 - y1)
            half_size = size // 2
            
            x1 = max(0, center_x - half_size)
            y1 = max(0, center_y - half_size)
            x2 = min(width, center_x + half_size)
            y2 = min(height, center_y + half_size)
            
            cropped = image[y1:y2, x1:x2]
        elif crop_mode == "aspect_ratio":
            # Cắt theo tỷ lệ khung hình chuẩn (3:4 cho ID card)
            bbox_width = x2 - x1
            bbox_height = y2 - y1
            center_x = (x1 + x2) // 2
            center_y = (y1 + y2) // 2
            
            # Tỷ lệ 3:4 cho ID card
            target_ratio = 3 / 4
            current_ratio = bbox_width / bbox_height
            
            if current_ratio > target_ratio:
                # Bbox quá rộng, giữ chiều cao
                new_width = int(bbox_height * target_ratio)
                half_width = new_width // 2
                x1 = max(0, center_x - half_width)
                x2 = min(width, center_x + half_width)
            else:
                # Bbox quá cao, giữ chiều rộng
                new_height = int(bbox_width / target_ratio)
                half_height = new_height // 2
                y1 = max(0, center_y - half_height)
                y2 = min(height, center_y + half_height)
            
            cropped = image[y1:y2, x1:x2]
        else:
            # Default: cắt theo bbox
            cropped = image[y1:y2, x1:x2]
        
        # Resize nếu có target_size
        if target_size:
            cropped = cv2.resize(cropped, target_size, interpolation=cv2.INTER_AREA)
        
        return cropped
    
    def process_single_image(self, image_path: Union[str, Path], output_dir: Path, 
                           confidence: float = 0.5, iou_threshold: float = 0.45,
                           crop_mode: str = "bbox", target_size: Tuple[int, int] = None,
                           padding: int = 10, card_counter: int = 0) -> Dict[str, Any]:
        """
        Xử lý một ảnh: detect ID cards, cắt và áp dụng augmentation
        
        Args:
            image_path: Đường dẫn đến ảnh input
            output_dir: Thư mục output
            apply_augmentation: Có áp dụng data augmentation không
            save_original: Có lưu ảnh gốc không
            confidence: Confidence threshold
            iou_threshold: IoU threshold
            crop_mode: Mode cắt ("bbox", "square", "aspect_ratio")
            target_size: Kích thước target (width, height) hoặc None
            padding: Padding thêm xung quanh bbox
            
        Returns:
            Dictionary chứa kết quả xử lý
        """
        image_path = Path(image_path)
        if not image_path.exists():
            self.logger.error(f"Image not found: {image_path}")
            return {}
        
        # Load ảnh
        image = load_image(str(image_path))
        if image is None:
            self.logger.error(f"Failed to load image: {image_path}")
            return {}
        
        # Detect ID cards
        detections = self.detect_id_cards(image, confidence, float(iou_threshold))
        
        if not detections:
            self.logger.warning(f"No ID cards detected in: {image_path}")
            return {
                'image_path': str(image_path),
                'detections': [],
                'processed_cards': []
            }
        
        # Tạo thư mục output
        output_dir.mkdir(parents=True, exist_ok=True)
        
        processed_cards = []
        current_card_counter = card_counter
        
        for i, detection in enumerate(detections):
            # Cắt ID card với options mới
            cropped_card = self.crop_id_card(
                image, 
                detection['bbox'], 
                padding=padding,
                crop_mode=crop_mode,
                target_size=target_size
            )
            
            # Tạo tên file unique cho mỗi ID card
            current_card_counter += 1
            card_filename = f"id_card_{current_card_counter:03d}.jpg"
            card_path = output_dir / card_filename
            
            # Lưu ảnh gốc
            save_image(cropped_card, card_path)
            processed_cards.append({
                'original_path': str(card_path),
                'detection_info': detection,
                'crop_info': {
                    'mode': crop_mode,
                    'target_size': target_size,
                    'padding': padding
                }
            })
        
        result = {
            'image_path': str(image_path),
            'detections': detections,
            'processed_cards': processed_cards,
            'total_cards': len(processed_cards),
            'crop_settings': {
                'mode': crop_mode,
                'target_size': target_size,
                'padding': padding
            }
        }
        
        self.logger.info(f"Processed {len(processed_cards)} cards from {image_path.name}")
        return result
    
    def batch_process(self, input_dir: Union[str, Path], output_dir: Union[str, Path],
                     confidence: float = 0.5, iou_threshold: float = 0.45,
                     crop_mode: str = "bbox", target_size: Tuple[int, int] = None,
                     padding: int = 10) -> Dict[str, Any]:
        """
        Xử lý batch nhiều ảnh
        
        Args:
            input_dir: Thư mục chứa ảnh input
            output_dir: Thư mục output
            apply_augmentation: Có áp dụng data augmentation không
            save_original: Có lưu ảnh gốc không
            confidence: Confidence threshold
            iou_threshold: IoU threshold
            crop_mode: Mode cắt ("bbox", "square", "aspect_ratio")
            target_size: Kích thước target (width, height) hoặc None
            padding: Padding thêm xung quanh bbox
            
        Returns:
            Dictionary chứa kết quả batch processing
        """
        input_dir = Path(input_dir)
        output_dir = Path(output_dir)
        
        if not input_dir.exists():
            self.logger.error(f"Input directory not found: {input_dir}")
            return {}
        
        # Tạo thư mục output
        output_dir.mkdir(parents=True, exist_ok=True)
        
        # Tìm tất cả ảnh
        supported_formats = self.config.get('supported_formats', ['.jpg', '.jpeg', '.png', '.bmp', '.tiff'])
        image_files = []
        for fmt in supported_formats:
            image_files.extend(input_dir.glob(f"*{fmt}"))
            image_files.extend(input_dir.glob(f"*{fmt.upper()}"))
        
        if not image_files:
            self.logger.warning(f"No supported images found in: {input_dir}")
            return {}
        
        self.logger.info(f"Found {len(image_files)} images to process")
        
        results = {}
        total_cards = 0
        global_card_counter = 0  # Counter để tạo tên file unique
        
        for i, image_path in enumerate(image_files):
            self.logger.info(f"Processing {i+1}/{len(image_files)}: {image_path.name}")
            
            # Xử lý ảnh - chỉ detect và crop, không augment
            result = self.process_single_image(
                image_path, 
                output_dir,
                confidence,
                iou_threshold,
                crop_mode,
                target_size,
                padding,
                global_card_counter
            )
            
            # Cập nhật counter
            global_card_counter += len(result.get('detections', []))
            
            results[image_path.name] = result
            total_cards += len(result.get('detections', []))  # Số lượng ID cards thực tế đã detect
            
            # Print progress
            print_progress(i + 1, len(image_files), f"Processed {image_path.name}")
        
        # Tạo summary
        summary = {
            'total_images': len(image_files),
            'total_cards_detected': total_cards,
            'images_with_cards': len([r for r in results.values() if r.get('detections')]),
            'images_without_cards': len([r for r in results.values() if not r.get('detections')]),
            'output_directory': str(output_dir),
            'crop_settings': {
                'mode': crop_mode,
                'target_size': target_size,
                'padding': padding
            },
            'results': results
        }
        
        # Lưu summary
        summary_path = output_dir / "processing_summary.json"
        with open(summary_path, 'w', encoding='utf-8') as f:
            json.dump(summary, f, indent=2, ensure_ascii=False)
        
        self.logger.info(f"Batch processing completed. Summary saved to: {summary_path}")
        return summary
    
    def get_detection_statistics(self, results: Dict[str, Any]) -> Dict[str, Any]:
        """
        Tính toán thống kê từ kết quả detection
        
        Args:
            results: Kết quả từ batch_process
            
        Returns:
            Dictionary chứa thống kê
        """
        if not results:
            return {}
        
        total_images = results.get('total_images', 0)
        total_cards = results.get('total_cards_detected', 0)
        images_with_cards = results.get('images_with_cards', 0)
        
        # Tính confidence statistics
        all_confidences = []
        for image_result in results.get('results', {}).values():
            for detection in image_result.get('detections', []):
                all_confidences.append(detection.get('confidence', 0))
        
        stats = {
            'total_images_processed': total_images,
            'total_cards_detected': total_cards,
            'images_with_cards': images_with_cards,
            'images_without_cards': total_images - images_with_cards,
            'average_cards_per_image': total_cards / total_images if total_images > 0 else 0,
            'detection_rate': images_with_cards / total_images if total_images > 0 else 0,
            'confidence_statistics': {
                'min': min(all_confidences) if all_confidences else 0,
                'max': max(all_confidences) if all_confidences else 0,
                'mean': np.mean(all_confidences) if all_confidences else 0,
                'std': np.std(all_confidences) if all_confidences else 0
            }
        }
        
        return stats
    
    def augment_cropped_cards(self, input_dir: Union[str, Path], output_dir: Union[str, Path],
                             num_augmentations: int = 3) -> Dict[str, Any]:
        """
        Augment tất cả ID cards đã crop trong thư mục input
        
        Args:
            input_dir: Thư mục chứa ID cards đã crop
            output_dir: Thư mục output cho augmented images
            num_augmentations: Số lượng augmentation cho mỗi card
            
        Returns:
            Dictionary chứa kết quả augmentation
        """
        input_dir = Path(input_dir)
        output_dir = Path(output_dir)
        
        if not input_dir.exists():
            self.logger.error(f"Input directory not found: {input_dir}")
            return {}
        
        # Tạo thư mục output
        output_dir.mkdir(parents=True, exist_ok=True)
        
        # Tìm tất cả ID cards đã crop
        card_files = list(input_dir.glob("id_card_*.jpg"))
        
        if not card_files:
            self.logger.warning(f"No ID card files found in: {input_dir}")
            return {}
        
        self.logger.info(f"Found {len(card_files)} ID cards to augment")
        
        results = {}
        total_augmented = 0
        
        for i, card_path in enumerate(card_files):
            self.logger.info(f"Augmenting {i+1}/{len(card_files)}: {card_path.name}")
            
            # Load ID card
            card_image = load_image(str(card_path))
            if card_image is None:
                self.logger.error(f"Failed to load card: {card_path}")
                continue
            
            # Augment card
            try:
                augmented_cards = self.data_augmentation.augment_single_image(
                    card_image, 
                    num_augmentations=num_augmentations
                )
                
                # Debug: Kiểm tra số lượng augmented cards
                self.logger.info(f"Generated {len(augmented_cards)} augmented cards for {card_path.name}")
                
                # Debug: Kiểm tra config
                self.logger.info(f"DataAugmentation config: {self.data_augmentation.config}")
                
            except Exception as e:
                self.logger.error(f"Error during augmentation: {e}")
                augmented_cards = []
            
            # Save augmented cards
            card_results = []
            for j, aug_card in enumerate(augmented_cards):
                aug_filename = f"{card_path.stem}_aug_{j+1}.jpg"
                aug_path = output_dir / aug_filename
                save_image(aug_card, aug_path)
                
                card_results.append({
                    'augmented_path': str(aug_path),
                    'augmentation_index': j+1
                })
            
            results[card_path.name] = {
                'original_path': str(card_path),
                'augmented_cards': card_results,
                'total_augmented': len(card_results)
            }
            
            total_augmented += len(card_results)
            
            # Print progress
            print_progress(i + 1, len(card_files), f"Augmented {card_path.name}")
        
        # Tạo summary
        summary = {
            'total_cards': len(card_files),
            'total_augmented': total_augmented,
            'output_directory': str(output_dir),
            'results': results
        }
        
        # Lưu summary
        summary_path = output_dir / "augmentation_summary.json"
        with open(summary_path, 'w', encoding='utf-8') as f:
            json.dump(summary, f, indent=2, ensure_ascii=False)
        
        self.logger.info(f"Augmentation completed. Summary saved to: {summary_path}")
        return summary
    
    def load_yolo_config(self, config_path: str = None) -> Dict[str, Any]:
        """
        Load config từ YOLO detector
        
        Args:
            config_path: Đường dẫn đến file config
            
        Returns:
            Config dictionary
        """
        if config_path is None:
            # Tìm config mặc định
            default_config_path = "src/model/ID_cards_detector/config.py"
            if os.path.exists(default_config_path):
                config_path = default_config_path
        
        config = {}
        
        try:
            # Import config từ YOLO detector
            import sys
            sys.path.append(str(Path("src/model/ID_cards_detector")))
            
            from config import DEFAULT_TRAINING_CONFIG, DEFAULT_INFERENCE_CONFIG
            
            config.update({
                'yolo_training_config': DEFAULT_TRAINING_CONFIG,
                'yolo_inference_config': DEFAULT_INFERENCE_CONFIG,
                'detection': {
                    'confidence_threshold': DEFAULT_INFERENCE_CONFIG.get('conf_threshold', 0.25),
                    'iou_threshold': DEFAULT_INFERENCE_CONFIG.get('iou_threshold', 0.45),
                    'padding': 10
                },
                'processing': {
                    'apply_augmentation': True,
                    'save_original': True,
                    'num_augmentations': 3,
                    'save_format': "jpg",
                    'quality': 95,
                    'target_size': [640, 640]
                },
                'crop_options': {
                    'crop_mode': 'bbox',  # bbox, square, aspect_ratio
                    'target_size': None,  # (width, height) hoặc None
                    'padding': 10
                }
            })
            
            self.logger.info("Loaded YOLO config successfully")
            
        except Exception as e:
            self.logger.warning(f"Failed to load YOLO config: {e}")
            # Fallback config
            config = {
                'detection': {
                    'confidence_threshold': 0.25,
                    'iou_threshold': 0.45,
                    'padding': 10
                },
                'processing': {
                    'apply_augmentation': True,
                    'save_original': True,
                    'num_augmentations': 3,
                    'save_format': "jpg",
                    'quality': 95,
                    'target_size': [640, 640]
                },
                'crop_options': {
                    'crop_mode': 'bbox',
                    'target_size': None,
                    'padding': 10
                }
            }
        
        return config 