import os
import torch
import numpy as np
import csv
from torch.utils.data import DataLoader, Dataset
import torch.nn as nn
from torch import optim
import time
import matplotlib.pyplot as plt

# -------------------- 自定义数据集类 --------------------
class CovidDataset(Dataset):
    """COVID-19数据集类，负责读取CSV、数据标准化、划分训练/验证/测试集"""

    def __init__(self, file_path, mode="train"):
        """
        参数:
            file_path: CSV文件路径
            mode: 'train'/'val'/'test'，决定数据划分方式
        """
        # 读取CSV文件
        with open(file_path, "r") as f:
            ori_data = list(csv.reader(f))          # 读取所有行，每行是一个列表
            csv_data = np.array(ori_data[1:])[:, 1:].astype(float)  # 跳过标题行和第一列（可能是索引），转为浮点数numpy数组

        # 根据模式划分数据
        if mode == "train":
            # 训练集：每5个样本中取前4个（索引对5取余不为0）
            indices = [i for i in range(len(csv_data)) if i % 5 != 0]
            data = torch.tensor(csv_data[indices, :-1])   # 取所有特征列（最后一列是标签）
            self.y = torch.tensor(csv_data[indices, -1])  # 最后一列作为标签

        elif mode == "val":
            # 验证集：每5个样本中取第5个（索引对5取余为0）
            indices = [i for i in range(len(csv_data)) if i % 5 == 0]
            data = torch.tensor(csv_data[indices, :-1])
            self.y = torch.tensor(csv_data[indices, -1])

        else:   # mode == "test"
            # 测试集：使用所有数据（没有标签）
            indices = [i for i in range(len(csv_data))]
            data = torch.tensor(csv_data[indices])   # 测试集只有特征，没有标签

        # 标准化：对特征进行z-score归一化（均值0，标准差1）
        mean = data.mean(dim=0, keepdim=True)   # 按列求均值，keepdim保持维度以便广播
        std = data.std(dim=0, keepdim=True)     # 按列求标准差
        std[std == 0] = 1                        # 避免除以0，若某列标准差为0则设为1（该列所有值相同）
        self.data = (data - mean) / std           # 应用标准化
        self.mode = mode

        # 注意：对于测试集，这里使用测试集自身的均值和标准差进行了标准化，这是不严谨的。
        # 正确的做法是使用训练集的均值和标准差来标准化测试集，以保证数据分布一致。
        # 在实际应用中，应该保存训练集的mean和std，并在测试时使用它们。

    def __getitem__(self, idx):
        """根据索引返回一个样本"""
        if self.mode != "test":
            # 训练/验证集返回 (特征, 标签)
            return self.data[idx].float(), self.y[idx].float()
        else:
            # 测试集只返回特征
            return self.data[idx].float()

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


# -------------------- 模型定义 --------------------
class MyModel(nn.Module):
    """简单的两层全连接神经网络，用于回归预测"""

    def __init__(self, inDim):
        """
        参数:
            inDim: 输入特征维度
        """
        super(MyModel, self).__init__()
        self.fc1 = nn.Linear(inDim, 64)    # 第一层：线性变换，输入维度inDim，输出维度64
        self.relu1 = nn.ReLU()              # ReLU激活函数
        self.fc2 = nn.Linear(64, 1)         # 第二层：线性变换，输入64，输出1（回归值）

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

        # 如果输出维度大于1（例如batch维度），则压缩最后一维，使其形状为(batch,)
        if len(x.size()) > 1:
            return x.squeeze(1)   # 去掉维度1，形状从(batch,1)变为(batch,)
        return x


