#!/usr/bin/env python
# coding=UTF-8
import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 或 ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
import datetime
import tkinter as tk
from tkinter import ttk
from matplotlib.figure import Figure
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg

# 设置随机种子保证可复现性
torch.manual_seed(42)

# 设备配置（自动选择GPU或CPU）
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 数据加载和预处理
CSV_FILE_PATH = './four_classification.csv'
df = pd.read_csv(CSV_FILE_PATH)

# 标签编码
df[' Label'] = pd.Categorical(df[' Label'])
df[' Label'] = df[' Label'].cat.codes

# 类型转换
for col in df.columns:
    if df[col].dtype != 'float64':
        df[col] = df[col].astype(float)

# 特征选择
features_considered = [' Label', ' Bwd Packet Length Min', ' Subflow Fwd Bytes',
                       'Total Length of Fwd Packets', ' Fwd Packet Length Mean',
                       ' Bwd Packet Length Std', ' Flow Duration', ' Flow IAT Std',
                       'Init_Win_bytes_forward', ' Bwd Packets/s',
                       ' PSH Flag Count', ' Average Packet Size']
feature = df[features_considered]

# 数据集划分
train, test = train_test_split(feature, test_size=0.2)
train, val = train_test_split(train, test_size=0.2)


# 标准化处理
def normalize_dataset(dataset, dataset_mean, dataset_std):
    return (dataset - dataset_mean) / dataset_std


train_features = train.drop(' Label', axis=1).values
train_mean = train_features.mean(axis=0)
train_std = train_features.std(axis=0)

train_features = normalize_dataset(train_features, train_mean, train_std)
val_features = normalize_dataset(val.drop(' Label', axis=1).values, train_mean, train_std)
test_features = normalize_dataset(test.drop(' Label', axis=1).values, train_mean, train_std)


# PyTorch Dataset
class TrafficDataset(Dataset):
    def __init__(self, features, labels):
        self.features = torch.FloatTensor(features)
        self.labels = torch.LongTensor(labels.values)

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        return self.features[idx], self.labels[idx]


train_dataset = TrafficDataset(train_features, train[' Label'])
val_dataset = TrafficDataset(val_features, val[' Label'])
test_dataset = TrafficDataset(test_features, test[' Label'])

# 数据加载器
batch_size = 50
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)


# 定义模型
class DNNModel(nn.Module):
    def __init__(self, input_size, num_classes):
        super(DNNModel, self).__init__()
        self.fc = nn.Sequential(
            nn.Linear(input_size, 20),
            nn.SELU(),
            nn.Linear(20, 20),
            nn.SELU(),
            nn.Linear(20, 20),
            nn.SELU(),
            nn.Linear(20, 20),
            nn.SELU(),
            nn.Linear(20, num_classes)
        )

    def forward(self, x):
        return self.fc(x)


model = DNNModel(input_size=11, num_classes=4).to(device)


# 训练函数
def train_model(epochs):
    start_time = datetime.datetime.now()
    train_losses = []
    val_losses = []
    val_accuracies = []

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters())

    for epoch in range(epochs):
        model.train()
        train_loss = 0.0

        for inputs, labels in train_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            train_loss += loss.item()

        # 验证集评估
        model.eval()
        val_loss = 0.0
        correct = 0
        total = 0

        with torch.no_grad():
            for inputs, labels in val_loader:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = model(inputs)
                loss = criterion(outputs, labels)
                val_loss += loss.item()

                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        # 记录指标
        epoch_train_loss = train_loss / len(train_loader)
        epoch_val_loss = val_loss / len(val_loader)
        epoch_val_acc = 100 * correct / total

        train_losses.append(epoch_train_loss)
        val_losses.append(epoch_val_loss)
        val_accuracies.append(epoch_val_acc)

        print(f'Epoch {epoch + 1}/{epochs} | '
              f'Train Loss: {epoch_train_loss:.4f} | '
              f'Val Loss: {epoch_val_loss:.4f} | '
              f'Val Acc: {epoch_val_acc:.2f}%')

    end_time = datetime.datetime.now()
    print(f"Training Time: {(end_time - start_time).seconds} seconds")
    return train_losses, val_losses, val_accuracies


