"""
ATM Anomaly Detection Model

YOLOv8-based anomaly detector for ATM images.
Detects: damage, graffiti, obstruction, skimming devices, out-of-service
"""

import torch
import torch.nn as nn
from ultralytics import YOLO
from pathlib import Path
from typing import List, Dict, Tuple, Optional
import numpy as np
import cv2

class ATMAnomalyDetector:
    """
    ATM Anomaly Detection using YOLOv8.
    
    Usage:
        detector = ATMAnomalyDetector(model_path='models/best.pt')
        results = detector.predict('path/to/image.jpg')
    """
    
    # Anomaly class mapping
    CLASS_NAMES = {
        0: 'physical_damage',
        1: 'vandalism',
        2: 'graffiti',
        3: 'dirt_debris',
        4: 'obstruction',
        5: 'skimming_device',
        6: 'suspicious_attachment',
        7: 'out_of_service',
        8: 'screen_damage',
        9: 'cash_jam',
        10: 'receipt_jam',
        11: 'lighting_failure',
        12: 'camera_blind',
        13: 'network_down'
    }
    
    SEVERITY_MAP = {
        'skimming_device': 5,
        'suspicious_attachment': 5,
        'physical_damage': 4,
        'vandalism': 4,
        'camera_blind': 4,
        'obstruction': 3,
        'out_of_service': 3,
        'screen_damage': 3,
        'cash_jam': 3,
        'lighting_failure': 3,
        'network_down': 3,
        'graffiti': 2,
        'dirt_debris': 2,
        'receipt_jam': 2
    }
    
    def __init__(
        self,
        model_path: Optional[str] = None,
        conf_threshold: float = 0.25,
        iou_threshold: float = 0.45,
        device: str = 'auto'
    ):
        """
        Initialize detector.
        
        Args:
            model_path: Path to trained YOLO model
            conf_threshold: Confidence threshold for detections
            iou_threshold: IoU threshold for NMS
            device: 'cpu', 'cuda', or 'auto'
        """
        self.conf_threshold = conf_threshold
        self.iou_threshold = iou_threshold
        
        # Set device
        if device == 'auto':
            self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        else:
            self.device = device
        
        # Load model
        if model_path and Path(model_path).exists():
            self.model = YOLO(model_path)
        else:
            # Load pretrained COCO model as base
            print("No trained model found. Loading YOLOv8n pretrained...")
            self.model = YOLO('yolov8n.pt')
        
        self.model.to(self.device)
    
    def predict(
        self,
        image_path: str,
        save: bool = False,
        save_dir: Optional[str] = None
    ) -> List[Dict]:
        """
        Run anomaly detection on image.
        
        Args:
            image_path: Path to image file
            save: Whether to save annotated image
            save_dir: Directory to save annotations
            
        Returns:
            List of detection dicts with keys:
                - class_id: int
                - class_name: str
                - confidence: float
                - bbox: [x1, y1, x2, y2]
                - severity: int
        """
        # Run inference
        results = self.model(
            image_path,
            conf=self.conf_threshold,
            iou=self.iou_threshold,
            device=self.device,
            verbose=False
        )
        
        detections = []
        
        for result in results:
            boxes = result.boxes
            if boxes is None:
                continue
            
            for box in boxes:
                class_id = int(box.cls.item())
                confidence = float(box.conf.item())
                bbox = box.xyxy[0].cpu().numpy().tolist()
                
                class_name = self.CLASS_NAMES.get(class_id, 'unknown')
                severity = self.SEVERITY_MAP.get(class_name, 1)
                
                detection = {
                    'class_id': class_id,
                    'class_name': class_name,
                    'confidence': round(confidence, 4),
                    'bbox': [round(x, 2) for x in bbox],
                    'severity': severity,
                    'requires_action': severity >= 4
                }
                detections.append(detection)
        
        # Sort by severity (highest first)
        detections.sort(key=lambda x: x['severity'], reverse=True)
        
        # Save annotated image if requested
        if save and save_dir:
            self._save_annotated(image_path, detections, save_dir)
        
        return detections
    
    def predict_batch(
        self,
        image_paths: List[str],
        batch_size: int = 8
    ) -> List[List[Dict]]:
        """
        Run detection on batch of images.
        
        Args:
            image_paths: List of image paths
            batch_size: Batch size for inference
            
        Returns:
            List of detection lists
        """
        all_detections = []
        
        for i in range(0, len(image_paths), batch_size):
            batch = image_paths[i:i + batch_size]
            results = self.model(
                batch,
                conf=self.conf_threshold,
                iou=self.iou_threshold,
                device=self.device,
                verbose=False
            )
            
            for result in results:
                detections = []
                boxes = result.boxes
                
                if boxes is not None:
                    for box in boxes:
                        class_id = int(box.cls.item())
                        confidence = float(box.conf.item())
                        bbox = box.xyxy[0].cpu().numpy().tolist()
                        class_name = self.CLASS_NAMES.get(class_id, 'unknown')
                        
                        detections.append({
                            'class_id': class_id,
                            'class_name': class_name,
                            'confidence': round(confidence, 4),
                            'bbox': [round(x, 2) for x in bbox],
                            'severity': self.SEVERITY_MAP.get(class_name, 1),
                            'requires_action': self.SEVERITY_MAP.get(class_name, 1) >= 4
                        })
                
                detections.sort(key=lambda x: x['severity'], reverse=True)
                all_detections.append(detections)
        
        return all_detections
    
    def _save_annotated(
        self,
        image_path: str,
        detections: List[Dict],
        save_dir: str
    ):
        """Save annotated image with bounding boxes."""
        import os
        os.makedirs(save_dir, exist_ok=True)
        
        image = cv2.imread(image_path)
        
        for det in detections:
            x1, y1, x2, y2 = map(int, det['bbox'])
            color = (0, 0, 255) if det['severity'] >= 4 else (0, 165, 255)
            
            cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
            
            label = f"{det['class_name']} {det['confidence']:.2f}"
            cv2.putText(
                image, label, (x1, y1 - 10),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2
            )
        
        filename = Path(image_path).name
        save_path = os.path.join(save_dir, f"annotated_{filename}")
        cv2.imwrite(save_path, image)
    
    def train(
        self,
        data_yaml: str,
        epochs: int = 100,
        batch_size: int = 16,
        img_size: int = 640,
        output_dir: str = 'models/checkpoints'
    ):
        """
        Train model on custom dataset.
        
        Args:
            data_yaml: Path to data.yaml for YOLO
            epochs: Number of training epochs
            batch_size: Batch size
            img_size: Input image size
            output_dir: Where to save checkpoints
        """
        self.model.train(
            data=data_yaml,
            epochs=epochs,
            batch=batch_size,
            imgsz=img_size,
            project=output_dir,
            name='atm_anomaly',
            device=self.device
        )
    
    def export(
        self,
        format: str = 'onnx',
        output_path: Optional[str] = None
    ):
        """
        Export model to deployment format.
        
        Args:
            format: 'onnx', 'torchscript', 'openvino', 'engine'
            output_path: Where to save exported model
        """
        self.model.export(format=format)
        
        if output_path:
            import shutil
            default_path = f"models/checkpoints/atm_anomaly/weights/best.{format}"
            if Path(default_path).exists():
                shutil.copy2(default_path, output_path)
                print(f"Exported to {output_path}")


