'''
author ：  Hatcher W
Email： hcwang@semi.ac.cn
use: 密度估计训练
'''
import torch
import torch.nn as nn
from torch import optim
from torch.utils.data import DataLoader
import numpy as np
from network import CSRNet,PKSNet,PKSNet_attention
from dataset import People_Dataset
from configs.config import *
from configs.config import model_head
import json
import os
from tqdm import tqdm
import cv2
os.environ['CUDA_VISIBLE_DEVICES'] = '3'
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
train_weigth_path = "../data_MPKFNet/weigth/BirdCounting2/"
if not os.path.exists(train_weigth_path):
    os.makedirs(train_weigth_path)
'''
    准备数据
'''
batch_size = 1
with open(trainA_json, 'r') as outfile:
    train_list = json.load(outfile)
with open(testA_json, 'r') as outfile:
    test_list = json.load(outfile)
train_dataset = People_Dataset(train_list)
train_loader = DataLoader(train_dataset, batch_size=batch_size)
test_dataset = People_Dataset(test_list)
test_loader = DataLoader(test_dataset, batch_size=1)

'''
    创建网络
'''
net = PKSNet().to(device)


'''
    训练设置
'''
lr=1e-10
criterion_desity = nn.MSELoss(reduction='sum')
criterion_spice = nn.CrossEntropyLoss()
criterion_posture = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.95, weight_decay=5*1e-4)
epoch_start = 0
# optimizer = optim.Adam(net.parameters(), lr=1e-7,  weight_decay=1e-4)
if True:
   # net.load_state_dict(torch.load('./weigth/bestmae.pth'))
    checkpoint = torch.load(train_weigth_path+'final.pth')
    net.load_state_dict(checkpoint['net'])
    #optimizer.load_state_dict(checkpoint['optimizer'])
    #epoch_start = checkpoint['epoch']
    print("load model weight ok")

'''
    开始训练
'''
mae_best = 1e10
num_lr=0
print(net)
for epoch in range(epoch_start,epochs):
    net.train()
    epoch_loss = 0
    with tqdm(total=len(train_dataset), desc=f'Epoch {epoch + 1}/{epochs}', unit='img') as pbar:
        if model_head:
            for img, target, count, posture, species in train_loader:
                img = img.to(device)
                target = target.to(device)
                posture = posture.to(device)
                species = species.to(device)
                target_pred, posture_pred, spice_pred = net(img)
                # print(target_pred.shape, posture_pred.shape, spice_pred.shape)
                # print(target.shape, posture.shape, species.shape)
                loss = 0.9*criterion_desity(target_pred, target) + 0.05*criterion_posture(posture_pred, posture) \
                       + 0.05*criterion_spice(spice_pred, species)
                epoch_loss += loss.item()
                pbar.set_postfix({'loss (batch)': loss.item()/img.shape[0]})
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                pbar.update(img.shape[0])
        else:
            for img, target, count in train_loader:
                img = img.to(device)
                target = target.to(device)
                target_pred= net(img)
                # print(target_pred.shape, posture_pred.shape, spice_pred.shape)
                # print(target.shape, posture.shape, species.shape)
                loss = criterion_desity(target_pred, target)
                epoch_loss += loss.item()
                pbar.set_postfix({'loss (batch)': loss.item() / img.shape[0]})
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                pbar.update(img.shape[0])
    print('train_data: {}, iter loss: {} for single image \n'.format(len(train_list), epoch_loss/len(train_list)))

    mae ,mse= 0,0
    if model_head:
        for img, traget, count, posture, species  in tqdm(test_loader):
            img = img.to(device)
            target = traget.to(device)
            targets_pred, posture_pred, spice_pred = net(img)
            #print(targets_pred.data.sum())
            #print(target.data.sum())
            mae += abs(targets_pred.data.sum() - target.data.sum())
            mse += ((target.data.sum() - targets_pred.data.sum()) * (target.data.sum() - targets_pred.data.sum()))
    else:
        for img, traget, count in tqdm(test_loader):
            img = img.to(device)
            target = traget.to(device)
            targets_pred = net(img)
            #print(targets_pred.data.sum())
            #print(target.data.sum())
            mae += abs(targets_pred.data.sum() - target.data.sum())
            mse += ((target.data.sum() - targets_pred.data.sum()) * (target.data.sum() - targets_pred.data.sum()))
    mae = mae / len(test_list)
    mse = mse /len(test_list)
    if mae < mae_best:
        num_lr=0
        print(epoch)
        state = {
            "epoch" :epoch,
            "net"   :net.state_dict(),
            "optimizer":optimizer.state_dict()

        }
        torch.save(state, train_weigth_path+'bestmaesheep{},mse{}.pth'.format(int(mae),int(mse)))
        torch.save(state, train_weigth_path+'bestmaesheep.pth'.format(int(mae)))
        mae_best = mae
    else:
        num_lr+=1
    if num_lr>5:
        lr*=0.1
        num_lr = 0
        if lr < 1e-10:
            lr = 1e-10
        print("lr is down to :"+str(lr))
        optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.95, weight_decay=5 * 1e-4)
    print('test_data:{},calculate MAE: \n'.format(len(test_list)))
    print(' * MAE {mae:.3f} \n'.format(mae=mae))
    print(' * MSE {mse:.3f} \n'.format(mse=mse))


'''
    模型加载与单张图片测试
'''
# net = CSRNet()
# net.load_state_dict(torch.load('CSRnet.pth'))
# net.eval()
# test_img = cv2.imread('./data/part_A_final/test_data/images/IMG_1.jpg.jpg')
# test_img = cv2.resize(test_img, (224, 224))
# test_img = np.transpose(test_img, [2, 0, 1])
# test_img = np.expand_dims(test_img,axis=0).astype(np.float)
# with torch.no_grad():
#     out = net(torch.from_numpy(test_img).float())
#     out = out.squeeze(0).numpy()
# cv2.imwrite('result.png',out)