# 导包
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import torch
from torch.utils.data import TensorDataset, DataLoader
import numpy as np
import time


# 获取数据并切割数据
def step1_get_data(datapath):
    # 读取原始数据
    df = pd.read_csv(datapath, sep=',')
    # print(df.head())
    # print(df.shape)
    # 数据预处理
    df = df.dropna(how='any')
    # print(df.shape)
    # 分别获取特征x和标签y
    x = df.iloc[:, 0:-1]
    y = df.iloc[:, -1]
    # TODO 提前把x转换为浮点数
    x = x.astype(np.float32)
    y = y.astype(np.int64)
    # print(x.shape, y.shape)
    # 切割数据
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=666)
    # TODO 优化1:特征标准化
    ss = StandardScaler()
    x_train = ss.fit_transform(x_train)
    x_test = ss.transform(x_test)
    return x_train, x_test, y_train, y_test


# 构建数据加载器
def step2_get_data_loader(batch_size, x_train, x_test, y_train, y_test):
    # 先构建dataset数据集:  把特征和标签打包构建完整数据集   注意: 需要先把pandas的df数据转换为张量
    train_dataset = TensorDataset(torch.from_numpy(x_train), torch.from_numpy(y_train.values))
    test_dataset = TensorDataset(torch.from_numpy(x_test), torch.from_numpy(y_test.values))
    # 再构建data_loader数据加载器  把完整数据集进行分批操作
    train_data_loder = DataLoader(train_dataset, batch_size, shuffle=True)
    test_data_loder = DataLoader(test_dataset, batch_size, shuffle=False)
    # 返回数据加载器
    return train_data_loder, test_data_loder


# 构建模型
class PhoneModel(torch.nn.Module):
    # 重写init魔法方法,创建模型的时候自动调用
    def __init__(self, input_num, output_num):
        super().__init__()
        self.linear1 = torch.nn.Linear(input_num, 128)
        self.linear2 = torch.nn.Linear(128, 256)
        # TODO 优化2:增加网络深度
        self.linear3 = torch.nn.Linear(256, 256)
        self.linear4 = torch.nn.Linear(256, 256)
        self.out = torch.nn.Linear(256, output_num)

    # 重写forward方法,使用模型的时候自动调用
    def forward(self, x):
        # print(f"进来一批数据,形状为:{x.shape}")  # (批次,特征数)
        # 加权求和+激活函数
        x = torch.relu(self.linear1(x))
        x = torch.relu(self.linear2(x))
        x = torch.relu(self.linear3(x))
        x = torch.relu(self.linear4(x))
        # 因为后续我要用多分类交叉熵损失函数,所以这里不加softmax
        x = self.out(x)
        # 返回结果
        return x


# TODO 模型训练
def train_model(data_loader, model, lr, modelpath):
    # 设置模型为训练模式
    model.train()
    # todo 1.准备数据(此处已经传入)
    # todo 2.准备模型(此处已经传入)
    # todo 3.准备损失函数
    loss_fn = torch.nn.CrossEntropyLoss(reduction='mean')
    # todo 4.准备优化器
    optimizer = torch.optim.Adam(lr=lr, betas=(0.9, 0.999), params=model.parameters())
    # todo 5.外层循环控制轮次
    for epoch in range(1, epochs + 1):
        # 为了打印日志, 记录总损失和轮次和时间
        total_loss, batch_cnt, start_time = 0, 0, time.time()
        # todo 6.内层循环控制训练批次
        for batch_x, batch_y in data_loader:
            # todo 7.前向传播
            logits_y = model(batch_x)
            # todo 8.计算损失
            loss = loss_fn(logits_y, batch_y)
            total_loss += loss.item()
            batch_cnt += 1
            # todo 9.梯度清零
            optimizer.zero_grad()
            # todo 10.反向传播
            loss.backward()
            # todo 11.参数更新
            optimizer.step()
        # 打印日志
        print(f"第{epoch}轮,平均损失:{total_loss / batch_cnt:.6f},用时:{time.time() - start_time:.4f}秒")
    # todo 12.保存模型
    torch.save(model.state_dict(), modelpath)


# TODO 模型评估
def eval_model(data_loader, modelpath):
    # todo 1.准备数据(此处已经传参)
    # todo 2.准备模型(创建一个新模型,加载已经训练的参数)
    model = PhoneModel(input_num, output_num)
    model.load_state_dict(torch.load(modelpath))
    # todo 设置模型为评估模式
    model.eval()  # 禁用dropout,同时改变bn计算方式
    # todo 禁用梯度计算以减少内存和加速
    with torch.no_grad():
        # todo 3.遍历数据加载器
        # 提前设置pred_true_cnt记录预测正确的数量
        pred_true_cnt = 0
        for batch_x, batch_y in data_loader:
            # todo 4.前向传播
            logits_y = model(batch_x)
            # todo 5.获取预测结果
            pred_y = logits_y.argmax(dim=-1)
            # todo 统计预测正确的数量
            pred_true_cnt += (pred_y == batch_y).sum()
        # todo 6.计算准确率   预测对的/总预测数
        acc = pred_true_cnt / len(data_loader.dataset)
        print(f"准确率:{acc:.4f}")


if __name__ == '__main__':
    # TODO 优化1:特征标准化
    # TODO 优化2:增加网络深度
    # TODO 优化3:修改优化器
    # TODO 优化4:调整初始学习率
    # TODO 优化5:调整批次大小
    # TODO 优化6:调整轮次大小
    # 1.提前各种超参数/路径
    datapath = "data/手机价格预测.csv"
    modelpath = "model/my_model.pth"
    lr = 0.0001
    batch_size = 16
    epochs = 80
    # 2.获取数据并切割数据
    x_train, x_test, y_train, y_test = step1_get_data(datapath)
    # todo 自动把原始整体数据分为多批数据
    # 3.提前构建数据加载器
    train_data_loder, test_data_loder = step2_get_data_loader(batch_size, x_train, x_test, y_train, y_test)
    # 4.提前准备模型
    input_num = x_train.shape[1]  # 20
    output_num = len(y_train.unique())  # 4
    model = PhoneModel(input_num, output_num)
    # TODO 模型训练(传入数据加载器和模型)
    train_model(train_data_loder, model, lr, modelpath)
    # TODO 模型评估
    eval_model(test_data_loder, modelpath)