alnn_rewrite/model/dataset.py

import math
import pathlib as pl
import random
import re
from typing import Callable, Dict, Iterator, List, Tuple, cast

import nibabel as nib
import pandas as pd
import torch
import torch.utils.data as data
from jaxtyping import Float
from torch.utils.data import DataLoader, Subset

from util.seeding import torch_generator

# Tokens that appear in MRI filenames to indicate the diagnostic class.
# On disk the cognitively-normal class is spelled "NL" (not "CN").
CLASS_TOKENS: Dict[str, torch.Tensor] = {
    "AD": torch.tensor([1.0, 0.0]),
    "NL": torch.tensor([0.0, 1.0]),
}


def _row_to_float_tensor(row: pd.DataFrame, *, image_id: int) -> torch.Tensor:
    values = row.drop(columns=["Image Data ID"]).iloc[0]
    numeric = pd.to_numeric(values, errors="coerce")

    if numeric.isna().any():  # pyright: ignore
        bad_columns = [column for column, value in numeric.items() if pd.isna(value)]  # pyright: ignore
        raise ValueError(
            f"Non-numeric Excel values for Image Data ID {image_id}: {bad_columns}"
        )

    return torch.tensor(numeric.to_numpy(dtype="float32"))  # pyright: ignore


def xls_pre(df: pd.DataFrame) -> pd.DataFrame:
    """
    Preprocess the Excel DataFrame.
    This function can be customized to filter or modify the DataFrame as needed.

    Returns a DataFrame with columns [Image Data ID, Sex, Age (current)] where
    Sex is encoded 0/1. "Image Data ID" is deliberately kept as a column (not an
    index) because downstream lookups filter on it directly
    (``xls_values["Image Data ID"] == img_id``).
    """

    # Work on an explicit copy so column assignment writes back reliably instead
    # of raising SettingWithCopyWarning against a view of ``df``.
    #
    # The ``cast``s are needed because pandas ships no type stubs: basedpyright
    # widens ``.copy()``/``.replace()``/``__getitem__`` to
    # ``DataFrame | Series | Unknown``, which then has no ``.str`` accessor and
    # isn't assignable to the declared return type. ``cast`` asserts the concrete
    # pandas type at each step without any runtime effect.
    data = cast(pd.DataFrame, df[["Image Data ID", "Sex", "Age (current)"]].copy())

    # ``.str`` is pandas' vectorized string accessor: it applies a str method
    # element-wise over every value in the Series. Here it strips surrounding
    # whitespace from each Sex entry (e.g. " M " -> "M") before encoding.
    sex_col = cast(pd.Series, data["Sex"])
    data["Sex"] = sex_col.str.strip()

    data = cast(pd.DataFrame, data.replace({"M": 0, "F": 1}))

    return data


class ADNIDataset(data.Dataset):  # type: ignore
    """
    A PyTorch Dataset class for loading
    and processing MRI and Excel data from the ADNI dataset.
    """

    def __init__(
        self,
        mri_data: Float[torch.Tensor, "n_samples channels width height depth"],
        xls_data: Float[torch.Tensor, "n_samples features"],
        expected_classes: Float[
            torch.Tensor, " classes"
        ],  # leading space is necessary so pyright doesn't complain
        filename_ids: List[int],
        device: str = "cuda",
    ):
        """
        Args:
            mri_data (torch.Tensor): 5D tensor of MRI data with shape (n_samples, channels, width, height, depth).
            xls_data (torch.Tensor): 2D tensor of Excel data with shape (n_samples, features).
        """
        # Keep full datasets on CPU and move only active batches to accelerator.
        # This avoids loading the entire training corpus into VRAM.
        self.mri_data = mri_data.float().cpu()
        self.xls_data = xls_data.float().cpu()
        self.expected_classes = expected_classes.float().cpu()
        self.filename_ids = filename_ids

    def __len__(self) -> int:
        """
        Returns the number of samples in the dataset.
        """
        return self.mri_data.shape[0]  # 0th dimension is the number of samples

    def __getitem__(
        self, idx: int
    ) -> Tuple[
        Float[torch.Tensor, "channels width height depth"],
        Float[torch.Tensor, " features"],
        Float[torch.Tensor, " classes"],
        int,
    ]:
        """
        Returns a sample from the dataset at the given index.

        Args:
            idx (int): Index of the sample to retrieve.

        Returns:
            tuple: A tuple containing the MRI data and Excel data for the sample.
        """

        # Slices the data on the 0th dimension, corresponding to the sample index
        mri_sample = self.mri_data[idx]
        xls_sample = self.xls_data[idx]
        # Assuming expected_classes is a tensor of classes, we return it as well
        expected_classes = self.expected_classes[idx]
        filename_id = self.filename_ids[idx]

        return mri_sample, xls_sample, expected_classes, filename_id


