#!/usr/bin/env python3
"""
YOLOv8 Training Script for French ID Card Detection
"""
import os
import sys
import argparse
from pathlib import Path
import logging
import torch

# Import config
from config import (
    DATA_YAML_PATH, TRAINING_LOG_PATH, create_directories
)

# Create necessary directories first
create_directories()

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler(TRAINING_LOG_PATH),
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger(__name__)

# Import modules
from modules.trainer import YOLOv8Trainer
from modules.data_preparator import DataPreparator

def check_dependencies():
    """Kiểm tra dependencies"""
    try:
        import ultralytics
        import torch
        import cv2
        import yaml
        logger.info("[OK] Dependencies checked")
        return True
    except ImportError as e:
        logger.error(f"[ERROR] Missing dependency: {e}")
        logger.info("Run: pip install -r requirements.txt")
        return False

def check_gpu():
    """Kiểm tra GPU"""
    try:
        import torch
        if torch.cuda.is_available():
            gpu_name = torch.cuda.get_device_name(0)
            gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1024**3
            logger.info(f"[OK] GPU: {gpu_name} ({gpu_memory:.1f} GB)")
            return True
        else:
            logger.warning("[WARNING] No GPU detected, using CPU")
            return False
    except Exception as e:
        logger.error(f"[ERROR] GPU check failed: {e}")
        return False

def validate_data(data_yaml_path):
    """Validate data trước khi training"""
    logger.info("[INFO] Validating data...")
    
    try:
        preparator = DataPreparator(data_yaml_path)
        preparator.run_full_validation()
        logger.info("[OK] Data validation completed")
        return True
    except Exception as e:
        logger.error(f"[ERROR] Data validation failed: {e}")
        return False

def main():
    """Main function"""
    parser = argparse.ArgumentParser(description='Train YOLOv8 for French ID Card Detection')
    parser.add_argument('--data', type=str, default=None, 
                       help='Path to data.yaml file (uses default if not specified)')
    parser.add_argument('--model-size', type=str, default='n', 
                       choices=['n', 's', 'm', 'l', 'x'],
                       help='Model size (n=nano, s=small, m=medium, l=large, x=xlarge)')
    parser.add_argument('--epochs', type=int, default=100, 
                       help='Number of training epochs')
    parser.add_argument('--batch-size', type=int, default=16, 
                       help='Batch size')
    parser.add_argument('--img-size', type=int, default=640, 
                       help='Input image size')
    parser.add_argument('--patience', type=int, default=50, 
                       help='Early stopping patience')
    parser.add_argument('--save-period', type=int, default=10, 
                       help='Save checkpoint every N epochs')
    parser.add_argument('--validate', action='store_true', 
                       help='Validate model after training')
    parser.add_argument('--export', type=str, default=None, 
                       help='Export model format (e.g., onnx, torchscript)')
    parser.add_argument('--model-path', type=str, default=None, 
                       help='Path to trained model for validation/export')
    parser.add_argument('--skip-validation', action='store_true',
                       help='Skip data validation')
    parser.add_argument('--validate-only', action='store_true',
                       help='Only validate data, skip training')
    
    args = parser.parse_args()
    
    logger.info("=" * 60)
    logger.info("YOLOv8 French ID Card Detection Training")
    logger.info("=" * 60)
    
    # Kiểm tra dependencies
    logger.info("\n1. Checking dependencies...")
    if not check_dependencies():
        sys.exit(1)
    
    # Kiểm tra GPU
    logger.info("\n2. Checking GPU...")
    check_gpu()
    
    # Kiểm tra data
    logger.info("\n3. Checking data...")
    data_path = Path(args.data) if args.data else DATA_YAML_PATH
    if not data_path.exists():
        logger.error(f"[ERROR] Data file not found: {data_path}")
        sys.exit(1)
    logger.info("[OK] Data configuration found")
    
    # Validate data (nếu không skip)
    if not args.skip_validation:
        logger.info("\n4. Validating data...")
        if not validate_data(str(data_path)):
            logger.error("Data validation failed. Please check your data.")
            if not args.validate_only:
                sys.exit(1)
    
    # Chạy training (nếu không chỉ validate)
    if not args.validate_only:
        logger.info("\n5. Starting training...")
        logger.info(f"Configuration:")
        logger.info(f"  - Model size: {args.model_size}")
        logger.info(f"  - Epochs: {args.epochs}")
        logger.info(f"  - Batch size: {args.batch_size}")
        logger.info(f"  - Image size: {args.img_size}")
        logger.info(f"  - Patience: {args.patience}")
        
        try:
            # Initialize trainer
            trainer = YOLOv8Trainer(str(data_path), args.model_size)
            
            # Train model
            if args.model_path is None:
                logger.info("Starting training...")
                results = trainer.train(
                    epochs=args.epochs,
                    batch=args.batch_size,  # Sửa từ batch_size thành batch
                    imgsz=args.img_size,
                    patience=args.patience,
                    save_period=args.save_period
                )
            
            # Validate model
            if args.validate:
                logger.info("Validating model...")
                trainer.validate(args.model_path)
            
            # Export model
            if args.export:
                logger.info(f"Exporting model to {args.export} format...")
                trainer.export_model(args.model_path, args.export)
            
            logger.info("[OK] Training completed successfully!")
            
        except Exception as e:
            logger.error(f"[ERROR] Training failed: {e}")
            sys.exit(1)
    
    logger.info("\n" + "=" * 60)
    logger.info("[SUCCESS] Process completed successfully!")
    logger.info("=" * 60)
    
    # Thông tin về kết quả
    if not args.validate_only:
        logger.info("\n[INFO] Training results:")
        logger.info(f"  - Model weights: runs/train/yolov8_*_french_id_card/weights/")
        logger.info(f"  - Training logs: {TRAINING_LOG_PATH}")
        logger.info(f"  - Plots: runs/train/yolov8_*_french_id_card/")
        
        logger.info("\n[INFO] To evaluate your model:")
        logger.info(f"  python eval.py --model-size {args.model_size}")
        
        logger.info("\n[INFO] To test your model:")
        logger.info(f"  python inference.py --model runs/train/yolov8_{args.model_size}_french_id_card/weights/best.pt --input path/to/image.jpg")

if __name__ == '__main__':
    main() 