import os
import time

import torch
import argparse
from torchvision import transforms

from block.dataset_get import CustomDataset

# -------------------------------------------------------------------------------------------------------------------- #
parser = argparse.ArgumentParser(description='|pt模型推理|')
parser.add_argument('--model_path', default='./checkpoints/Alexnet/best.pt', type=str, help='|pt模型位置|')
parser.add_argument('--data_path', default='./dataset/CIFAR-10/test_cifar10_JPG', type=str, help='|验证集文件夹位置|')
parser.add_argument('--input_size', default=32, type=int, help='|模型输入图片大小|')
parser.add_argument('--batch', default=200, type=int, help='|输入图片批量|')
parser.add_argument('--device', default='cuda', type=str, help='|推理设备|')
parser.add_argument('--num_worker', default=0, type=int, help='|CPU处理数据的进程数，0只有一个主进程，一般为0、2、4、8|')
parser.add_argument('--float16', default=False, type=bool, help='|推理数据类型，要支持float16的GPU，False时为float32|')
args, _ = parser.parse_known_args()  # 防止传入参数冲突，替代args = parser.parse_args()
# -------------------------------------------------------------------------------------------------------------------- #
assert os.path.exists(args.model_path), f'! model_path不存在:{args.model_path} !'
assert os.path.exists(args.data_path), f'! data_path不存在:{args.data_path} !'
if args.float16:
    assert torch.cuda.is_available(), 'cuda不可用，因此无法使用float16'


# -------------------------------------------------------------------------------------------------------------------- #
def predict_pt(args):
    # 加载模型
    model_dict = torch.load(args.model_path, map_location='cpu')
    model = model_dict['model']
    model.half().eval().to(args.device) if args.float16 else model.float().eval().to(args.device)
    epoch = model_dict['epoch_finished']
    accuracy = round(model_dict['standard'], 4)
    print(f'| 模型加载成功:{args.model_path} | epoch:{epoch} | accuracy:{accuracy} |')
    # 推理
    start_time = time.time()
    with torch.no_grad():
        print(f"加载测试集至内存...")
        transform = transforms.Compose([
            transforms.ToTensor(),  # 将图像转换为PyTorch张量
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])  # 标准化
        ])
        dataset = CustomDataset(data_dir=args.data_path, image_size=(args.input_size, args.input_size), transform=transform)
        dataloader = torch.utils.data.DataLoader(dataset=dataset, batch_size=args.batch,
                                                 shuffle=False, drop_last=False, pin_memory=False,
                                                 num_workers=args.num_worker)
        print(f"加载测试集完成，开始预测...")
        correct = 0
        total = 0
        epoch = 0
        for index, (image_batch, true_batch) in enumerate(dataloader):
            image_batch = image_batch.to(args.device)
            pred_batch = model(image_batch).detach().cpu()
            # 获取指标项
            _, predicted = torch.max(pred_batch, 1)
            total += true_batch.size(0)
            correct += (predicted == true_batch).sum().item()
            epoch = epoch + 1
        # 计算指标
        accuracy = correct / total
        end_time = time.time()
        print(f'\n| 验证 | accuracy:{accuracy:.4f} | 图片总数:{total} | 每张耗时:{(end_time - start_time) / total} ')


if __name__ == '__main__':
    predict_pt(args)