# -------------------- 训练与验证函数 --------------------
def train_val(model, train_loader, val_loader, device, epochs, optimizer, loss_fn, save_path):
    """
    执行训练和验证，每个epoch后打印损失，并保存验证损失最小的模型
    """
    model = model.to(device)   # 将模型移动到指定设备（CPU/GPU）

    plt_train_loss = []        # 记录每个epoch的训练损失，用于绘图
    plt_val_loss = []          # 记录每个epoch的验证损失
    min_val_loss = float("inf")  # 初始化最小验证损失为无穷大

    for epoch in range(epochs):
        train_loss = 0.0        # 累计训练损失
        val_loss = 0.0          # 累计验证损失
        start_time = time.time()

        # ---------- 训练阶段 ----------
        model.train()           # 设置为训练模式（启用dropout、batch norm等）
        for batch_x, batch_y in train_loader:
            # 将数据移动到设备
            x, target = batch_x.to(device), batch_y.to(device)

            # 前向传播、计算损失、反向传播、参数更新
            optimizer.zero_grad()               # 清空梯度
            pred = model(x)                      # 前向传播
            train_bat_loss = loss_fn(pred, target)  # 计算损失
            train_bat_loss.backward()            # 反向传播计算梯度
            optimizer.step()                     # 更新参数

            train_loss += train_bat_loss.item()  # 累计损失

        # 计算平均训练损失
        avg_train_loss = train_loss / len(train_loader)
        plt_train_loss.append(avg_train_loss)

        # ---------- 验证阶段 ----------
        model.eval()            # 设置为评估模式（关闭dropout、batch norm等）
        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_bat_loss = loss_fn(pred, target)
                val_loss += val_bat_loss.item()

        avg_val_loss = val_loss / len(val_loader)
        plt_val_loss.append(avg_val_loss)

        # 如果当前验证损失低于历史最低，则保存模型
        if avg_val_loss < min_val_loss:
            torch.save(model.state_dict(), save_path)  # 只保存模型参数
            min_val_loss = avg_val_loss

        # 打印当前epoch信息
        print("[%03d/%03d] %2.2f sec(s) TrainLoss: %.6f | ValLoss: %.6f" %
              (epoch + 1, epochs, time.time() - start_time, avg_train_loss, avg_val_loss))

    # 绘制训练和验证损失曲线
    plt.plot(plt_train_loss)
    plt.plot(plt_val_loss)
    plt.title("Loss")
    plt.legend(["train", "val"])
    plt.show()


# -------------------- 测试与输出函数 --------------------
def evaluate(save_path, test_loader, device, rel_path):
    """
    加载训练好的模型，对测试集进行预测，并将结果写入CSV文件
    """
    # 创建模型并加载保存的参数
    model = MyModel(inDim=93).to(device)   # 注意输入维度必须与训练时一致（93）
    model.load_state_dict(torch.load(save_path, map_location=device, weights_only=True))
    model.eval()

    rel = []   # 保存所有预测结果
    with torch.no_grad():
        for x in test_loader:
            pred = model(x.to(device))
            rel.extend(pred.cpu().tolist())   # 将预测值转为列表并添加到rel

    # 将预测结果写入CSV文件
    with open(rel_path, "w", newline='') as f:
        csvWriter = csv.writer(f)
        csvWriter.writerow(["id", "tested_positive"])   # 写表头
        for i, y in enumerate(rel):
            csvWriter.writerow([i, y])                  # 每一行：样本id, 预测值

    print("文件已经保存到 " + rel_path)


# -------------------- 主程序入口 --------------------
if __name__ == "__main__":   # 确保脚本被直接运行时执行以下代码

    # 文件路径
    train_file = "covid.train.csv"
    test_file = "covid.test.csv"

    # 创建数据集实例
    train_dataset = CovidDataset(train_file, "train")
    val_dataset = CovidDataset(train_file, "val")
    test_dataset = CovidDataset(test_file, "test")

    # 超参数设置
    batch_size = 16
    # 创建数据加载器（DataLoader）用于批量加载数据
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)   # 训练集打乱顺序
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)      # 验证集不打乱
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)    # 测试集不打乱

    # 检测可用设备（GPU优先）
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(device)

    # 配置参数
    config = {
        "lr": 0.001,               # 学习率
        "epochs": 20,               # 训练轮数
        "momentum": 0.9,            # SGD动量
        "save_path": "model_save/best_model.pth",   # 最佳模型保存路径
        "rel_path": "pred.csv"       # 预测结果输出路径
    }

    # 创建保存模型的目录（如果不存在）
    os.makedirs(os.path.dirname(config["save_path"]), exist_ok=True)

    # 初始化模型、损失函数、优化器
    model = MyModel(inDim=93).to(device)          # 输入维度93
    loss_fn = nn.MSELoss()                         # 均方误差损失（回归任务）
    optimizer = optim.SGD(
        model.parameters(),
        lr=config["lr"],
        momentum=config["momentum"],
        weight_decay=0.00075                        # L2正则化系数
    )

    # 开始训练和验证
    train_val(model, train_loader, val_loader, device, config["epochs"], optimizer, loss_fn, config["save_path"])

    # 在测试集上评估并保存结果
    evaluate(config["save_path"], test_loader, device, config["rel_path"])