import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

# ========== 全局配置 中文+随机种子 ==========
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
torch.manual_seed(42)
np.random.seed(42)

# ========== 1.读取并预处理数据 ==========
df = pd.read_csv("pima-diabetes.csv")
X = df.drop("Outcome", axis=1).values
y = df["Outcome"].values.reshape(-1, 1)

# 划分训练集、测试集，分层抽样保证正负样本比例
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)

# 转为torch张量
X_train = torch.tensor(X_train, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.float32)
X_test = torch.tensor(X_test, dtype=torch.float32)
y_test = torch.tensor(y_test, dtype=torch.float32)

# ========== 2.定义网络 ==========
class DiabetesNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(8, 64),
            nn.ReLU(),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, 1),
            nn.Sigmoid()
        )
    def forward(self, x):
        return self.net(x)

loss_fn = nn.BCELoss()
lr = 0.001
max_epoch = 150

# ===================== 第一部分：无早停，完整训练（观察过拟合） =====================
print("===== 开始无早停训练 =====")
model1 = DiabetesNet()
opt1 = optim.Adam(model1.parameters(), lr=lr)
train_loss_no_early = []
test_loss_no_early = []

for epoch in range(max_epoch):
    model1.train()
    opt1.zero_grad()
    pred_train = model1(X_train)
    tr_loss = loss_fn(pred_train, y_train)
    tr_loss.backward()
    opt1.step()

    model1.eval()
    with torch.no_grad():
        pred_test = model1(X_test)
        te_loss = loss_fn(pred_test, y_test)

    train_loss_no_early.append(tr_loss.item())
    test_loss_no_early.append(te_loss.item())
    if (epoch+1) % 20 == 0:
        print(f"Epoch {epoch+1:3d} | Train:{tr_loss:.4f} | Test:{te_loss:.4f}")

# 绘图：无早停loss曲线
plt.figure(figsize=(8, 5))
plt.plot(train_loss_no_early, label="Train Loss")
plt.plot(test_loss_no_early, label="Test Loss")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.title("Training and Testing Loss Curves（无早停，出现过拟合）")
plt.legend()
plt.grid(alpha=0.3)
plt.tight_layout()
plt.savefig("loss_no_earlystop.jpg", dpi=300)
plt.show()

# ===================== 第二部分：带早停法训练 =====================
print("\n===== 开始带早停训练 =====")
model2 = DiabetesNet()
opt2 = optim.Adam(model2.parameters(), lr=lr)
train_loss_early = []
test_loss_early = []
patience = 15          # 连续15轮测试loss不下降就停止
best_test_loss = float("inf")
wait = 0
stop_epoch = max_epoch # 记录停止轮次

for epoch in range(max_epoch):
    model2.train()
    opt2.zero_grad()
    pred_train = model2(X_train)
    tr_loss = loss_fn(pred_train, y_train)
    tr_loss.backward()
    opt2.step()

    model2.eval()
    with torch.no_grad():
        pred_test = model2(X_test)
        te_loss = loss_fn(pred_test, y_test)

    train_loss_early.append(tr_loss.item())
    test_loss_early.append(te_loss.item())

    # 早停判断逻辑
    if te_loss < best_test_loss:
        best_test_loss = te_loss
        wait = 0
    else:
        wait += 1
        if wait >= patience:
            stop_epoch = epoch + 1
            print(f"✅ 触发早停，Epoch={stop_epoch}，最佳测试损失：{best_test_loss:.4f}")
            break

    if (epoch+1) % 20 == 0:
        print(f"Epoch {epoch+1:3d} | Train:{tr_loss:.4f} | Test:{te_loss:.4f}")

# 绘图：带早停loss曲线，标记停止点
plt.figure(figsize=(8, 5))
plt.plot(train_loss_early, label="Train Loss")
plt.plot(test_loss_early, label="Test Loss")
plt.axvline(x=stop_epoch-1, color="red", linestyle="--", label=f"早停点 Epoch={stop_epoch}")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.title("Training and Testing Loss Curves（加入早停法 Early Stopping）")
plt.legend()
plt.grid(alpha=0.3)
plt.tight_layout()
plt.savefig("loss_with_earlystop.jpg", dpi=300)
plt.show()

print("\n✅ 全部训练完成，两张图片已保存：")
print("loss_no_earlystop.jpg    无早停（过拟合对照）")
print("loss_with_earlystop.jpg  带早停法")