dataset.py

import json
from pathlib import Path
from typing import List, Tuple

import numpy as np
import torch
from torch_geometric.data import Data, Dataset


class VRPLabelDataset(Dataset):
    """
    Supervised Edge Co-Membership Dataset for Multi-Objective VRP.
    Loads pre-generated JSON graphs and corresponding NSGA Pareto files.
    """
    def __init__(
        self, 
        graphs_dir: str, 
        labels_dir: str, 
        max_nodes: int = 200, 
        param_filter: str = "100_0p50_0p20",
        seed_filter: str = None,
        knn_k: int = 20,
    ):
        super().__init__(root=None, transform=None, pre_transform=None)
        
        self.graphs_dir = Path(graphs_dir)
        self.labels_dir = Path(labels_dir)
        self.max_nodes = max_nodes
        self.param_filter = param_filter
        self.seed_filter = seed_filter
        self.knn_k = knn_k

        # Collect and filter matching dataset items
        self.data_items = self._collect_data()
        
    def _collect_data(self) -> List[dict]:
        items = []
        for instance_dir in self.labels_dir.iterdir():
            if not instance_dir.is_dir():
                continue
                
            instance_name = instance_dir.name
            
            # Check node count from instance name (e.g. X-n101-k25 -> 101)
            try:
                n_str = instance_name.split("-")[1].replace("n", "")
                n_count = int(n_str)
                if n_count > self.max_nodes:
                    continue
            except (IndexError, ValueError):
                continue
                
            # Find the corresponding graph JSON
            graph_path = self.graphs_dir / f"{instance_name}.json"
            if not graph_path.exists():
                print(f"Warning: Missing graph JSON for {instance_name}. Skipping.")
                continue

            # Look for labeled Pareto sets that match the parameter filter
            search_pattern = f"*{self.param_filter}_{self.seed_filter}.sol" if self.seed_filter else f"*{self.param_filter}*.sol"
            for label_file in instance_dir.glob(search_pattern):
                # We load the json just to check pareto solutions inside
                with open(label_file, "r", encoding="utf-8") as f:
                    try:
                        nsga_data = json.load(f)
                    except json.JSONDecodeError:
                        continue
                        
                solutions = nsga_data.get("solutions", [])
                if not solutions:
                    continue
                    
                # Compute min/max for Instance-Local Normalization
                f0_list = [s["f0"] for s in solutions]
                f1_list = [s["f1"] for s in solutions]
                f0_min, f0_max = min(f0_list), max(f0_list)
                f1_min, f1_max = min(f1_list), max(f1_list)
                
                # Each pareto solution becomes a training sample!
                for sol_idx, sol in enumerate(solutions):
                    items.append({
                        "instance": instance_name,
                        "graph_path": graph_path,
                        "routes": sol["routes_customers_0based"],
                        "f0": sol["f0"],
                        "f1": sol["f1"],
                        "f0_min": f0_min,
                        "f0_max": f0_max,
                        "f1_min": f1_min,
                        "f1_max": f1_max,
                    })
        return items

    def len(self) -> int:
        return len(self.data_items)

    def get(self, idx: int) -> Data:
        item = self.data_items[idx]
        
        # 1. Load Graph 
        with open(item["graph_path"], "r", encoding="utf-8") as f:
            graph_data = json.load(f)

        # Parse Node Features -> [x, y, demand, capacity]
        capacity = float(graph_data["capacity"])
        num_nodes = int(graph_data["num_nodes"])
        
        # We append [f0_norm, f1_norm] as global condition to every node
        f0_norm = (item["f0"] - item["f0_min"]) / (item["f0_max"] - item["f0_min"] + 1e-8)
        f1_norm = (item["f1"] - item["f1_min"]) / (item["f1_max"] - item["f1_min"] + 1e-8)

        node_features = []
        coords = []
        
        # 1. Coordinate Min-Max for Graph Normalization
        x_raw = [float(graph_data["coordinates"][i][0]) for i in range(num_nodes)]
        y_raw = [float(graph_data["coordinates"][i][1]) for i in range(num_nodes)]
        x_min, x_max = min(x_raw), max(x_raw)
        y_min, y_max = min(y_raw), max(y_raw)
        
        for i in range(num_nodes):
            cx_raw, cy_raw = float(graph_data["coordinates"][i][0]), float(graph_data["coordinates"][i][1])
            cx_norm = (cx_raw - x_min) / (x_max - x_min + 1e-8)
            cy_norm = (cy_raw - y_min) / (y_max - y_min + 1e-8)
            
            dem = float(graph_data["demands"][str(i)])
            dem_norm = dem / capacity # Convert to unified scaling metric
            
            # Feature vector now unified: [x_norm, y_norm, demand_norm, f0_norm, f1_norm]
            feat = [cx_norm, cy_norm, dem_norm, f0_norm, f1_norm]
            node_features.append(feat)
            coords.append([cx_norm, cy_norm])

        x = torch.tensor(node_features, dtype=torch.float)
        pos = torch.tensor(coords, dtype=torch.float)

        # 2. Build KNN Graph (to limit edges and GPU RAM) using torch.cdist
        # Avoids torch-cluster installation on Windows
        num_neighbors = min(self.knn_k + 1, num_nodes) # +1 for self
        dists = torch.cdist(pos, pos)
        _, topk_idx = torch.topk(dists, k=num_neighbors, dim=1, largest=False)
        
        # Remove self-loops (the 0-th index of topk is always the node itself)
        neighbors = topk_idx[:, 1:]
        
        src_nodes = torch.arange(num_nodes).repeat_interleave(num_neighbors - 1)
        dst_nodes = neighbors.flatten()
        edge_index = torch.stack([src_nodes, dst_nodes], dim=0)

        # Build continuous edge features for the KNN edges
        # Edge feature map dictionary for O(1) lookup
        f0_edge_map = {}
        f1_edge_map = {}
        for e in graph_data["edges"]:
            f0_edge_map[(e["source"], e["target"])] = float(e["distance"])
            f1_edge_map[(e["source"], e["target"])] = float(e["random_cost"])
            
        f0_vals = list(f0_edge_map.values())
        f1_vals = list(f1_edge_map.values())
        dist_min, dist_max = min(f0_vals), max(f0_vals)
        rc_min, rc_max = min(f1_vals), max(f1_vals)

        edge_attr = []
        for i in range(edge_index.shape[1]):
            src = int(edge_index[0, i])
            dst = int(edge_index[1, i])
            # if directed json, might need sorting or bidirectional fallback
            c0 = f0_edge_map.get((src, dst), f0_edge_map.get((dst, src), 0.0))
            c1 = f1_edge_map.get((src, dst), f1_edge_map.get((dst, src), 0.0))
            
            # Normalize to proportional magnitude inside graph
            c0_norm = (c0 - dist_min) / (dist_max - dist_min + 1e-8)
            c1_norm = (c1 - rc_min) / (rc_max - rc_min + 1e-8)
            edge_attr.append([c0_norm, c1_norm])
        
        edge_attr = torch.tensor(edge_attr, dtype=torch.float)

        # 3. Build Ground Truth Co-membership Matrix Y
        # y_ij = 1 if node i and node j belong to the same route, 0 otherwise
        # Edge index targets have shape [num_edges]
        target_labels = torch.zeros(edge_index.shape[1], dtype=torch.float)
        
        # Build node -> vehicle_id mapping
        node_to_route = {0: -1} # Depot is -1, it exists in all routes, we generally don't cluster it.
        for route_idx, route in enumerate(item["routes"]):
            for node_id in route:
                node_to_route[node_id] = route_idx
                
        # Assign edge labels
        for i in range(edge_index.shape[1]):
            src = int(edge_index[0, i])
            dst = int(edge_index[1, i])
            # If they are the same actual class (and not the depot)
            if src != 0 and dst != 0 and node_to_route.get(src) == node_to_route.get(dst):
                target_labels[i] = 1.0

        # Construct PyG Data object
        data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr, y=target_labels)
        
        # Optional metadata attached to the returned object
        data.instance_name = item["instance"]
        data.f0 = item["f0"]
        data.f1 = item["f1"]
        return data


