import torch
import pandas as pd
import numpy as np
from torch_geometric.data import Data
import config
import random

def build_graphs_from_dataframe(df, feature_cols, label_col):
    print("Building graphs... FINAL PAPER-LEVEL VERSION")

    # 1. 时间窗口改小一点，保证图的数量
    df["time_window"] = (df["time"] // config.TIME_WINDOW).astype(int)
    window_groups = df.groupby("time_window")
    pyg_graph_list = []

    for window_id, window_df in window_groups:
        n = len(window_df)
        # 【修复1】不限制上限，只限制下限
        if n < 3:
            continue

        x = torch.tensor(window_df[feature_cols].values, dtype=torch.float)
        y = torch.tensor(window_df[label_col].values, dtype=torch.long)

        # 【修复2】快速构图：每个节点连前后5个
        edge_index = []
        for i in range(n):
            for j in range(i+1, min(i+6, n)):
                edge_index.append([i, j])
                edge_index.append([j, i])
        edge_index = torch.tensor(edge_index, dtype=torch.long).t().contiguous() if edge_index else torch.empty((2, 0), dtype=torch.long)

        # 【修复3】掩码先全设为True，后面全局划分
        train_mask = torch.ones(n, dtype=torch.bool)
        val_mask = torch.zeros(n, dtype=torch.bool)
        test_mask = torch.zeros(n, dtype=torch.bool)

        graph = Data(x=x, edge_index=edge_index, y=y,
                     train_mask=train_mask, val_mask=val_mask, test_mask=test_mask)
        pyg_graph_list.append(graph)

        # 最多15000张图，足够了
        if len(pyg_graph_list) >= 15000:
            break

    # 【修复4】全局彻底打乱！保证训练集和测试集都有正常和攻击
    print(f"Shuffling {len(pyg_graph_list)} graphs...")
    random.shuffle(pyg_graph_list)
    random.shuffle(pyg_graph_list)

    split = int(len(pyg_graph_list) * 0.7)
    train_graphs = pyg_graph_list[:split]
    test_graphs = pyg_graph_list[split:]

    # 给测试集打开 test_mask
    for g in test_graphs:
        g.train_mask[:] = False
        g.test_mask[:] = True

    print(f"Done! Train: {len(train_graphs)}, Test: {len(test_graphs)}")
    return train_graphs, test_graphs