import os

import argparse

import os

def find_ckpt_files(base_dir, name):
    """
    在指定的 base_dir 目录下查找所有包含特定名称的 ckpt 文件的路径。
    
    参数:
    base_dir (str): 搜索的基本目录，例如 'trained_models'
    name (str): 需要查找的目录名称部分

    返回:
    list: 包含符合条件的 ckpt 文件路径的列表
    """
    ckpt_files = []
    # 遍历 base_dir 下的所有文件和目录
    for root, dirs, files in os.walk(base_dir):
        # 检查当前目录名称中是否包含指定的 name
        if name in os.path.basename(root):
            # 遍历当前目录下的所有文件
            for file in files:
                # 如果文件是 ckpt 文件，则收集其路径
                if file.endswith('.ckpt'):
                    ckpt_files.append(os.path.join(root, file))
    
    return ckpt_files

# 使用例子
# 假设你要找的目录名包含 'method1' 的所有 ckpt 文件
ckpt_paths = find_ckpt_files('trained_models', 'method1')
print(ckpt_paths)
                
# l = ['byol','nnbyol','mocov2plus','mocov3','dino']
# l = ['byol','dino']

resnet18s = ['byol-cifar100_none_0.01','byol-cifar100_risp_0.001','byol-cifar100_so_1e-06','dino-cifar100_none_0.001','dino-cifar100_risp_0.001',
'dino-cifar100_so_1e-06','mocov2plus-cifar100_none_0.001''mocov2plus-cifar100_risp_0.001','mocov2plus-cifar100_so_1e-06','mocov3-cifar100_none_0.001','mocov3-cifar100_risp_0.001','mocov3-cifar100_so_1e-06'
'nnbyol-cifar100_none_0.001','nnbyol-cifar100_risp_0.001','nnbyol-cifar100_so_1e-06']
resnet50s = ['byol-cifar100_none_0.001_resnet50','byol-cifar100_risp_0.0001_resnet50','byol-cifar100_so_1e-06_resnet50','dino-cifar100_none_0.001_resnet50',
'dino-cifar100_risp_0.001_resnet50','dino-cifar100_so_1e-06_resnet50','mocov2plus-cifar100_none_0.001_resnet50','mocov2plus-cifar100_risp_0.001_resnet50','mocov2plus-cifar100_so_1e-06_resnet50',
'mocov3-cifar100_none_0.001_resnet50','mocov3-cifar100_risp_0.001_resnet50','mocov3-cifar100_so_1e-06_resnet50','nnbyol-cifar100_none_0.001_resnet50',
'nnbyol-cifar100_risp_0.001_resnet50','nnbyol-cifar100_so_1e-06_resnet50']
wideresnets = ['byol-cifar100_none_0.01_wide_resnet28w2','byol-cifar100_risp_0.0001_wide_resnet28w2','byol-cifar100_so_1e-07_wide_resnet28w2','dino-cifar100_none_1e-05_wide_resnet28w2',
'dino-cifar100_risp_1e-05_wide_resnet28w2','dino-cifar100_so_1e-06_wide_resnet28w2','mocov2plus-cifar100_none_1e-05_wide_resnet28w2','mocov2plus-cifar100_risp_1e-05_wide_resnet28w2','mocov2plus-cifar100_so_1e-06_wide_resnet28w2',
'mocov3-cifar100_none_1e-05_wide_resnet28w2','mocov3-cifar100_risp_1e-05_wide_resnet28w2','mocov3-cifar100_so_1e-06_wide_resnet28w2','nnbyol-cifar100_none_0.001_wide_resnet28w2',
'nnbyol-cifar100_risp_1e-05_wide_resnet28w2','nnbyol-cifar100_so_1e-06_wide_resnet28w2']


print(len(resnet18s),len(resnet50s),len(wideresnets))

finetune_names = resnet18s+resnet50s+wideresnets
finetune_datasets = 'cifar100'


# l =['byol']
# base = l[0]
# for base in l:
#     for method in ['none','risp','so']:
       
#         # if method=='none' and base =='mocov2plus':
#         #     continue
#      #   cmd = 'python3 -u main_pretrain.py  --config-path scripts/pretrain/cifar --config-name {}.yaml ++regular_method={} ++regular_weight={}'.format(base,method,weight)
#         cmd = 'python3 -u main_linear.py  --config-path scripts/linear/imagenet-100 --config-name {}.yaml ++regular_method={} ++regular_weight={}'.format(base,method,weight)
#         print(cmd)
#         os.system(cmd)
cmd = 'python3 -u main_linear.py  --config-path scripts/linear/cifar100 --config-name byol.yaml ++regular_method={}'