src/dst_gnn/data.py

"""Data loading utilities for temporal COSTE/SSS graphs."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

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

GROUP_TO_STAGE = {"HD": "T1", "LA": "T2", "MA": "T3"}
STAGE_ORDER = ("T1", "T2", "T3")
STAGE_LABELS = {"T1": "healthy", "T2": "less_affected", "T3": "more_affected"}
REQUIRED_COLUMNS = ("row", "column", "value", "sample", "group")


@dataclass
class TemporalGraphBundle:
    """Container holding stage graphs and related metadata."""

    node_names: list[str]
    sample_stage_map: dict[str, str]
    sample_matrices: dict[str, np.ndarray]
    stage_matrices: dict[str, np.ndarray]
    stage_graphs: list[Data]
    sample_counts: dict[str, int]


def load_relationship_table(csv_path: str | Path) -> pd.DataFrame:
    """Load and validate the flattened COSTE/SSS relationship table."""

    csv_path = Path(csv_path)
    df = pd.read_csv(csv_path)
    missing_columns = [column for column in REQUIRED_COLUMNS if column not in df.columns]
    if missing_columns:
        raise ValueError(f"Missing required columns: {missing_columns}")

    unknown_groups = sorted(set(df["group"]) - set(GROUP_TO_STAGE))
    if unknown_groups:
        raise ValueError(f"Unexpected group labels: {unknown_groups}")

    return df.copy()


def build_sample_matrices(
    df: pd.DataFrame,
    node_names: list[str] | None = None,
) -> tuple[list[str], dict[str, str], dict[str, np.ndarray]]:
    """Build full per-sample directed SSS matrices.

    Missing off-diagonal pairs are filled with 1.0, matching the manuscript's
    treatment of absent spatial associations. Diagonal values are set to 0.0.
    """

    if node_names is None:
        node_names = sorted(set(df["row"]).union(set(df["column"])))

    node_to_idx = {name: idx for idx, name in enumerate(node_names)}
    num_nodes = len(node_names)

    sample_stage_map: dict[str, str] = {}
    sample_matrices: dict[str, np.ndarray] = {}

    for sample, sample_df in df.groupby("sample", sort=True):
        stage = GROUP_TO_STAGE[str(sample_df["group"].iloc[0])]
        matrix = np.ones((num_nodes, num_nodes), dtype=np.float32)
        np.fill_diagonal(matrix, 0.0)

        for row in sample_df.itertuples(index=False):
            source = node_to_idx[str(row.row)]
            target = node_to_idx[str(row.column)]
            matrix[source, target] = float(row.value)

        sample_stage_map[str(sample)] = stage
        sample_matrices[str(sample)] = matrix

    return node_names, sample_stage_map, sample_matrices


def aggregate_stage_matrices(
    sample_stage_map: dict[str, str],
    sample_matrices: dict[str, np.ndarray],
) -> tuple[dict[str, np.ndarray], dict[str, int]]:
    """Average per-sample matrices within each temporal stage."""

    grouped: dict[str, list[np.ndarray]] = {stage: [] for stage in STAGE_ORDER}
    for sample, matrix in sample_matrices.items():
        grouped[sample_stage_map[sample]].append(matrix)

    stage_matrices: dict[str, np.ndarray] = {}
    sample_counts: dict[str, int] = {}
    for stage in STAGE_ORDER:
        matrices = grouped[stage]
        if not matrices:
            raise ValueError(f"No sample matrices found for stage {stage}")
        stage_matrices[stage] = np.mean(np.stack(matrices, axis=0), axis=0).astype(np.float32)
        sample_counts[stage] = len(matrices)

    return stage_matrices, sample_counts


def matrix_to_graph(stage: str, matrix: np.ndarray, node_names: list[str]) -> Data:
    """Convert a directed SSS matrix to a PyG Data object.

    Message passing uses an affinity transform of ``1 - SSS`` so that closer
    spatial associations produce larger edge weights.
    """

    num_nodes = matrix.shape[0]
    sss_matrix = np.clip(matrix, 0.0, 1.0).astype(np.float32)
    affinity = 1.0 - sss_matrix
    np.fill_diagonal(affinity, 0.0)

    source_nodes, target_nodes = np.where(affinity > 0)
    edge_index = torch.tensor(
        np.vstack([source_nodes, target_nodes]),
        dtype=torch.long,
    )
    edge_weight = torch.tensor(affinity[source_nodes, target_nodes], dtype=torch.float32)
    x = torch.eye(num_nodes, dtype=torch.float32)

    graph = Data(
        x=x,
        edge_index=edge_index,
        edge_weight=edge_weight,
        y=torch.from_numpy(sss_matrix),
    )
    graph.stage_name = stage
    graph.stage_index = torch.tensor([STAGE_ORDER.index(stage)], dtype=torch.long)
    graph.node_names = node_names
    graph.sss_matrix = torch.from_numpy(sss_matrix)
    graph.affinity_matrix = torch.from_numpy(affinity.astype(np.float32))
    return graph


def build_temporal_graphs(csv_path: str | Path) -> TemporalGraphBundle:
    """Build stage graphs from the flattened COSTE/SSS CSV."""

    df = load_relationship_table(csv_path)
    node_names, sample_stage_map, sample_matrices = build_sample_matrices(df)
    stage_matrices, sample_counts = aggregate_stage_matrices(sample_stage_map, sample_matrices)
    stage_graphs = [matrix_to_graph(stage, stage_matrices[stage], node_names) for stage in STAGE_ORDER]

    return TemporalGraphBundle(
        node_names=node_names,
        sample_stage_map=sample_stage_map,
        sample_matrices=sample_matrices,
        stage_matrices=stage_matrices,
        stage_graphs=stage_graphs,
        sample_counts=sample_counts,
    )