import os
import shutil
from pathlib import Path

def group_files_by_scale(src_dir: str):
    """
    按照文件名中的 scale 值（如 0.25, 0.5, 1.0, 2, 4）将 .pt 文件分组，
    分别移动到 src_dir_0.25、src_dir_0.5 等子目录中。
    
    参数:
        src_dir (str): 原始文件所在目录
    """
    src_path = Path(src_dir).resolve()
    if not src_path.is_dir():
        raise ValueError(f"{src_dir} 不是有效目录")

    # 遍历目录下的所有 .pt 文件
    for file in src_path.glob("*.pt"):
        # 从文件名中提取 scale 值
        # 例如：100_scale_0.25_dh-0_raw_results_2025-09-15-15-13-54.pt
        scale_val = file.stem.split("scale_")[-1].split("_")[0]

        # 构造目标目录：原目录名 + "_" + scale_val
        dst_dir = src_path.parent / f"{src_path.name}_{scale_val}"
        dst_dir.mkdir(exist_ok=True)

        # 构造目标文件路径
        dst_file = dst_dir / file.name

        # 移动文件
        shutil.move(str(file), str(dst_file))
        # print(f"Moved: {file} -> {dst_file}")
        # input()

# 示例用法
if __name__ == "__main__":
    for dirs in os.listdir("."):
        if not os.path.isdir(dirs) or "init_7" not in dirs:
            continue
        print(dirs)
        k = input()
        if k == "q":
            continue
        for dataset in os.listdir(dirs):
            group_files_by_scale(os.path.join(dirs, dataset))