# 数据需准备成以下格式
# ├── 数据集路径：data_path
#     └── image：存放所有图片
#     └── train.txt：训练图片的绝对路径(或相对data_path下路径)和类别号，(image/mask/0.jpg 0 2\n)表示该图片类别为0和2，空类别图片无类别号
#     └── val.txt：验证图片的绝对路径(或相对data_path下路径)和类别
#     └── class.txt：所有的类别名称
# class.txt：
# 类别1
# 类别2
# ...
# -------------------------------------------------------------------------------------------------------------------- #
# 分布式数据并行训练：
# python -m torch.distributed.launch --master_port 9999 --nproc_per_node n run.py --distributed True
# master_port为GPU之间的通讯端口，空闲的即可
# n为GPU数量
# -------------------------------------------------------------------------------------------------------------------- #
import os
import wandb
import torch
import argparse
from block.data_get import data_get
from block.loss_get import loss_get
from block.model_get import model_get
from block.train_get import train_get
from block.train_embeder import train_embeder

# 获取当前文件路径
pwd = os.getcwd()

# -------------------------------------------------------------------------------------------------------------------- #
# 模型加载/创建的优先级为：加载已有模型>创建剪枝模型>创建timm库模型>创建自定义模型
parser = argparse.ArgumentParser(description='|针对分类任务，添加水印机制，包含数据隐私、模型水印|')
parser.add_argument('--wandb', default=False, type=bool, help='|是否使用wandb可视化|')
parser.add_argument('--wandb_project', default='classification', type=str, help='|wandb项目名称|')
parser.add_argument('--wandb_name', default='train', type=str, help='|wandb项目中的训练名称|')
parser.add_argument('--wandb_image_num', default=16, type=int, help='|wandb保存图片的数量|')

# 训练时是否嵌入白盒水印，默认为False，即不嵌入白盒水印
parser.add_argument('--wm_embed', default=True, type=bool,
                    help='|训练时是否嵌入白盒水印，默认为False，即不嵌入白盒水印|')

# new_added
parser.add_argument('--data_path', default=f'{pwd}/dataset', type=str, help='Root path to datasets')
parser.add_argument('--dataset_name', default='CIFAR-10', type=str, help='Specific dataset name')
parser.add_argument('--input_channels', default=3, type=int)
parser.add_argument('--output_num', default=10, type=int)
# parser.add_argument('--input_size', default=32, type=int)
# 触发集标签定义，黑盒水印植入,这里需要调用它，用于处理部分数据的
parser.add_argument('--trigger_label', type=int, default=2,
                    help='The NO. of trigger label (int, range from 0 to 10, default: 2)')
# 设置数据集中添加水印的文件占比，这里可以直接选择水印控制,看看如何选择调用进来
parser.add_argument('--watermarking_portion', type=float, default=0.1,
                    help='poisoning portion (float, range from 0 to 1, default: 0.1)')

# 待修改
parser.add_argument('--input_size', default=32, type=int, help='|输入图片大小|')
# 待修改
parser.add_argument('--output_class', default=10, type=int, help='|输出的类别数|')
parser.add_argument('--weight', default='last.pt', type=str, help='|已有模型的位置，没找到模型会创建剪枝/新模型|')

# 剪枝的处理部分
parser.add_argument('--prune', default=False, type=bool, help='|模型剪枝后再训练(部分模型有)，需要提供prune_weight|')
parser.add_argument('--prune_weight', default='best.pt', type=str, help='|模型剪枝的参考模型，会创建剪枝模型和训练模型|')
parser.add_argument('--prune_ratio', default=0.5, type=float, help='|模型剪枝时的保留比例|')
parser.add_argument('--prune_save', default='prune_best.pt', type=str, help='|保存最佳模型，每轮还会保存prune_last.pt|')

# 模型处理的部分了
parser.add_argument('--timm', default=False, type=bool, help='|是否使用timm库创建模型|')
parser.add_argument('--model', default='Alexnet', type=str, help='|自定义模型选择，timm为True时为timm库中模型|')
parser.add_argument('--model_type', default='s', type=str, help='|自定义模型型号|')
parser.add_argument('--save_path', default=None, type=str, help='|保存最佳模型，除此之外每轮还会保存last.pt|')
parser.add_argument('--save_path_last', default=None, type=str, help='|保存最佳模型，除此之外每轮还会保存last.pt|')

