import yaml, sys
import torch
from pathlib import Path

# --- 工具函数：动态设备选择 (Dynamic Device Selection) ---
def get_auto_device(verbose = False):
    """自动检测并返回最合适的设备"""
    if torch.cuda.is_available():
        gpu_count = torch.cuda.device_count()
        if gpu_count > 1:
            # 如果有多个GPU，使用所有GPU
            device_str = ",".join(str(i) for i in range(gpu_count))
            if verbose:
                print(f"检测到 {gpu_count} 个可用的GPU。将使用所有GPU: {device_str}")
            return device_str
        else:
            # 如果只有1个GPU
            if verbose:
                print("检测到 1 个可用的GPU。将使用 GPU: 0")
            return '0'
    else:
        # 如果没有可用的GPU，使用CPU
        if verbose:
            print("未检测到可用的GPU。将使用CPU。")
        return 'cpu'
# --- 1. 核心目录设置 (Core Directories) ---
HARDISK_DIR = Path.home() / "Desktop" / "Seg" / "Hardisk" # 硬盘根目录
BASE_DIR = Path(__file__).parent.parent # 当前脚本所在目录上级
DATASET_YAML_PATH = Path(__file__).parent / "dataset.yaml" # 数据集的 YAML 配置文件路径
OUTPUTS_DIR = BASE_DIR / "DataSet_Public_outputs" # 输出结果目录 # 最终为：OUTPUTS_DIR / dataset_name
PREDICT_ALL_BEST_MODELS_DIR = BASE_DIR / 'BestMode_Predict_Results_DataSet_Public'   # 所有最佳模型存放目录

TEST_IMAGE_DIR = None # 初始化 TEST_IMAGE_DIR # 测试文件地址 dataset.path / TODO "val" / "test" TODO
# try:
# 3. 读取并解析 YAML 文件
with open(DATASET_YAML_PATH, 'r', encoding='utf-8') as f:
    yaml_data = yaml.safe_load(f)
# 4. 从解析后的数据中获取 'path' 的值
relative_path_from_yaml = yaml_data.get('path')
test_path_from_yaml = yaml_data.get('test')
dataset_name = Path(relative_path_from_yaml).name
if relative_path_from_yaml:
    # 5. 【核心步骤】构建绝对路径
    #    relative_path_from_yaml 是相对于 .yaml 文件本身的路径。
    #    所以，我们需要获取 .yaml 文件所在的目录，然后与这个相对路径拼接。
    yaml_file_directory = DATASET_YAML_PATH.parent
    TEST_IMAGE_DIR = (DATASET_YAML_PATH.parent / relative_path_from_yaml / test_path_from_yaml).resolve() # 这里val 或 test在dataset.yaml中定义
    # 使用 .exists() 方法来检查路径是否存在
    if not TEST_IMAGE_DIR.exists():
        # 如果路径不存在，则执行这里的代码
        print(f"警告: 测试图片目录不存在: {TEST_IMAGE_DIR}")
        sys.exit(1)
    # 6. 获取OUTPUTS_DIR
    if dataset_name and dataset_name not in {'.', '..'}:
        dataset_name_ = dataset_name + "-Yolo"
        # 设定输出路径
        OUTPUTS_DIR = OUTPUTS_DIR / dataset_name_
        # ../BestMode_Predict_Results_DataSet_Public/"dataset_name+"-Yolo"" # 需要用到 Tool_Yolo_Copy_Best_Model.sh
        PREDICT_BEST_MODEL_DIR = PREDICT_ALL_BEST_MODELS_DIR / dataset_name_     # 最优模型保存位置 PREDICT_ALL_BEST_MODELS_DIR / dataset_name_ / weights / best.pt
    else:
        print(f"警告: 提取的dataset_name: '{dataset_name}' 无效（为空、'.' 或 '..'）。")
        sys.exit(1)
else:
    print(f"警告: 在 '{DATASET_YAML_PATH}' 文件中没有找到 'path' 键。")
    sys.exit(1)
# except FileNotFoundError:
#     print(f"错误: YAML 配置文件未找到: '{DATASET_YAML_PATH}'")
# except Exception as e:
#     print(f"读取或解析 YAML 文件时发生错误: {e}")