# GUI创建函数
def create_gui(model, train_mean, train_std, train_loss, val_loss, val_acc):
    root = tk.Tk()
    root.title("DNN网络流量分类器")
    root.geometry("1000x800")

    # 设备信息
    device_frame = ttk.Frame(root)
    device_frame.pack(pady=10)
    ttk.Label(device_frame, text="运行设备:", font=('Arial', 12)).pack(side=tk.LEFT)
    ttk.Label(device_frame, text=str(device), font=('Arial', 12, 'bold')).pack(side=tk.LEFT)

    # 训练曲线
    fig = Figure(figsize=(10, 6), dpi=100)
    ax1 = fig.add_subplot(211)
    ax1.plot(train_loss, label='训练损失', marker='o')
    ax1.plot(val_loss, label='验证损失', marker='s')
    ax1.set_title('训练过程可视化')
    ax1.set_ylabel('损失值')
    ax1.legend()

    ax2 = fig.add_subplot(212)
    ax2.plot(val_acc, label='验证准确率', color='green', marker='^')
    ax2.set_xlabel('训练轮次')
    ax2.set_ylabel('准确率(%)')
    ax2.legend()

    canvas = FigureCanvasTkAgg(fig, master=root)
    canvas.draw()
    canvas.get_tk_widget().pack(pady=10)

    # 输入表单
    input_frame = ttk.LabelFrame(root, text="流量特征输入")
    input_frame.pack(padx=20, pady=10, fill=tk.BOTH)

    entries = []
    features = [
        'Bwd Packet Length Min', 'Subflow Fwd Bytes',
        'Total Length of Fwd Packets', 'Fwd Packet Length Mean',
        'Bwd Packet Length Std', 'Flow Duration', 'Flow IAT Std',
        'Init_Win_bytes_forward', 'Bwd Packets/s',
        'PSH Flag Count', 'Average Packet Size'
    ]

    for i, feat in enumerate(features):
        row = i // 3
        col = (i % 3) * 2
        frame = ttk.Frame(input_frame)
        frame.grid(row=row, column=col, padx=5, pady=5)
        ttk.Label(frame, text=feat + ":").pack(side=tk.LEFT)
        entry = ttk.Entry(frame, width=15)
        entry.pack(side=tk.LEFT)
        entries.append(entry)

    # 预测功能
    def predict():
        input_data = []
        for entry in entries:
            try:
                input_data.append(float(entry.get()))
            except:
                input_data.append(0.0)

        # 预处理
        input_array = np.array([input_data])
        normalized_input = (input_array - train_mean) / train_std
        tensor_input = torch.FloatTensor(normalized_input).to(device)

        # 预测
        with torch.no_grad():
            output = model(tensor_input)
            _, predicted = torch.max(output.data, 1)

        class_names = ['正常流量', 'DDoS攻击', '端口扫描', '恶意软件']
        result_label.config(text=f"预测结果: {class_names[predicted.item()]}",
                            font=('Arial', 12, 'bold'), foreground='blue')

    # 控制面板
    control_frame = ttk.Frame(root)
    control_frame.pack(pady=10)

    predict_btn = ttk.Button(control_frame, text="执行预测", command=predict)
    predict_btn.pack(side=tk.LEFT, padx=10)

    result_label = ttk.Label(control_frame, text="等待输入...", font=('Arial', 12))
    result_label.pack(side=tk.LEFT)

    root.mainloop()


if __name__ == "__main__":
    # 训练模型
    train_loss, val_loss, val_acc = train_model(epochs=20)

    # 测试集评估
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for inputs, labels in test_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    print(f'\n最终测试准确率: {100 * correct / total:.2f}%')

    # 启动GUI
    create_gui(model, train_mean, train_std, train_loss, val_loss, val_acc)