#!/usr/bin/env python3
"""
Evaluation script for YOLOv8 French ID Card Detection
"""
import os
import sys
import argparse
import logging
from pathlib import Path
import yaml
from ultralytics import YOLO

# Import config
sys.path.append(str(Path(__file__).parent))
from config import (
    DATA_YAML_PATH, EVAL_LOG_PATH, get_best_model_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(EVAL_LOG_PATH),
        logging.StreamHandler(sys.stdout)
    ]
)

logger = logging.getLogger(__name__)

def check_dependencies():
    """Check if required dependencies are installed"""
    try:
        import ultralytics
        import torch
        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():
    """Check GPU availability"""
    try:
        import torch
        if torch.cuda.is_available():
            gpu_name = torch.cuda.get_device_name(0)
            logger.info(f"[OK] GPU available: {gpu_name}")
            return True
        else:
            logger.warning("[WARNING] No GPU available, using CPU")
            return False
    except Exception as e:
        logger.error(f"[ERROR] GPU check failed: {e}")
        return False

def make_data_yaml_absolute(data_yaml_path):
    """Tạo file data.yaml tạm với các đường dẫn tuyệt đối cho train/val/test"""
    with open(data_yaml_path, 'r') as f:
        data = yaml.safe_load(f)
    
    # Lấy thư mục chứa data.yaml (data/)
    yaml_dir = Path(data_yaml_path).parent.resolve()
    
    # Map các đường dẫn tương đối sang đúng cấu trúc thư mục
    path_mapping = {
        '../train/images': 'train/images',
        '../valid/images': 'valid/images', 
        '../test/images': 'test/images'
    }
    
    for key in ['train', 'val', 'test']:
        if key in data:
            rel_path = data[key]
            # Kiểm tra nếu là đường dẫn tương đối
            if not str(rel_path).startswith('/') and not str(rel_path).startswith('C:'):
                # Map sang đường dẫn đúng trong thư mục data/
                if rel_path in path_mapping:
                    correct_path = path_mapping[rel_path]
                    abs_path = yaml_dir / correct_path
                    data[key] = str(abs_path.resolve())
                else:
                    # Fallback: resolve như cũ
                    abs_path = (yaml_dir / rel_path).resolve()
                    data[key] = str(abs_path)
    
    abs_yaml_path = yaml_dir / 'data_abs.yaml'
    with open(abs_yaml_path, 'w') as f:
        yaml.safe_dump(data, f)
    return str(abs_yaml_path)

# Sửa lại load_data_config để trả về đường dẫn tuyệt đối

def load_data_config():
    """Load and validate data configuration, trả về đường dẫn data_abs.yaml"""
    try:
        abs_yaml_path = make_data_yaml_absolute(DATA_YAML_PATH)
        with open(abs_yaml_path, 'r') as f:
            data_config = yaml.safe_load(f)
        # Check test path
        test_path = Path(data_config.get('test', ''))
        if not test_path.exists():
            logger.error(f"[ERROR] Test path does not exist: {test_path}")
            return None
        logger.info(f"[INFO] Test path: {test_path}")
        logger.info(f"[INFO] Classes: {data_config['names']}")
        return abs_yaml_path
    except Exception as e:
        logger.error(f"[ERROR] Failed to load data config: {e}")
        return None

# Sửa lại evaluate_model để nhận data_yaml_path là file tuyệt đối

def evaluate_model(model_path: str, data_yaml_path: str, conf_threshold: float = 0.25, iou_threshold: float = 0.45):
    """
    Evaluate model on test set
    
    Args:
        model_path: Path to trained model
        data_yaml_path: Path to data.yaml (absolute paths)
        conf_threshold: Confidence threshold
        iou_threshold: IoU threshold
    """
    try:
        logger.info(f"[INFO] Loading model: {model_path}")
        model = YOLO(model_path)
        logger.info("[INFO] Starting evaluation on test set...")
        results = model.val(
            data=data_yaml_path,
            split='test',  # Use test split
            conf=conf_threshold,
            iou=iou_threshold,
            verbose=True,
            save_json=True,  # Save results as JSON
            save_txt=True,   # Save results as TXT
            save_conf=True,  # Save confidence scores
            project='runs/val',
            name='test_evaluation',
            exist_ok=True
        )
        logger.info("[SUCCESS] Evaluation completed!")
        logger.info(f"[INFO] Results saved to: runs/val/test_evaluation/")
        if hasattr(results, 'results_dict'):
            metrics = results.results_dict
            logger.info(f"[INFO] mAP50: {metrics.get('metrics/mAP50', 'N/A')}")
            logger.info(f"[INFO] mAP50-95: {metrics.get('metrics/mAP50-95', 'N/A')}")
            logger.info(f"[INFO] Precision: {metrics.get('metrics/precision', 'N/A')}")
            logger.info(f"[INFO] Recall: {metrics.get('metrics/recall', 'N/A')}")
        return results
    except Exception as e:
        logger.error(f"[ERROR] Evaluation failed: {e}")
        return None

# Sửa lại main để lấy abs_yaml_path từ load_data_config

def main():
    """Main evaluation function"""
    parser = argparse.ArgumentParser(description='Evaluate YOLOv8 French ID Card Detection Model')
    parser.add_argument('--model', type=str, default=None, 
                       help='Path to trained model (if None, uses best model from runs/train)')
    parser.add_argument('--data', type=str, default=None,
                       help='Path to data.yaml (if None, uses default)')
    parser.add_argument('--conf', type=float, default=0.25,
                       help='Confidence threshold')
    parser.add_argument('--iou', type=float, default=0.45,
                       help='IoU threshold')
    parser.add_argument('--model-size', type=str, default='n',
                       help='Model size (n, s, m, l, x)')
    args = parser.parse_args()
    logger.info("=" * 50)
    logger.info("YOLOv8 French ID Card Detection - Evaluation")
    logger.info("=" * 50)
    if not check_dependencies():
        return
    check_gpu()
    # Lấy đường dẫn data.yaml tuyệt đối
    abs_yaml_path = load_data_config()
    if not abs_yaml_path:
        return
    if args.model:
        model_path = args.model
    else:
        model_path = get_best_model_path(args.model_size)
        if not model_path:
            logger.error("[ERROR] No trained model found. Please train a model first.")
            return
    logger.info(f"[INFO] Model: {model_path}")
    logger.info(f"[INFO] Data: {abs_yaml_path}")
    logger.info(f"[INFO] Confidence threshold: {args.conf}")
    logger.info(f"[INFO] IoU threshold: {args.iou}")
    results = evaluate_model(
        model_path=model_path,
        data_yaml_path=abs_yaml_path,
        conf_threshold=args.conf,
        iou_threshold=args.iou
    )
    if results:
        logger.info("[SUCCESS] Evaluation completed successfully!")
        logger.info(f"[INFO] Results saved to: runs/val/test_evaluation/")
    else:
        logger.error("[ERROR] Evaluation failed!")

if __name__ == '__main__':
    main() 