def load_adni_data_from_file(
    mri_files: Iterator[pl.Path],  # List of nibablel files
    xls_file: pl.Path,  # Path to the Excel file
    device: str = "cuda",
    xls_preprocessor: Callable[[pd.DataFrame], pd.DataFrame] = lambda x: x,
) -> ADNIDataset:
    """
    Loads MRI and Excel data from the ADNI dataset.

    Args:
        mri_files (List[pl.Path]): List of paths to the MRI files.
        xls_file (pl.Path): Path to the Excel file.

    Returns:
        ADNIDataset: The loaded dataset.
    """
    # Load the Excel data
    xls_values = xls_preprocessor(pd.read_csv(xls_file))  # type: ignore

    # Load the MRI data
    mri_data_unstacked: List[torch.Tensor] = []
    expected_classes_unstacked: List[torch.Tensor] = []
    xls_data_unstacked: List[torch.Tensor] = []
    img_ids: List[int] = []
    for file in mri_files:
        filename = file.stem
        match re.search(r".+?(?=_I)_I(\d+).+", filename):
            case None:
                raise ValueError(
                    f"Filename {filename} does not match expected pattern."
                )
            case m:
                img_id = int(m.group(1))

        file_mri_data = torch.from_numpy(nib.load(file).get_fdata())  # type: ignore # type checking does not work well with nibabel

        # Read the filename to determine the expected class. On disk the classes
        # are spelled "AD" and "NL"; see CLASS_TOKENS.
        file_expected_class: torch.Tensor | None = None
        for token, one_hot in CLASS_TOKENS.items():
            if token in filename:
                file_expected_class = one_hot.clone()
                break

        if file_expected_class is None:
            raise ValueError(
                f"Filename {filename} does not contain a valid class identifier "
                f"({' or '.join(CLASS_TOKENS)})."
            )

        mri_data_unstacked.append(file_mri_data)
        expected_classes_unstacked.append(file_expected_class)
        # Extract the corresponding row from the Excel data using the img_id
        xls_row = xls_values.loc[xls_values["Image Data ID"] == img_id]
        if xls_row.empty:
            raise ValueError(
                f"No matching row found in Excel data for Image Data ID {img_id}."
            )
        elif len(xls_row) > 1:
            raise ValueError(
                f"Multiple rows found in Excel data for Image Data ID {img_id}."
            )
        file_xls_data = _row_to_float_tensor(xls_row, image_id=img_id)

        xls_data_unstacked.append(file_xls_data)
        img_ids.append(img_id)

    mri_data = torch.stack(mri_data_unstacked).unsqueeze(1)
    # Stack the list of tensors into a single tensor and unsqueeze along the channel dimension

    xls_data = torch.stack(
        xls_data_unstacked
    )  # Stack the list of tensors into a single tensor

    expected_classes = torch.stack(
        expected_classes_unstacked
    )  # Stack the list of expected classes into a single tensor

    return ADNIDataset(mri_data, xls_data, expected_classes, img_ids, device=device)


def divide_dataset(
    dataset: ADNIDataset,
    ratios: Tuple[float, float, float],
    seed: int,
) -> List[data.Subset[ADNIDataset]]:
    """
    Divides the dataset into training, validation, and test sets.

    Args:
        dataset (ADNIDataset): The dataset to divide.
        train_ratio (float): The ratio of the training set.
        val_ratio (float): The ratio of the validation set.
        test_ratio (float): The ratio of the test set.

    Returns:
        Result[List[data.Subset[ADNIDataset]], str]: A Result object containing the subsets or an error message.
    """
    if not math.isclose(sum(ratios), 1.0, rel_tol=1e-9, abs_tol=1e-9):
        raise ValueError(f"Ratios must sum to 1.0, got {ratios} (sum={sum(ratios)}).")

    # Set the random seed for reproducibility
    gen = torch.Generator().manual_seed(seed)
    return data.random_split(dataset, ratios, generator=gen)


