#!/usr/bin/env python3
"""
YOLOv8 Inference Module for French ID Card Detection
"""
import os
import sys
import argparse
from pathlib import Path
import logging
import cv2
import numpy as np
from ultralytics import YOLO
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from PIL import Image, ImageDraw, ImageFont
import json

# Import config
sys.path.append(str(Path(__file__).parent.parent))
from config import (
    INFERENCE_RESULTS_DIR, EVALUATION_RESULTS_DIR, 
    VISUALIZATION_RESULTS_DIR, DEFAULT_INFERENCE_CONFIG
)

logger = logging.getLogger(__name__)

class YOLOv8Inference:
    """
    YOLOv8 Inference for French ID Card Detection
    """
    
    def __init__(self, model_path: str, conf_threshold: float = None, iou_threshold: float = None):
        """
        Initialize YOLOv8 Inference
        
        Args:
            model_path: Path to trained model
            conf_threshold: Confidence threshold (uses default if None)
            iou_threshold: IoU threshold for NMS (uses default if None)
        """
        self.model_path = Path(model_path)
        self.conf_threshold = conf_threshold or DEFAULT_INFERENCE_CONFIG['conf_threshold']
        self.iou_threshold = iou_threshold or DEFAULT_INFERENCE_CONFIG['iou_threshold']
        
        if not self.model_path.exists():
            raise FileNotFoundError(f"Model not found: {model_path}")
        
        # Load model
        self.model = YOLO(model_path)
        logger.info(f"Model loaded: {model_path}")
        logger.info(f"Confidence threshold: {self.conf_threshold}")
        logger.info(f"IoU threshold: {self.iou_threshold}")
    
    def predict_single_image(self, image_path: str, save_result: bool = True, 
                           output_dir: str = None) -> dict:
        """
        Predict on a single image
        
        Args:
            image_path: Path to input image
            save_result: Whether to save result image
            output_dir: Output directory for results (uses default if None)
            
        Returns:
            Prediction results
        """
        if output_dir is None:
            output_dir = INFERENCE_RESULTS_DIR
        
        image_path = Path(image_path)
        if not image_path.exists():
            raise FileNotFoundError(f"Image not found: {image_path}")
        
        logger.info(f"Processing image: {image_path}")
        
        # Run inference
        results = self.model.predict(
            source=str(image_path),
            conf=self.conf_threshold,
            iou=self.iou_threshold,
            save=save_result,
            project=output_dir,
            name='predictions'
        )
        
        # Process results
        result = results[0] if results else None
        
        if result is None:
            logger.warning(f"No detections found in {image_path}")
            return {'detections': [], 'image_path': str(image_path)}
        
        # Extract detection information
        detections = []
        if result.boxes is not None:
            boxes = result.boxes.xyxy.cpu().numpy()  # x1, y1, x2, y2
            confidences = result.boxes.conf.cpu().numpy()
            class_ids = result.boxes.cls.cpu().numpy()
            
            for i in range(len(boxes)):
                detection = {
                    'bbox': boxes[i].tolist(),  # [x1, y1, x2, y2]
                    'confidence': float(confidences[i]),
                    'class_id': int(class_ids[i]),
                    'class_name': 'french'  # Based on your data.yaml
                }
                detections.append(detection)
        
        logger.info(f"Found {len(detections)} detections in {image_path.name}")
        
        return {
            'detections': detections,
            'image_path': str(image_path),
            'result_path': str(result.save_dir) if hasattr(result, 'save_dir') else None
        }
    
    def predict_batch(self, input_dir: str, output_dir: str = None) -> dict:
        """
        Predict on a batch of images
        
        Args:
            input_dir: Input directory containing images
            output_dir: Output directory for results (uses default if None)
            
        Returns:
            Batch prediction results
        """
        if output_dir is None:
            output_dir = INFERENCE_RESULTS_DIR
        
        input_path = Path(input_dir)
        if not input_path.exists():
            raise FileNotFoundError(f"Input directory not found: {input_dir}")
        
        # Find all image files
        image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff'}
        image_files = []
        
        for file_path in input_path.rglob('*'):
            if file_path.is_file() and file_path.suffix.lower() in image_extensions:
                image_files.append(file_path)
        
        if not image_files:
            logger.warning(f"No images found in {input_dir}")
            return {'total_images': 0, 'processed_images': 0, 'results': []}
        
        logger.info(f"Processing {len(image_files)} images from {input_dir}")
        
        results = {
            'total_images': len(image_files),
            'processed_images': 0,
            'results': []
        }
        
        # Process each image
        for i, image_path in enumerate(image_files):
            try:
                logger.info(f"Processing {i+1}/{len(image_files)}: {image_path.name}")
                
                result = self.predict_single_image(
                    str(image_path),
                    save_result=True,
                    output_dir=output_dir
                )
                
                results['results'].append(result)
                results['processed_images'] += 1
                
            except Exception as e:
                logger.error(f"Error processing {image_path}: {e}")
        
        # Summary
        total_detections = sum(len(r['detections']) for r in results['results'])
        logger.info(f"Batch processing completed:")
        logger.info(f"  - Total images: {results['total_images']}")
        logger.info(f"  - Processed: {results['processed_images']}")
        logger.info(f"  - Total detections: {total_detections}")
        
        return results
    
    def visualize_detections(self, image_path: str, detections: list, 
                           save_path: str = None, show: bool = False):
        """
        Visualize detections on image
        
        Args:
            image_path: Path to input image
            detections: List of detection dictionaries
            save_path: Path to save visualization (uses default if None)
            show: Whether to show the plot
        """
        if save_path is None:
            save_path = VISUALIZATION_RESULTS_DIR / f"viz_{Path(image_path).stem}.png"
        
        # Load image
        image = cv2.imread(image_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        # Create figure
        fig, ax = plt.subplots(1, 1, figsize=(12, 8))
        ax.imshow(image)
        
        # Draw detections
        for detection in detections:
            bbox = detection['bbox']
            confidence = detection['confidence']
            class_name = detection['class_name']
            
            # Create rectangle
            x1, y1, x2, y2 = bbox
            width = x2 - x1
            height = y2 - y1
            
            rect = patches.Rectangle(
                (x1, y1), width, height,
                linewidth=2, edgecolor='red', facecolor='none'
            )
            ax.add_patch(rect)
            
            # Add text
            text = f"{class_name}: {confidence:.2f}"
            ax.text(x1, y1-10, text, color='red', fontsize=12, 
                   bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.8))
        
        ax.set_title(f"Detections: {len(detections)}")
        ax.axis('off')
        
        if save_path:
            plt.savefig(save_path, bbox_inches='tight', dpi=300)
            logger.info(f"Visualization saved to {save_path}")
        
        if show:
            plt.show()
        
        plt.close()
    
    def evaluate_on_test_set(self, test_dir: str, labels_dir: str = None) -> dict:
        """
        Evaluate model on test set
        
        Args:
            test_dir: Directory containing test images
            labels_dir: Directory containing ground truth labels (optional)
            
        Returns:
            Evaluation results
        """
        test_path = Path(test_dir)
        if not test_path.exists():
            raise FileNotFoundError(f"Test directory not found: {test_dir}")
        
        # Get test images
        image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff'}
        test_images = []
        
        for file_path in test_path.rglob('*'):
            if file_path.is_file() and file_path.suffix.lower() in image_extensions:
                test_images.append(file_path)
        
        if not test_images:
            logger.warning(f"No test images found in {test_dir}")
            return {}
        
        logger.info(f"Evaluating on {len(test_images)} test images")
        
        # Run predictions
        results = self.predict_batch(test_dir, EVALUATION_RESULTS_DIR)
        
        # Calculate metrics
        total_detections = sum(len(r['detections']) for r in results['results'])
        avg_detections = total_detections / len(test_images) if test_images else 0
        
        evaluation_results = {
            'total_images': len(test_images),
            'total_detections': total_detections,
            'avg_detections_per_image': avg_detections,
            'detection_rate': len([r for r in results['results'] if r['detections']]) / len(test_images),
            'results': results['results']
        }
        
        logger.info("Evaluation results:")
        logger.info(f"  - Total images: {evaluation_results['total_images']}")
        logger.info(f"  - Total detections: {evaluation_results['total_detections']}")
        logger.info(f"  - Avg detections per image: {evaluation_results['avg_detections_per_image']:.2f}")
        logger.info(f"  - Detection rate: {evaluation_results['detection_rate']:.2f}")
        
        return evaluation_results
    
    def export_results(self, results: dict, output_file: str = None):
        """
        Export results to JSON file
        
        Args:
            results: Results dictionary
            output_file: Output file path (uses default if None)
        """
        if output_file is None:
            output_file = INFERENCE_RESULTS_DIR / "inference_results.json"
        
        with open(output_file, 'w') as f:
            json.dump(results, f, indent=2)
        
        logger.info(f"Results exported to {output_file}") 