def create_data_yaml(
    train_dir: str,
    val_dir: str,
    test_dir: Optional[str] = None,
    class_names: Optional[List[str]] = None,
    output_path: str = 'data/data.yaml'
):
    """
    Create YOLO data.yaml configuration file.
    
    Args:
        train_dir: Path to train directory
        val_dir: Path to validation directory
        test_dir: Path to test directory (optional)
        class_names: List of class names
        output_path: Where to save yaml
    """
    if class_names is None:
        class_names = list(ATMAnomalyDetector.CLASS_NAMES.values())
    
    import yaml
    
    data = {
        'path': str(Path(train_dir).parent),
        'train': str(Path(train_dir).relative_to(Path(train_dir).parent)),
        'val': str(Path(val_dir).relative_to(Path(val_dir).parent)),
        'nc': len(class_names),
        'names': class_names
    }
    
    if test_dir:
        data['test'] = str(Path(test_dir).relative_to(Path(test_dir).parent))
    
    with open(output_path, 'w') as f:
        yaml.dump(data, f, default_flow_style=False)
    
    print(f"Created {output_path}")


if __name__ == "__main__":
    # Example usage
    detector = ATMAnomalyDetector()
    
    # Single image prediction
    results = detector.predict('data/test/atm_001.jpg', save=True, save_dir='output')
    print(f"Found {len(results)} anomalies")
    for r in results:
        print(f"  - {r['class_name']}: {r['confidence']:.2f} (severity: {r['severity']})")