import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch
from sklearn.model_selection import train_test_split
from torch import nn
from torch.autograd import Variable
from torch.utils.data import TensorDataset
from sklearn.metrics import f1_score, accuracy_score,roc_auc_score,confusion_matrix
# torch.set_default_tensor_type(torch.DoubleTensor)

df = pd.read_csv('C:/data.csv')
X=df.drop(['label'][0], axis=1)
X=X.values
Y=df[['label'][0]]
Y=Y.values
x_train, x_test, y_train, y_test = train_test_split(X, Y,train_size=0.8, random_state=1)

train_x = torch.from_numpy(x_train)
train_y = torch.from_numpy(y_train)
#inputsize代表特征维度
class lstm(nn.Module):
    def __init__(self, input_size=180, hidden_size=4, output_size=1, num_layer=2):
        super(lstm, self).__init__()
        self.layer1 = nn.LSTM(input_size, hidden_size, num_layer)
        self.layer2 = nn.Linear(hidden_size, output_size)
        #self.layer3 = nn.Softmax(dim = 0)

    def forward(self, x):
        x, _ = self.layer1(x)
        s, b, h = x.size()
        x = x.view(s * b, h)
        x = self.layer2(x)
        x = x.view(s, b, -1)
        #x = self.layer3(x)
        return x


model = lstm(180, 4, 1, 2)
criterion = nn.MSELoss()
#如果用二分类交叉熵，则要进行归一化
#criterion = nn.BCELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
# 开始训练

train_x = train_x.reshape(-1, 1, 180).to(torch.float32)
train_y = train_y.reshape(-1, 1, 1).to(torch.float32)
for e in range(2000):
    var_x = Variable(train_x)
    var_y = Variable(train_y)
    # 前向传播
    out = model(var_x)
    loss = criterion(out, var_y)
    #
    # loss=loss.clone().detach().requires_grad_()
    # 反向传播
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if (e + 1) % 100 == 0:  # 每 100 次输出结果
        print('Epoch: {}, Loss: {:.5f}'.format(e + 1, loss.item()))

torch.save(model.state_dict(), 'lstm.pth')

model.load_state_dict(torch.load('lstm'))

model = model.eval（)  # 转换成测试模式

data_X = x_test.reshape(-1, 1, 180)
data_X = torch.from_numpy(data_X).to(torch.float32)
var_data = Variable(data_X)
pred_test = model(var_data)  # 测试集的预测结果
pred_test = pred_test.view(-1).data.numpy()

# 改变输出的格式
auc = roc_auc_score(y_test, pred_test)
print(auc)