# train_dataloader TODO
def generate_train_dataloader(batch_size_default=None, num_workers_default=None):
    batch_size = None
    if batch_size_default != None:
        while True:
            user_input = input(f"请输入 batch size (默认为 {batch_size_default}): ")
            
            # 如果用户没有输入内容，使用默认值
            if not user_input.strip():
                batch_size = batch_size_default
                break
            
            # 尝试将输入转换为整数
            try:
                batch_size = int(user_input)
                if batch_size > 0:
                    break  # 输入正确，退出循环
                else:
                    print("Batch size 必须是正整数，请重新输入。")
            except ValueError:
                print("输入无效，请输入一个有效的整数。")
        print(f"将train_dataloader的batch_size设置为{batch_size}")
    
    num_workers = None
    if num_workers_default != None:
        while True:
            user_input = input(f"请输入 num workers (默认为 {num_workers_default}): ")
            
            # 如果用户没有输入内容，使用默认值
            if not user_input.strip():
                num_workers = num_workers_default
                break
            
            # 尝试将输入转换为整数
            try:
                num_workers = int(user_input)
                if num_workers > 0:
                    break  # 输入正确，退出循环
                else:
                    print("Num workers 必须是正整数，请重新输入。")
            except ValueError:
                print("输入无效，请输入一个有效的整数。")
        print(f"将train_dataloader的num_workers设置为{num_workers}")
    
    # 返回包含 batch_size 的字典
    train_dataloader = {}
    if num_workers == None:
        train_dataloader['batch_size'] = batch_size
        return train_dataloader, batch_size
    
    if batch_size == None:
        train_dataloader['num_workers'] = num_workers
        return train_dataloader, num_workers
    
    train_dataloader['batch_size'] = batch_size
    train_dataloader['num_workers'] = num_workers
    return train_dataloader, batch_size, num_workers

    