# 训练控制
parser.add_argument('--epoch', default=20, type=int, help='|训练总轮数(包含之前已训练轮数)|')
parser.add_argument('--batch', default=600, type=int, help='|训练批量大小，分布式时为总批量|')
# parser.add_argument('--loss', default='bce', type=str, help='|损失函数|')
parser.add_argument('--loss', default='cross', type=str, help='|损失函数|')
parser.add_argument('--warmup_ratio', default=0.01, type=float, help='|预热训练步数占总步数比例，最少5步，基准为0.01|')
parser.add_argument('--lr_start', default=0.001, type=float, help='|初始学习率，adam算法，批量小时要减小，基准为0.001|')
parser.add_argument('--lr_end_ratio', default=0.01, type=float, help='|最终学习率=lr_end_ratio*lr_start，基准为0.01|')
parser.add_argument('--lr_end_epoch', default=100, type=int, help='|最终学习率达到的轮数，每一步都调整，余玄下降法|')
parser.add_argument('--regularization', default='L2', type=str, help='|正则化，有L2、None|')
parser.add_argument('--r_value', default=0.0005, type=float, help='|正则化权重系数，基准为0.0005|')
parser.add_argument('--device', default='cuda', type=str, help='|训练设备|')
parser.add_argument('--latch', default=True, type=bool, help='|模型和数据是否为锁存，True为锁存|')
parser.add_argument('--num_worker', default=0, type=int, help='|CPU处理数据的进程数，0只有一个主进程，一般为0、2、4、8|')
parser.add_argument('--ema', default=True, type=bool, help='|使用平均指数移动(EMA)调整参数|')
parser.add_argument('--amp', default=False, type=bool, help='|混合float16精度训练，CPU时不可用，出现nan可能与GPU有关|')
parser.add_argument('--noise', default=0.5, type=float, help='|训练数据加噪概率|')
parser.add_argument('--class_threshold', default=0.5, type=float, help='|计算指标时，大于阈值判定为图片有该类别|')
parser.add_argument('--distributed', default=False, type=bool, help='|单机多卡分布式训练，分布式训练时batch为总batch|')
parser.add_argument('--local_rank', default=0, type=int, help='|分布式训练使用命令后会自动传入的参数|')
args = parser.parse_args()
args.device_number = max(torch.cuda.device_count(), 1)  # 使用的GPU数，可能为CPU

if args.wm_embed:
    args.save_path = f'{pwd}/weight/{args.model}/watermark/best.pt' if args.save_path is None else args.save_path
    args.save_path_last = f'{pwd}/weight/{args.model}/watermark/last.pt' if args.save_path_last is None else args.save_path_last
else:
    args.save_path = f'{pwd}/weight/{args.model}/clean/best.pt' if args.save_path is None else args.save_path
    args.save_path_last = f'{pwd}/weight/{args.model}/clean/last.pt' if args.save_path_last is None else args.save_path_last
# 创建模型对应的检查点目录
# checkpoint_dir = os.path.join('./checkpoints', args.model)
# if not os.path.exists(checkpoint_dir):
#     os.makedirs(checkpoint_dir)
# print(f"{args.model}模型检查点保存路径已创建: {checkpoint_dir}")

# 创建模型保存权重目录
weight_dir = f'{pwd}/weight/{args.model}'
if not os.path.exists(weight_dir):
    os.makedirs(weight_dir)
print(f"{args.model}模型权重保存路径已创建: {weight_dir}")

# 为CPU设置随机种子
torch.manual_seed(999)
# 为所有GPU设置随机种子
torch.cuda.manual_seed_all(999)
# 固定每次返回的卷积算法
torch.backends.cudnn.deterministic = True
# cuDNN使用非确定性算法
torch.backends.cudnn.enabled = True
# 训练前cuDNN会先搜寻每个卷积层最适合实现它的卷积算法，加速运行；但对于复杂变化的输入数据，可能会有过长的搜寻时间，对于训练比较快的网络建议设为False
torch.backends.cudnn.benchmark = False
# -------------------------------------------------------------------------------------------------------------------- #
# wandb可视化:https://wandb.ai
# if args.wandb and args.local_rank == 0:  # 分布式时只记录一次wandb
#     args.wandb_run = wandb.init(project=args.wandb_project, name=args.wandb_name, config=args)
# -------------------------------------------------------------------------------------------------------------------- #
# 混合float16精度训练
if args.amp:
    args.amp = torch.cuda.amp.GradScaler()
# 分布式训练
if args.distributed:
    torch.distributed.init_process_group(backend='nccl')  # 分布式训练初始化
    args.device = torch.device("cuda", args.local_rank)
# -------------------------------------------------------------------------------------------------------------------- #

# 判定数据集文件夹内信息是否齐全
if args.local_rank == 0:
    print(f'| args:{args} |')
    assert os.path.exists(f'{args.data_path}/{args.dataset_name}'), '! data_path中缺少:{args.dataset_name} !'
    assert os.path.exists(f'{args.data_path}/{args.dataset_name}/train.txt'), '! data_path中缺少:train.txt !'
    assert os.path.exists(f'{args.data_path}/{args.dataset_name}/test.txt'), '! data_path中缺少:test.txt !'
    assert os.path.exists(f'{args.data_path}/{args.dataset_name}/class.txt'), '! data_path中缺少:class.txt !'
    if os.path.exists(args.weight):  # 优先加载已有模型args.weight继续训练
        print(f'| 加载已有模型:{args.weight} |')
    elif args.prune:
        print(f'| 加载模型+剪枝训练:{args.prune_weight} |')
    # elif args.timm:  # 创建timm库中模型args.timm
    #     import timm
    #
    #     assert timm.list_models(args.model), f'! timm中没有模型:{args.model}，使用timm.list_models()查看所有模型 !'
    #     print(f'| 创建timm库中模型:{args.model} |')
    else:  # 创建自定义模型args.model
        assert os.path.exists(f'model/{args.model}.py'), f'! 没有自定义模型:{args.model} !'
        print(f'| 创建自定义模型:{args.model} | 型号:{args.model_type} |')
# -------------------------------------------------------------------------------------------------------------------- #
if __name__ == '__main__':
    # 摘要
    print(f'| args:{args} |') if args.local_rank == 0 else None
    # 数据
    data_dict = data_get(args)
    # 模型
    model_dict = model_get(args)
    # 损失
    loss = loss_get(args)
    # 训练
    if args.wm_embed:
        train_embeder(args, data_dict, model_dict, loss)
    else:
        train_get(args, data_dict, model_dict, loss)