
import os
from ultralytics import YOLO
import torch
import argparse
from model_utils import get_project_root


def train_yolo(config_path=None, model_name='yolov8n.pt', epochs=100, imgsz=640,
               batch=16, device=None, lr0=0.01, patience=50,
               project='runs/detect', name='my_model', exist_ok=True):
    """训练 YOLO 模型
    
    Args:
        config_path: dataset.yaml 配置文件路径
    """
    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 device is None:
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"使用设备: {device}")
    
    try:
        # 加载模型
        print(f"\n加载模型: {model_name}")
        model = YOLO(model_name)
        
        # 训练 - 确保 project 路径是相对于项目根目录的
        print(f"开始训练...")
        results = model.train(
            data=config_path,
            epochs=epochs,
            imgsz=imgsz,
            batch=batch,
            device=device,
            lr0=lr0,
            patience=patience,
            save=True,
            plots=True,
            project=os.path.join(project_root, project),
            name=name,
            exist_ok=exist_ok,
        )
        
        best_model_path = results.save_dir / 'weights' / 'best.pt'
        print(f"\n训练完成！")
        print(f"最佳模型保存在: {best_model_path}")
        print(f"\n现在可以运行:")
        print(f"  评估模型: python src/evaluate.py --model {best_model_path}")
        print(f"  进行预测: python src/predict.py --model {best_model_path}")
        
        return results
        
    except Exception as e:
        print(f"\n训练失败: {e}")
        return None


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='训练 YOLO 模型')
    parser.add_argument('--config', default=None,
                        help='dataset.yaml 配置文件路径（默认使用项目根目录）')
    parser.add_argument('--model', default='yolov8n.pt',
                        help='预训练模型 (默认: yolov8n.pt)')
    parser.add_argument('--epochs', type=int, default=100,
                        help='训练轮数 (默认: 100)')
    parser.add_argument('--imgsz', type=int, default=640,
                        help='输入图片尺寸 (默认: 640)')
    parser.add_argument('--batch', type=int, default=16,
                        help='批次大小 (默认: 16)')
    parser.add_argument('--device', default=None,
                        help='设备 (cuda/cpu, 默认: 自动检测)')
    parser.add_argument('--lr', type=float, default=0.01,
                        help='初始学习率 (默认: 0.01)')
    parser.add_argument('--patience', type=int, default=50,
                        help='早停耐心值 (默认: 50)')
    parser.add_argument('--project', default='runs/detect',
                        help='项目目录 (默认: runs/detect)')
    parser.add_argument('--name', default='my_model',
                        help='实验名称 (默认: my_model)')
    
    args = parser.parse_args()
    
    train_yolo(
        config_path=args.config,
        model_name=args.model,
        epochs=args.epochs,
        imgsz=args.imgsz,
        batch=args.batch,
        device=args.device,
        lr0=args.lr,
        patience=args.patience,
        project=args.project,
        name=args.name
    )