# --- 修改：模型、图像大小和批处理大小的集成配置 ---
# 将所有模型的配置（权重、图像大小、批处理大小）集中管理
MODEL_CONFIGS = {
    # YOLOv8
    'YOLOv8n-seg': {'weights': 'yolov8n-seg.pt', 'image_size': 640, 'batch_size': 16}, 
    'YOLOv8s-seg': {'weights': 'yolov8s-seg.pt', 'image_size': 640, 'batch_size': 16},
    'YOLOv8m-seg': {'weights': 'yolov8m-seg.pt', 'image_size': 640, 'batch_size': 16},
    'YOLOv8l-seg': {'weights': 'yolov8l-seg.pt', 'image_size': 640, 'batch_size': 16}, # 640，16，12.5GB
    'YOLOv8x-seg': {'weights': 'yolov8x-seg.pt', 'image_size': 640, 'batch_size': 16}, # 640，16，15.5GB # 示例：X模型使用1280分辨率
    # YOLOv9
    'YOLOv9c-seg': {'weights': 'yolov9c-seg.pt', 'image_size': 640, 'batch_size': 16}, # 640，16，13GB
    'YOLOv9e-seg': {'weights': 'yolov9e-seg.pt', 'image_size': 640, 'batch_size': 8}, # 640，16，内存超了 # 示例：E模型（最大）使用1280分辨率和更小的batch
    # YOLOv11 (假设)
    'YOLO11n-seg': {'weights': 'yolo11n-seg.pt', 'image_size': 640, 'batch_size': 16},
    'YOLO11s-seg': {'weights': 'yolo11s-seg.pt', 'image_size': 640, 'batch_size': 16},
    'YOLO11m-seg': {'weights': 'yolo11m-seg.pt', 'image_size': 640, 'batch_size': 16},
    'YOLO11l-seg': {'weights': 'yolo11l-seg.pt', 'image_size': 640, 'batch_size': 16}, # 640，16，12.5GB
    'YOLO11x-seg': {'weights': 'yolo11x-seg.pt', 'image_size': 640, 'batch_size': 16}, # 640，16，19.5GB
    # YOLOv12 (假设)
    'YOLO12-seg': {'weights': str(Path(__file__).parent / 'yolo12-seg.yaml'), 'image_size': 640, 'batch_size': 16}, # 640，16，3GB
}

# --- 4. 训练超参数 (Training Hyperparameters) ---
DEVICE = get_auto_device(verbose=False) # 调用函数来设置设备
EPOCHS = 300          # 训练轮次 # TODO
PATIENCE = 150        # 提前停止训练的轮数
SAVE_PERIOD = 10     # 每隔多少轮保存一次模型 # TODO
# BATCH_SIZE 已在上面根据模型自动设置
LEARNING_RATE = 0.01  # 初始学习率（'SGD:0.01', 'Adam:1e-4', 'AdamW:1e-4', 'auto:0.01'）
OPTIMIZER = 'auto'    # 优化器 (TODO 'SGD', 'Adam', 'AdamW', 'auto' TODO)
WORKERS = 4           # 数据加载的工作线程数

# --- 5. 预测设置 (Prediction Settings) ---
SHOW_LABELS = True # 是否在预测结果上显示类别标签
SHOW_CONF = True # 是否在预测结果上显示置信度和边界框
SAVE_PREDICTIONS = True # 是否保存预测结果图像

# ==============================================================================
# --- 主执行块：只在直接运行时才打印配置信息 ---
# ==============================================================================
def show_config_summary():
    """打印所有配置摘要信息"""
    print(f"设定输出路径为: '{str(OUTPUTS_DIR)}'")
    print(f"最佳模型路径为: '{str(PREDICT_BEST_MODEL_DIR)}'")
    # 再次调用 get_auto_device 并设置 verbose=True 来打印设备信息
    get_auto_device(verbose=True)


# 当这个脚本被直接执行时（python yolo_config.py），__name__ 的值是 '__main__'
# 当它被其他脚本import时，__name__ 的值是 'yolo_config'
if __name__ == '__main__':
    show_config_summary()