# -*- coding: utf-8 -*-
"""
Created on Tue Dec  1 10:11:41 2020

@author: ALW
"""
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import model as ml
from torch import optim
import dataread
from torch.autograd import Variable
import time
import os
#import loss
import random
from rich.progress import track
batch_size = 400
batch_size1 = 100
lr = 1e-3
os.environ['CUDA_VISIBLE_DEVICES'] = '5,6'   ####gpu选择

net = ml.SKnet(1,5)
#net = torch.load("/data/nfs_rt16/luyuan/code/interspeech_classifition/ResNet_0207_augment/7_my_model.pth")#.module
#print(net)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  ###如果没有gpu选择cpu

if torch.cuda.device_count() > 1:
  net = nn.DataParallel(net)  ####gpu并行训练
# Assuming that we are on a CUDA machine, this should print a CUDA device:
net.to(device)  ####网络采用gpu
print(device)
net=net.double()  

criterion = torch.nn.CrossEntropyLoss()  #####CE准则loss
#criterion = torch.nn.SmoothL1Loss()
#criterion = loss.FocalLoss(5)
optimizer = optim.Adam(net.parameters(), lr=lr)  ########优化器
Train_list = np.loadtxt("./train.txt", dtype=str, delimiter=',')
devel_list = np.loadtxt("./test.txt", dtype=str, delimiter=',')
#Train_list = Train_list[1:]
#devel_list = devel_list[1:]
#devel_list1 = devel_list[1:500]
#devel_list = devel_list[500:]
#Train_list = np.concatenate([Train_list,devel_list],0)



all_time = 0          #####一共所用时间
num_epochs = 300   #######训练轮数

#print(net)
#net.load_state_dict(checkpoint)

start_epoch = 0
for epoch in range(start_epoch,num_epochs):
    start_time = time.time()
    permutation = np.random.permutation(Train_list.shape[0])####数据随机
    shuffled_dataset = Train_list[permutation, :]
    if epoch%5 == 0:   ####学习率衰减
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr
            lr = lr * 0.99
    net.train()
    for i in range(1):
        running_corrects = torch.zeros(1).squeeze().cuda()
        data_path = shuffled_dataset[i*batch_size:i*batch_size+batch_size]
        data,label = dataread.batch_data_read(data_path)
        print(data.shape,label.shape)
        data = data.reshape((-1,1,10))
        inputs1 = Variable(torch.from_numpy(data)).to(device)  #######Train
        target = Variable(torch.from_numpy(label)).to(device) 
        #target = target.long()
        out = net(inputs1)
        loss = criterion(out,target) #####out*inputs1.double()
        _ , prediction =  torch.max(out,1)
        running_corrects = torch.sum(prediction == target)
        optimizer.zero_grad()           #归零梯度，每次反向传播前都要归零梯度，不然梯度会累积，造成结果不收敛
        loss.backward()                 #反向传播
        optimizer.step()                #更新参数
        train_time = time.time() - start_time 
        print('Epoch[{}/{}],Process[{}/{}],loss:{:.6f},ACC:{:.6f},time:{:.3f},leanring_rate:{:.6f}'.format(epoch + 1, num_epochs, i + 1, int(len(Train_list)/batch_size)+1, loss.item(), running_corrects/len(label), train_time, param_group['lr']))
    running_corrects = torch.zeros(1).squeeze().cuda()
    total = torch.zeros(1).squeeze().cuda()
   
    
    net.eval()
    for i in range(1):
        data_path = devel_list[i*batch_size1:i*batch_size1+batch_size1]
        data,label = dataread.batch_data_read(data_path) 
        data = data.reshape((-1,1,10))
        inputs1 = Variable(torch.from_numpy(data)).to(device)  #######Train
        target = Variable(torch.from_numpy(label)).to(device)  #######输入phase*cos
        #target = target.long()
        out = net(inputs1)
        _ , prediction =  torch.max(out,1)
        running_corrects += torch.sum(prediction == target)
        total += len(label)  
    print('Accuracy: %f'%((running_corrects/total).cpu().detach().data.numpy()))
    #print('UAR: %f'%(UAR.cpu().detach().data.numpy()))
    
    
    #torch.save(net, "/data/nfs_rt16/luyuan/code/interspeech_classifition/SKnet_0224_noresize_0.4/"+str(epoch + 1)+"_my_model.pth")  # 保存模型