def initalize_dataloaders(
    datasets: List[Subset[ADNIDataset]],
    batch_size: int = 64,
    seed: int | None = None,
) -> List[DataLoader[ADNIDataset]]:
    """
    Initializes the DataLoaders for the [train, val, test] datasets.

    The three splits are configured differently:
      - train: shuffled (for SGD) with a seeded generator for reproducibility,
        and ``drop_last=True`` so a size-1 final batch cannot crash BatchNorm1d.
      - val / test: NOT shuffled, so sample order is stable. Stable order is
        required for evaluation, where saved model outputs must align with a
        fixed sample axis (see the netCDF evaluation schema).

    Args:
        datasets: Subsets in ``[train, val, test]`` order.
        batch_size: The batch size for the DataLoaders.
        seed: Seed for the training shuffle generator. ``None`` leaves shuffling
            to global RNG state.

    Returns:
        List[DataLoader[ADNIDataset]]: DataLoaders in the same order as ``datasets``.
    """
    loader_configs = [
        {"shuffle": True, "drop_last": True},  # train
        {"shuffle": False, "drop_last": False},  # val
        {"shuffle": False, "drop_last": False},  # test
    ]
    if len(datasets) != len(loader_configs):
        raise ValueError(
            f"Expected 3 datasets ([train, val, test]), got {len(datasets)}."
        )

    pin_memory = torch.cuda.is_available()
    loaders: List[DataLoader[ADNIDataset]] = []
    for dataset, cfg in zip(datasets, loader_configs):
        generator = torch_generator(seed) if cfg["shuffle"] else None
        loaders.append(
            DataLoader(
                dataset,
                batch_size=batch_size,
                shuffle=cfg["shuffle"],
                drop_last=cfg["drop_last"],
                pin_memory=pin_memory,
                generator=generator,
            )
        )
    return loaders


def _patient_index(
    dataset: ADNIDataset, ptids: List[Tuple[int, str]]
) -> Dict[str, List[int]]:
    """Map patient id -> dataset sample indices, validating PTID consistency.

    Shared by the single-split and k-fold partitioners so both enforce the same
    patient-grouping and leakage checks.
    """
    if not ptids:
        raise ValueError("ptids list cannot be empty.")

    image_to_patient: Dict[int, str] = {}
    for image_id, patient_id in ptids:
        image_id_int = int(image_id)
        patient_id_str = str(patient_id).strip()
        if not patient_id_str or patient_id_str.lower() == "nan":
            raise ValueError(f"Invalid PTID for Image Data ID {image_id_int}.")
        if (
            image_id_int in image_to_patient
            and image_to_patient[image_id_int] != patient_id_str
        ):
            raise ValueError(
                f"Conflicting PTIDs for Image Data ID {image_id_int}: "
                f"{image_to_patient[image_id_int]} vs {patient_id_str}."
            )
        image_to_patient[image_id_int] = patient_id_str

    patient_to_indices: Dict[str, List[int]] = {}
    for idx, image_id in enumerate(dataset.filename_ids):
        if image_id not in image_to_patient:
            raise ValueError(
                f"Missing PTID mapping for dataset Image Data ID {image_id}."
            )
        patient_to_indices.setdefault(image_to_patient[image_id], []).append(idx)

    return patient_to_indices


def divide_dataset_by_patient_id(
    dataset: ADNIDataset,
    ptids: List[Tuple[int, str]],
    ratios: Tuple[float, float, float],
    seed: int,
) -> List[data.Subset[ADNIDataset]]:
    """
    Divides the dataset into training, validation, and test sets based on patient IDs.
    Ensures that all samples from the same patient are in the same set.

    Args:
        dataset (ADNIDataset): The dataset to divide.
        ptids (List[Tuple[int, str]]): A list of tuples containing image file IDs and their corresponding patient IDs.
        ratios (Tuple[float, float, float]): The ratios for training, validation, and test sets.
        seed (int): The random seed for reproducibility.

    Returns:
        List[data.Subset[ADNIDataset]]: A list of subsets for training, validation, and test sets.

    Notes:
        This split is grouped by PTID, so all images from the same patient are assigned
        to exactly one partition to avoid patient-level leakage across train/val/test.
    """
    if not math.isclose(sum(ratios), 1.0, rel_tol=1e-9, abs_tol=1e-9):
        raise ValueError(f"Ratios must sum to 1.0, got {ratios} (sum={sum(ratios)}).")

    patient_to_indices = _patient_index(dataset, ptids)

    shuffled_patients = list(patient_to_indices.keys())
    random.Random(seed).shuffle(shuffled_patients)

    train_cutoff = int(len(shuffled_patients) * ratios[0])
    val_cutoff = train_cutoff + int(len(shuffled_patients) * ratios[1])

    train_patients = shuffled_patients[:train_cutoff]
    val_patients = shuffled_patients[train_cutoff:val_cutoff]
    test_patients = shuffled_patients[val_cutoff:]

    train_patient_set = set(train_patients)
    val_patient_set = set(val_patients)
    test_patient_set = set(test_patients)

    if (
        train_patient_set & val_patient_set
        or train_patient_set & test_patient_set
        or val_patient_set & test_patient_set
    ):
        raise ValueError("Patient separation violated across train/val/test splits.")

    all_patients = set(patient_to_indices.keys())
    if train_patient_set | val_patient_set | test_patient_set != all_patients:
        raise ValueError("Not all patients were assigned to a split.")

    train_indices = [
        idx for patient_id in train_patients for idx in patient_to_indices[patient_id]
    ]
    val_indices = [
        idx for patient_id in val_patients for idx in patient_to_indices[patient_id]
    ]
    test_indices = [
        idx for patient_id in test_patients for idx in patient_to_indices[patient_id]
    ]

    train_index_set = set(train_indices)
    val_index_set = set(val_indices)
    test_index_set = set(test_indices)

    if (
        train_index_set & val_index_set
        or train_index_set & test_index_set
        or val_index_set & test_index_set
    ):
        raise ValueError("Sample index overlap detected across train/val/test splits.")

    all_split_indices = train_index_set | val_index_set | test_index_set
    expected_indices = set(range(len(dataset)))
    if all_split_indices != expected_indices:
        raise ValueError(
            "Split coverage check failed: not all dataset samples are assigned exactly once."
        )

    return [
        Subset(dataset, train_indices),
        Subset(dataset, val_indices),
        Subset(dataset, test_indices),
    ]


