import torch
import matplotlib.pyplot as plt #画图
import matplotlib
import numpy as np  #矩阵相关
import csv     #csv文件
import pandas  #csv文件
from torch.utils.data import Dataset, DataLoader
import torch.nn as nn
from torch import optim
import time

#定义COVID数据集类,继承自PyTorch的Dataset类
class CovidDataset(Dataset):
    def __init__(self, file_path, mode):    #重新自定义初始化
        with open(file_path, "r") as f:     #用读的方式打开文件
            ori_data = list(csv.reader(f))   #读取文件并将其转化为列表
            csv_data = np.array(ori_data)[1:, 1:].astype(float)        #将数据列表转化为矩阵不要第一行和第一列,并且将里面的内容转化为float
            #逢五取1，不推荐,一般训练集,验证集,测试集都是随机选取的
            if mode == "train":           #根据模式选择不同的数据索引
                indices = [i for i in range(len(csv_data)) if i % 5 != 0]

            elif mode == "val":
                indices = [i for i in range(len(csv_data)) if i % 5 == 0]

            elif mode == "test":
                indices = [i for i in range(len(csv_data))]

            X = torch.tensor(csv_data[indices, :93])        #获取特征数据的前93列,并将其转化为张量放在张量网上
            if mode != "test":    #如果是训练集或者验证集,获取标签数据(最后一列)
                self.Y = torch.tensor(csv_data[indices, -1])    # #获取特征数据的最后一列,并将其转化为张量放在张量网上
                #对特征数据进行标准化处理
            self.X = (X - X.mean(dim=0, keepdim=True)) / X.std(dim=0, keepdim=True)   #dim表示维度,0表示行,1表示列,2表示深度等
                                                                                      #keepdim表示维持维度
            self.mode = mode

    def __getitem__(self, item):
        #获得单个样本,item为样本索引
        if self.mode == "test":
            return self.X[item].float()   #即使我们使用.astype(float)将NumPy数组转化为浮点数,但是将其转化为Pytorch张量
        #时，数据类型确实有可能变成整数类型,所以我们要明确返回张量类型是浮点类型，以避免在神经网络计算中可能出现的问题(确保数据类型一致性,提高代码
        # 健壮性,优化性能,清洗意图)
        else:
            return self.X[item].float(), self.Y[item].float()

    def __len__(self):
        return len(self.X)    #返回数据集大小


#定义神经网络模型
class myModel(nn.Module):
    def __init__(self, inDim):  #初始化模型,inDim是输入特征维度
        super(myModel, self).__init__()   #调用父类(nn.Module)的构造函数,确保正确初始化
        self.fc1 = nn.Linear(inDim, 128)  #如果不重写自定义,这行可能无法正常运行
     #第一层全连接层(输入维度--->128维)
        self.relu1 = nn.ReLU()
     #ReLU激活函数
        self.fc2 = nn.Linear(128, 1)
     #第二层全连接层(128维--->1维(输出层))

    def forward(self, x):   #前向传播,输入x
        x = self.fc1(x)
        x = self.relu1(x)
        x = self.fc2(x)

        if len(x.size()) > 1:
            x = x.squeeze(1)       #如果维度大于1， 就去掉第二个维度(将形状从[batch_size,1]变成[batch_size])
        return x

#训练和验证函数
def train_val(model, train_loader, val_loader,  lr, optimizer,   device,        epochs, save_path):
          #要训练的模型 训练数据加载器  验证数据加载器 学习率  优化器   计算设备(CPU或GPU)  训练轮数  模型保存路径
    model = model.to(device)    #防止意外,将模型移动到指定设备上
    plt_train_loss = []  #总训练loss
    plt_val_loss = []    #总验证loss
    min_val_loss = 999999999999999999.9  #初始化最小验证损失

    for epoch in range(epochs):          #发枪指令， 冲锋的号角， 模型训练的开始
  #模型转化为训练模式,model.train()是pytorch中nn.Module类自带的方法,当我们调用时,会将模型设置成训练模型,这会影响到一些在训练和评估时
  #行为不同的层
        model.train()
        start_time = time.time()  #记录此时本轮开始的时间
        train_loss = 0.0  #浮点形式
        for x, y in train_loader:  #遍历训练数据批次
            x, y = x.to(device), y.to(device)  #将x和y都放在设备上
            y_pred = model(x)    #前向传播
     # pytorch的特殊方法_call_实现了包装,所以当写model(x)时,实际上调用了:
       # model._call_(x)--->model.forward(x)
            bat_loss = loss(y_pred, y, model)  #计算批次损失
            bat_loss.backward()       #反向传播计算偏导
            optimizer.step()         #更新参数
            optimizer.zero_grad()     #清空梯度,保证不影响下一批次的损失
            train_loss += bat_loss.cpu().item()   #累加损失
           #计算平均训练损失并且加入总训练loss列表
        plt_train_loss.append(train_loss/train_loader.__len__())


            #验证阶段
        model.eval()      #转化为验证模式
        val_loss = 0.0     #验证损失
        with torch.no_grad():      #不需要计算梯度,节省内存
            for val_x, val_y in val_loader:
                val_x, val_y = val_x.to(device), val_y.to(device)   #都放在设备上
                val_pred_y = model(val_x)     #计算验证的pre_y
                val_bat_loss = loss(val_pred_y, val_y, model)   #计算批次损失
                val_loss += val_bat_loss.cpu().item()      #累加损失

       # 计算平均训练损失并且加入总验证loss列表
        plt_val_loss.append(val_loss / val_loader.__len__())

        #保存最佳模型
        if val_loss < min_val_loss:
            min_val_loss = val_loss
            torch.save(model, save_path)

       #打印训练信息
        print("[%03d/%03d] %2.2f sec(s)  train_loss: %.6f val_loss:%.6f" % \
              (epoch,   epochs, time.time()-start_time, plt_train_loss[-1], plt_val_loss[-1]))
             #当前的轮次   总轮次        本轮此训练花费时间         训练集中真实的y        预测的y