# --- Simple Test Code ---
if __name__ == "__main__":
    DATA_DIR = Path(__file__).resolve().parent.parent / "data"
    
    print("\nInitializing PyTorch Geometric Dataset...")
    dataset = VRPLabelDataset(
        graphs_dir=str(DATA_DIR / "xset_graphs"),
        labels_dir=str(DATA_DIR / "nsga_codex_labels"),
        max_nodes=200,                # Filter to instances N <= 200
        param_filter="100_0p50_0p20", # Filter to the best param setup
        seed_filter="2",              # Filter strictly to seed 2
        knn_k=20                      # Dense enough without exploding logic
    )
    
    print(f"Dataset successfully created!")
    print(f"Total training samples: {len(dataset)}")
    
    if len(dataset) > 0:
        sample = dataset[0]
        print("\n--- Sample 0 Specifications ---")
        print(f"Instance: {sample.instance_name}")
        print(f"Node Features (N, F) : {sample.x.shape}")
        print(f"Edge Index (2, E)    : {sample.edge_index.shape}")
        print(f"Edge Features (E, 2) : {sample.edge_attr.shape}")
        print(f"Labels (E)           : {sample.y.shape}")
        
        # Statistics
        nodes_with_demand = (sample.x[:, 2] > 0).sum().item()
        total_capacity = sample.x[0, 3].item()
        
        # Target Normalization Check
        # Nodes 0 to N will all have identically appended f0_norm and f1_norm
        f0_norm = sample.x[0, -2].item()
        f1_norm = sample.x[0, -1].item()
        
        pos_edges = int(sample.y.sum().item())
        total_edges = sample.y.shape[0]
        
        print("\n--- Internal Data Check ---")
        print(f"Nodes w/ Demand: {nodes_with_demand}")
        print(f"Total Capacity: {total_capacity}")
        print(f"Target Given: f0_norm={f0_norm:.2f}, f1_norm={f1_norm:.2f}")
        print(f"Positive Edges (same cluster): {pos_edges} / {total_edges} ({(pos_edges/total_edges)*100:.1f}%)")