def kfold_split_by_patient_id(
    dataset: ADNIDataset,
    ptids: List[Tuple[int, str]],
    k_folds: int,
    fold: int,
    seed: int,
    val_fraction: float = 0.15,
) -> List[data.Subset[ADNIDataset]]:
    """Patient-grouped k-fold partition for a single fold.

    Patients are shuffled once (deterministically from ``seed``) and dealt
    round-robin into ``k_folds`` folds, so every patient is in exactly one fold
    regardless of which fold is requested. For the requested ``fold``:
      - test  = that fold's patients (the held-out slice),
      - the remaining patients are split into val/train by ``val_fraction``.
    All splits are patient-disjoint (no leakage), and across all folds every
    sample is a test point exactly once.

    Args:
        dataset: The dataset to partition.
        ptids: (Image Data ID, PTID) pairs covering the dataset.
        k_folds: Number of folds (>= 2).
        fold: Which fold to hold out as test (0 <= fold < k_folds).
        seed: Seed for the (shared across folds) patient shuffle.
        val_fraction: Fraction of the NON-test patients used for validation.

    Returns:
        ``[train, val, test]`` subsets.
    """
    if k_folds < 2:
        raise ValueError(f"k_folds must be >= 2, got {k_folds}.")
    if not 0 <= fold < k_folds:
        raise ValueError(f"fold must be in [0, {k_folds}), got {fold}.")
    if not 0.0 <= val_fraction < 1.0:
        raise ValueError(f"val_fraction must be in [0.0, 1.0), got {val_fraction}.")

    patient_to_indices = _patient_index(dataset, ptids)
    patients = list(patient_to_indices.keys())
    if len(patients) < k_folds:
        raise ValueError(
            f"Cannot make {k_folds} folds from {len(patients)} patients."
        )
    # Deterministic shuffle shared across all folds, then round-robin deal so
    # fold sizes are balanced and the same partition is reproduced for every fold.
    random.Random(seed).shuffle(patients)
    fold_patients = [patients[i::k_folds] for i in range(k_folds)]

    test_patients = fold_patients[fold]
    rest = [p for f in range(k_folds) if f != fold for p in fold_patients[f]]
    n_val = int(round(len(rest) * val_fraction))
    val_patients = rest[:n_val]
    train_patients = rest[n_val:]

    def _indices(group: List[str]) -> List[int]:
        return [idx for patient in group for idx in patient_to_indices[patient]]

    train_indices = _indices(train_patients)
    val_indices = _indices(val_patients)
    test_indices = _indices(test_patients)

    # Safety: patient-disjoint and full, exact coverage of the dataset.
    train_set, val_set, test_set = (
        set(train_indices),
        set(val_indices),
        set(test_indices),
    )
    if train_set & val_set or train_set & test_set or val_set & test_set:
        raise ValueError("Sample index overlap detected across k-fold splits.")
    if train_set | val_set | test_set != set(range(len(dataset))):
        raise ValueError(
            "k-fold coverage check failed: not all samples assigned exactly once."
        )

    return [
        Subset(dataset, train_indices),
        Subset(dataset, val_indices),
        Subset(dataset, test_indices),
    ]