
import os
from ultralytics import YOLO
import argparse
from model_utils import get_project_root, find_latest_model


def evaluate_model(model_path=None, config_path=None, split='val'):
    """评估模型
    
    Args:
        model_path: 模型路径（如果为 None 则自动查找）
        config_path: dataset.yaml 配置文件路径
        split: 评估集（train/val/test）
    """
    project_root = get_project_root()
    
    # 如果没提供 config_path，使用项目根目录的 dataset.yaml
    if config_path is None:
        config_path = os.path.join(project_root, 'dataset.yaml')
    
    print(f"配置文件: {config_path}")
    
    # 检查配置文件是否存在
    if not os.path.exists(config_path):
        print(f"\n错误: 配置文件不存在！")
        print(f"路径: {config_path}")
        print(f"\n请先配置 dataset.yaml:")
        print(f"  1. 复制 configs/dataset-traditional.yaml 到项目根目录")
        print(f"  2. 重命名为 dataset.yaml")
        print(f"  3. 修改 path 指向你的数据集目录")
        return None
    
    # 如果没提供模型路径，自动查找最新模型
    if model_path is None or not os.path.exists(model_path):
        found_model = find_latest_model()
        if found_model:
            model_path = found_model
            print(f"找到最新模型: {model_path}")
        else:
            print(f"\n错误: 找不到模型文件！")
            print(f"\n可能的原因:")
            print(f"  1. 还没有运行训练，先运行: python src/train.py")
            print(f"  2. 使用 --model 参数指定正确路径")
            return None
    else:
        print(f"模型路径: {model_path}")
    
    try:
        # 加载模型
        model = YOLO(model_path)
        
        # 评估
        print(f"\n开始评估 {split} 集...")
        metrics = model.val(data=config_path, split=split)
        
        print("\n" + "=" * 50)
        print("模型评估结果")
        print("=" * 50)
        print(f"mAP50: {metrics.box.map50:.4f}")
        print(f"mAP50-95: {metrics.box.map:.4f}")
        print(f"Precision: {metrics.box.mp:.4f}")
        print(f"Recall: {metrics.box.mr:.4f}")
        print("=" * 50)
        
        return metrics
        
    except Exception as e:
        print(f"\n评估失败: {e}")
        return None


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='评估 YOLO 模型')
    parser.add_argument('--model', default=None,
                        help='模型路径（默认：自动查找最新模型）')
    parser.add_argument('--config', default=None,
                        help='dataset.yaml 配置文件路径（默认使用项目根目录）')
    parser.add_argument('--split', default='val',
                        help='评估集（train/val/test，默认: val）')
    
    args = parser.parse_args()
    
    evaluate_model(args.model, args.config, args.split)
