#!/usr/bin/env python3
"""
YOLOv8 Trainer Module
"""
import os
import sys
import yaml
import argparse
from pathlib import Path
import logging
from ultralytics import YOLO
import torch
import shutil

# Import config
sys.path.append(str(Path(__file__).parent.parent))
from config import (
    DATA_YAML_PATH, TRAINING_LOG_PATH, DEFAULT_TRAINING_CONFIG, get_best_model_path
)

logger = logging.getLogger(__name__)

class YOLOv8Trainer:
    """
    YOLOv8 Trainer for French ID Card Detection
    """
    
    def __init__(self, data_yaml_path: str = None, model_size: str = 'n'):
        """
        Initialize YOLOv8 Trainer
        
        Args:
            data_yaml_path: Path to data.yaml file (optional, uses default if None)
            model_size: Model size ('n', 's', 'm', 'l', 'x')
        """
        self.data_yaml_path = Path(data_yaml_path) if data_yaml_path else DATA_YAML_PATH
        self.model_size = model_size
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        
        logger.info(f"Using device: {self.device}")
        logger.info(f"Model size: {model_size}")
        
        # Validate data.yaml
        self._validate_data_yaml()
    
    def _validate_data_yaml(self):
        """Validate data.yaml file"""
        if not self.data_yaml_path.exists():
            raise FileNotFoundError(f"data.yaml not found at {self.data_yaml_path}")
        
        with open(self.data_yaml_path, 'r') as f:
            data_config = yaml.safe_load(f)
        
        # Check required fields
        required_fields = ['train', 'val', 'nc', 'names']
        for field in required_fields:
            if field not in data_config:
                raise ValueError(f"Missing required field '{field}' in data.yaml")
        
        # Check if paths exist
        train_path = Path(data_config['train'])
        val_path = Path(data_config['val'])
        
        if not train_path.exists():
            logger.warning(f"Training path does not exist: {train_path}")
        
        if not val_path.exists():
            logger.warning(f"Validation path does not exist: {val_path}")
        
        logger.info(f"Data configuration validated:")
        logger.info(f"  - Classes: {data_config['nc']}")
        logger.info(f"  - Class names: {data_config['names']}")
        logger.info(f"  - Training path: {data_config['train']}")
        logger.info(f"  - Validation path: {data_config['val']}")
    
    def train(self, epochs: int = None, batch: int = None, imgsz: int = None,
              patience: int = None, save_period: int = None, **kwargs):
        """
        Train YOLOv8 model
        
        Args:
            epochs: Number of training epochs
            batch: Batch size
            imgsz: Input image size
            patience: Early stopping patience
            save_period: Save checkpoint every N epochs
            **kwargs: Additional training arguments
        """
        logger.info("Starting YOLOv8 training...")
        
        # Initialize model - chỉ dùng YOLOv8
        model = YOLO(f'yolov8{self.model_size}.pt')
        
        # Get training configuration
        train_args = DEFAULT_TRAINING_CONFIG.copy()
        
        # Update with provided arguments
        if epochs is not None:
            train_args['epochs'] = epochs
        if batch is not None:
            train_args['batch'] = batch
        if imgsz is not None:
            train_args['imgsz'] = imgsz
        if patience is not None:
            train_args['patience'] = patience
        if save_period is not None:
            train_args['save_period'] = save_period
        
        # Update with additional kwargs
        train_args.update(kwargs)
        
        # Set specific paths
        train_args['data'] = str(self.data_yaml_path)
        train_args['device'] = self.device
        train_args['name'] = f'yolov8_{self.model_size}_french_id_card'
        
        logger.info("Training configuration:")
        for key, value in train_args.items():
            if key in ['data', 'epochs', 'batch', 'imgsz', 'patience', 'device']:
                logger.info(f"  {key}: {value}")
        
        try:
            # Start training
            results = model.train(**train_args)
            
            logger.info("Training completed successfully!")
            logger.info(f"Best model saved at: {results.save_dir}")
            
            return results
            
        except Exception as e:
            logger.error(f"Training failed: {e}")
            raise
    
    def validate(self, model_path: str = None):
        """
        Validate trained model
        
        Args:
            model_path: Path to trained model (if None, uses best model from runs/train)
        """
        if model_path is None:
            # Use best model from runs/train
            model_path = get_best_model_path(self.model_size)
        
        if not model_path or not Path(model_path).exists():
            logger.error(f"Model not found: {model_path}")
            return
        
        logger.info(f"Validating model: {model_path}")
        
        # Load model and validate
        model = YOLO(model_path)
        results = model.val(data=str(self.data_yaml_path))
        
        logger.info("Validation completed!")
        return results
    
    def export_model(self, model_path: str = None, format: str = 'onnx'):
        """
        Export trained model to different formats
        
        Args:
            model_path: Path to trained model
            format: Export format ('onnx', 'torchscript', 'tflite', etc.)
        """
        if model_path is None:
            # Use best model from runs/train
            model_path = get_best_model_path(self.model_size)
        
        if not model_path or not Path(model_path).exists():
            logger.error(f"Model not found: {model_path}")
            return
        
        logger.info(f"Exporting model: {model_path} to {format}")
        
        # Load model and export
        model = YOLO(model_path)
        exported_path = model.export(format=format)
        
        logger.info(f"Model exported to: {exported_path}")
        return exported_path
    
    def get_latest_model(self, model_size: str = None) -> str:
        """
        Get path to latest trained model
        
        Args:
            model_size: Model size (if None, uses self.model_size)
        
        Returns:
            Path to latest model
        """
        if model_size is None:
            model_size = self.model_size
        
        model_path = TRAINED_MODELS_DIR / f"yolov8_{model_size}_french_id_card.pt"
        
        if model_path.exists():
            return str(model_path)
        else:
            logger.warning(f"No trained model found for size {model_size}")
            return None 