#绘制损失曲线
    plt.plot(plt_train_loss)
    plt.plot(plt_val_loss)
    plt.title("loss")
    plt.legend(["train", "val"])
    plt.show()


#评估函数
def evaluate(model_path, test_loader,  rel_path,   device):
         #   模型文件路径  测试数据加载器   结果保存路径   计算设备

    model = torch.load(model_path).to(device)        #加载训练好的模型并且将它放到设备上

    rel = []         #记录预测结果
    model.eval()    #将模型设置为评估模式

     #使用测试数据进行预测
    with torch.no_grad():    #评估模式不需要计算梯度
        for x in test_loader:
            x = x.to(device)    #将x放到设备上
            pred = model(x)     #得到预测的pred_y
            rel.append(pred.cpu().item())  #先将此时的结果从设备或者张量网上移下来放到cpu上，将他的数据内容放入列表

        # 将预测结果保存到CSV文件
    with open(rel_path, "w", newline="") as f:  #用写的方式打开文件,并且中间没有间隔
        csv_writer = csv.writer(f)          #定义一个写指针
        csv_writer.writerow(["id", "tested_positive"])     #写入第一行内容
        for i, pred in enumerate(rel):     #同时得到索引和预测值
            csv_writer.writerow([str(i), str(pred)])
        print("结果保存到了"+rel_path)

#数据文件路径
train_file = r"D:\项目深度学习\第三节，回归实战代码\regression\covid\covid.train.csv"
test_file = r"D:\项目深度学习\第三节，回归实战代码\regression\covid\covid.test.csv"

batch_size = 16   #设置批次大小

#创建数据集实例
train_set = CovidDataset(train_file, "train")
val_set = CovidDataset(train_file, "val")
test_set = CovidDataset(test_file, "test")


#创建数据加载器,数据加载器能够批处理,打乱数据,并行加载,自定义数据加载顺序,自动将数据转化为张量,内存管理
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_set, batch_size=1, shuffle=False)


#自定义损失函数(带L2正则化)  正则化能够有效的防止过拟合,使曲线更平滑
def mseLoss(pred,  target, model):
           #预测值   真实值   模型
    loss = nn.MSELoss(reduction='mean')    #平方差损失
    ''' Calculate loss '''
    regularization_loss = 0                    # 正则项
    for param in model.parameters():
        # TODO: you may implement L1/L2 regularization here
        # 使用L2正则项
        # regularization_loss += torch.sum(abs(param))
        regularization_loss += torch.sum(param ** 2)                  # 计算所有参数平方
            #总损失 = MSE损失 + 正则化项
    return loss(pred, target) + 0.00075 * regularization_loss             # 返回损失。

#训练参数设置
loss = mseLoss   #使用自定义损失函数
epochs = 20   #运行轮次
lr = 0.001     #学习率
device = "cuda" if torch.cuda.is_available() else "cpu"   #自动选择设备
print(device)


#模型初始化
data_dim = 93     #特征维度
model = myModel(data_dim).to(device)    #模型实例化后放到设备上
save_path = "model_save/best_model.pth"    #模型的保存路径
rel_path = "pred.csv"           #预测结果保存路径

#使用SGD优化器,带动量
optimizer = optim.SGD(params=model.parameters(), lr=lr, momentum=0.9)

#开始训练和验证
train_val(model, train_loader, val_loader, lr, optimizer, device, epochs, save_path)

#使用训练好的模型进行预测
evaluate(save_path, test_loader, rel_path, device)