import matplotlib.pyplot as plt
import torch #框架
import numpy as np #矩阵处理
import csv #读excel文件
from torch.utils.data import DataLoader, Dataset #两个与数据处理相关的包，类Dataset
import torch.nn as nn #类nn.Module需要用，损失函数需要用
from torch import optim #引入优化器
import time #记录时间
# 设置显示中文字体（可视化时图片有中文报错）
from pylab import mpl
mpl.rcParams["font.sans-serif"] = ["SimHei"]
#设置显示中文字体⬆
# 数据处理部分
# # 文件路径赋予
train_file = "covid.train.csv"
test_file = "covid.test.csv"
#
# file = pd.read_csv(train_file)
# print(file.head()) #输出数据前5行
class CovidDataset(Dataset): #dataset基类,CovidDataset子类
#init函数完成的是，从一个路径转入文件开始，得到想要的x（self.data），和y
def __init__(self, file_path, mode): #初始化,mode:对于训练集验证集测试集来说,他们的操作是不一样的,所以要用mode区分他们是什么集
with open(file_path, "r") as f: #打开文件,只读,文件名为f; with:在进入代码块之前自动打开文件，并确保在退出代码块后自动关闭文件
ori_data = list(csv.reader(f)) #原始数据
#要去掉原始数据的第一行和第一列,但列表是没法取一个方格的,所以要转换成矩阵
csv_data = np.array(ori_data)[1:, 1:].astype(float) #转换成矩阵,去掉第一行第一列,astype(float)转换为浮点型(本来是字符串数据)
# 逢5取1（行）
if mode == "train":
indices = [i for i in range(len(csv_data)) if i % 5 != 0]
self.y = torch.tensor(csv_data[indices, -1]) #对训练集来说y是最后一列，转换成张量，加self是为了让其在这个类中的每个函数都可使用
data = torch.tensor(csv_data[indices, :-1]) #真正需要的数据，并转换成张量（想进入神经网络模型必须转换成张量）
elif mode == "val":
indices = [i for i in range(len(csv_data)) if i % 5 == 0]
self.y = torch.tensor(csv_data[indices, -1]) #对验证集来说y是最后一列，转换成张量
data = torch.tensor(csv_data[indices, :-1]) #真正需要的数据，并转换成张量（想进入神经网络模型必须转换成张量）
else:
indices = [i for i in range(len(csv_data))] #测试集取所有,测试集用的是另一个文件的数据
data = torch.tensor(csv_data[indices, :]) #对测试集来说，不需要去掉最后一列
#数据须归一化（减去平均值，除以标准差），因为每列数据之间可能差距过大，比如一列都是零点几，另一列都是几十（量纲不一样），所以数据需要调整，调整后不影响训练，因为每个样本（行）都是列内比较，不会列间比较
self.data = (data - data.mean(dim=0, keepdim=True))/data.std(dim=0, keepdim=True) #mean()平均值，dim=0维度为0，keepdim=True保持维度不变std()标准差
self.mode = mode #为了让类的所有函数都知道传入的文件是什么mode（训练，验证，测试）
#传入一个下标，传出数据
def __getitem__(self, idx):
if self.mode == "test": #测试集
return self.data[idx].float() #后面加.float()把数据转换成32位的，节省训练资源
else: #验证集/训练集
return self.data[idx].float(), self.y[idx].float()
#回传data长度
def __len__(self):
return len(self.data)
train_dataset = CovidDataset(train_file, "train")
val_dataset = CovidDataset(train_file, "val") #验证集用的也是训练集文件的数据,只是训练集和验证集取的数据行数不同
test_dataset = CovidDataset(test_file, "test")
# 测试
# for data in train_dataset:
# print(data)
#DaraLoader
batch_size = 16
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) #(数据集，一批数量，数据打乱)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False) #测试集把batchsize改成1，测试集不能打乱，因为x，y要一一对应
#测试
# for batch_x, batch_y in train_loader:
# print(batch_x, batch_y)
# 至此完成数据处理部分。⬆
# 模型部分
class Mymodel(nn.Module):
def __init__(self, inDim): #输出维度已经确定为1，输入维度inDim需要传入， 模型框架规定
super(Mymodel, self).__init__()
self.fc1 = nn.Linear(inDim, 64) #fc1表示全连接层第一层 （输入维度，输出维度）输出维度自己随便指定
self.relu1 = nn.ReLU() #激活函数relu
self.fc2 = nn.Linear(64, 1) #fc2全连接第二层
def forward(self, x): #模型前向过程
x = self.fc1(x)
x = self.relu1(x)
x = self.fc2(x)
#经过调试发现，return的x即pred_y维度是（16，1），而batch_y维度是（16，），1代表2维，所以如果x维度大于1，要减一个维度
if len(x.size()) > 1: #size()维度
return x.squeeze(1) #给x减一个维度
return x
# 至此完成模型部分。⬆
# 超参部分
device = "cuda" if torch.cuda.is_available() else "cpu" #如果显卡可用用显卡，否则用cpu
#把超参数放在字典中，可以随时调用
config = {
"lr" : 0.001,
"epochs": 20,
"momentum": 0.9, #动量
"save_path": "model_save/best_model.pth", #保存路径，保存效果最好的模型
"res_path": "pred.csv"
}
model = Mymodel(inDim=93).to(device) #(Mymodel类的init函数的输入inDim) to(device)把模型放在gpu上
loss = nn.MSELoss() #损失函数,平方损失（相减再平方）
optimizer = optim.SGD(model.parameters(), lr=config["lr"], momentum=config["momentum"]) #优化器（随机梯度下降） （模型参数，学习率，动量）
#训练流程
def train_val(model, train_loader, val_loader, device, epochs, optimizer, loss, save_path): #(模型，训练集，验证集，设备，迭代轮数，优化器，损失函数，最好模型的保存路径)
model = model.to(device)
plt_train_loss = [] #存储训练过程中所有的loss，对于回归任务来说，只有loss值能判断模型的好坏
plt_val_loss = [] #存储验证过程中所有的loss，命名前面加了plt，表示画图用
min_val_loss = 9999999999999 #记录最小的loss值，为了之后保存最好的模型
#训练开始~~~~~~~~~~~！！！！！！！！！！
for epoch in range(epochs): #冲锋号角
train_loss = 0.0 #记录每一轮的loss
val_loss = 0.0
start_time = time.time() #记录训练开始时间,time.time()是当前时间的意思
model.train() #模型调为训练模式（对于模型来说，训练和测试是不一样的，训练时会关闭部分节点，测试时会开启所有节点）
for batch_x, batch_y in train_loader:
x, target = batch_x.to(device), batch_y.to(device) #先把x，y数据放设备上
pred = model(x) #数据x通过模型得到预测值
train_batch_loss = loss(pred, target) #计算loss, (y尖，y)
train_batch_loss.backward() #梯度回传
optimizer.step() #更新模型，起到了我mylinear.py代码中sgd([w_0, b_0], lr)的作用
optimizer.zero_grad() #更新模型后（梯度已经使用过），清零梯度
#所有训练步骤完成⬆
train_loss += train_batch_loss.cpu().item() #每轮迭代loss相加。.cpu()把train_batch_loss放在cpu上，因为本来放在gpu上，没法相加。.item()把数值取出来
plt_train_loss.append(train_loss / train_loader.__len__()) #记录这一轮的loss。train_loader.__len__()表示train_loader有多少批次。train_loss / train_loader.__len__()：因为train_loss是该轮所有批次的loss相加，所以要除以该轮的批次数
#每轮训练完后验证
model.eval() #调整为验证模式
with torch.no_grad(): #张量网上的所有计算都会记录下来，但是验证的时候不能更新模型，所以要这一句。属于这句代码的部分不计算梯度。
for batch_x, batch_y in val_loader:
x, target = batch_x.to(device), batch_y.to(device)
pred = model(x)
val_batch_loss = loss(pred, target) #验证的loss反映模型训练后的表现如何
val_loss += val_batch_loss.cpu().item()
plt_val_loss.append(val_loss / val_loader.__len__()) #append()加在列表最后
#若模型优秀，保存模型
if val_loss < min_val_loss:
min_val_loss = val_loss
torch.save(model, save_path) #保存模型（模型，位置）
#打印每轮训练情况
print("[%03d/%03d] %2.2f sec(s) Trainloss: %.6f|Valloss: %.6f"%(epoch, epochs, time.time() - start_time, plt_train_loss[-1], plt_val_loss[-1])) # %格式化输出
# 可视化
plt.plot(plt_train_loss) #画loss的变化
plt.plot(plt_val_loss)
plt.title("loss图")
plt.legend(["train", "val"]) #图例
plt.show()
#得出测试结果文件
def evaluate(save_path, test_loader, device, res_path):
model = torch.load(save_path).to(device) #加载模型
res = [] #记录结果
with torch.no_grad(): #f(x)=y
for x in test_loader:
x = x.to(device)
y = model(x)
res.append(y.cpu().item()) #把值从gpu上取下
print(res)
with open(res_path, "w", newline='') as f: #打开文件,写,文件名为f。每次写完会换行，和下面的writerow加在一起就会换两行，所以要加newline=‘’表示不换行
csvWriter = csv.writer(f) #定义 写指针，csv.writer(f)往f里写
csvWriter.writerow(["id","tested_positive"]) #writerow写一行,writerow会自动换行
# for i in range(len(res)): 或者：
for i, value in enumerate(res): #取列表res里的下标和值
csvWriter.writerow([str(i), str(value)]) #写一行（下标，值）str转换成字符串
print("文件已经保存到{}".format(res_path)) # {}格式化输出，format()里的东西替换大括号
# 或者： print("文件已经保存到“+res_path) python的字符串可以直接相加
#模型训练
train_val(model, train_loader, val_loader, device, config["epochs"], optimizer, loss, config["save_path"]) #训练，验证
#模型评价（测试）
evaluate(config["save_path"], test_loader, device, config["res_path"])
