"""
1. 准备数据
    - 批次
    - 打乱
2. 创建模型
3. 确定损失函数
4. 模型训练
5. 测试
"""
import os

import torch.optim
import tqdm

from dataset import FTCDataset
from torch.utils.data import DataLoader
from net import Net
from torch import nn
from torch.utils.tensorboard import SummaryWriter

best_weight_path = 'weights/best_v6.pt'
log_path = 'logs_v6'
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
epoch_num = 200


class Trainer:
    def __init__(self):
        super().__init__()
        # 1. 准备数据
        # 获取数据集
        train_set = FTCDataset(isTrain=True)
        test_set = FTCDataset(isTrain=False)
        # 数据加载器：批次处理 打乱顺序
        self.train_loader = DataLoader(dataset=train_set, batch_size=50, shuffle=True, num_workers=2)
        self.test_loader = DataLoader(dataset=test_set, batch_size=20, shuffle=True, num_workers=2)
        # 2. 创建模型
        net = Net()
        # 加载参数
        if os.path.exists(best_weight_path):
            net.load_state_dict(torch.load(best_weight_path))
        self.best_train_acc = 0
        # 切换设备处理数据
        net.to(device)
        self.net = net
        # 3. 确定损失函数
        # self.loss_fn = nn.MSELoss()
        self.loss_fn = nn.CrossEntropyLoss()
        # 优化器
        self.opt = torch.optim.Adam(net.parameters())
        # 可视化工具
        self.writer = SummaryWriter(log_path)

    def train(self, epoch):
        sum_loss = 0
        sum_acc = 0
        for img_vector, target in tqdm.tqdm(self.train_loader, desc='train', total=len(self.train_loader)):
            # 模型输出
            img_vector, target = img_vector.to(device), target.to(device)
            pred_out = self.net(img_vector)
            # 前向传播
            loss = self.loss_fn(pred_out, target)
            # 梯度清零
            self.opt.zero_grad()
            # 反向传播
            loss.backward()
            # 更新参数
            self.opt.step()
            sum_loss += loss.item()
            # 准确率统计
            pred_cls = torch.argmax(pred_out, dim=1)
            # 损失函数是交叉熵 target直接是索引
            # target_cls = torch.argmax(target, dim=1)
            acc = torch.mean(torch.eq(pred_cls, target).to(torch.float32))
            sum_acc += acc.item()

        train_avg_loss = sum_loss / len(self.train_loader)
        print(f'epoch:{epoch} train_avg_loss:{train_avg_loss}')
        train_avg_acc = sum_acc / len(self.train_loader)
        print(f'epoch:{epoch} train_avg_acc:{train_avg_acc}')
        # self.writer.add_scalar('train_avg_loss', train_avg_loss, epoch)
        # self.writer.add_scalar('train_avg_acc', train_avg_acc, epoch)
        self.writer.add_scalars('loss', {'train_avg_loss': train_avg_loss}, epoch)
        self.writer.add_scalars('acc', {'train_avg_acc': train_avg_acc}, epoch)
        # 保存参数 文件后缀 .pt 或 .pth
        if self.best_train_acc < train_avg_acc:
            self.best_train_acc = train_avg_acc
            torch.save(self.net.state_dict(), best_weight_path)

    def test(self, epoch):
        sum_loss = 0
        sum_acc = 0
        for img_vector, target in tqdm.tqdm(self.test_loader, desc='test', total=len(self.test_loader)):
            # 模型输出
            img_vector, target = img_vector.to(device), target.to(device)
            pred_out = self.net(img_vector)
            # 前向传播
            loss = self.loss_fn(pred_out, target)
            sum_loss += loss.item()
            # 准确率统计
            pred_cls = torch.argmax(pred_out, dim=1)
            # 损失函数是交叉熵 target直接是索引
            # target_cls = torch.argmax(target, dim=1)
            acc = torch.mean(torch.eq(pred_cls, target).to(torch.float32))
            sum_acc += acc.item()

        test_avg_loss = sum_loss / len(self.test_loader)
        print(f'epoch:{epoch} test_avg_loss:{test_avg_loss}')
        test_avg_acc = sum_acc / len(self.test_loader)
        print(f'epoch:{epoch} test_avg_acc:{test_avg_acc}')
        # self.writer.add_scalar('test_avg_loss', test_avg_loss, epoch)
        # self.writer.add_scalar('test_avg_acc', test_avg_acc, epoch)
        self.writer.add_scalars('loss', {'test_avg_loss': test_avg_loss}, epoch)
        self.writer.add_scalars('acc', {'test_avg_acc': test_avg_acc}, epoch)

    def run(self):
        for epoch in range(epoch_num):
            self.train(epoch)
            self.test(epoch)


if __name__ == '__main__':
    trainer = Trainer()
    trainer.run()
    pass