config.py

"""
Vbai-2.6AD — configuration.

Paired multimodal early-Alzheimer's detection: a 3D MRI volume plus a panel of
13 biomarkers, fused into one representation.

--------------------------------------------------------------------------
YOU MUST SET YOUR OWN PATHS.
--------------------------------------------------------------------------
No data location is hard-coded. Point the environment variables below at your
own files before running anything:

    VBAI_DATASET_ROOT     root of your imaging + tabular data
    VBAI_VOLUME_ROOT      root of the volume files referenced by the manifest
    VBAI_MODEL_SAVE_ROOT  where checkpoints are written

If a variable is unset, the loader walks up the project tree looking for a
`Datasets` directory. If that fails too, the data-preparation step raises an
explicit error rather than guessing.

Nothing here describes or names a particular corpus. Bring your own data; the
expected column contract is FEATURE_NAMES below.
"""
import os
from dataclasses import dataclass, field
from typing import List

PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))


def _walk_up_for_marker(start: str, marker_subdirs=("Datasets",), max_levels: int = 6):
    """Walk upwards from `start` until a directory containing all markers is found."""
    cur = os.path.abspath(start)
    for _ in range(max_levels):
        if all(os.path.isdir(os.path.join(cur, m)) for m in marker_subdirs):
            return cur
        parent = os.path.dirname(cur)
        if parent == cur:
            break
        cur = parent
    return None


def _resolve_dataset_root() -> str:
    """Resolution order: environment variable → project-tree search → default."""
    env = os.environ.get("VBAI_DATASET_ROOT")
    if env and os.path.isdir(env):
        return env

    walked = _walk_up_for_marker(PROJECT_ROOT, marker_subdirs=("Datasets",))
    if walked is not None:
        return os.path.join(walked, "Datasets")

    # Fall through to a relative default; data preparation will report clearly
    # if nothing is there. SET VBAI_DATASET_ROOT TO YOUR OWN PATH.
    return os.path.normpath(os.path.join(PROJECT_ROOT, "..", "Datasets"))


def _resolve_model_save_root() -> str:
    env = os.environ.get("VBAI_MODEL_SAVE_ROOT")
    if env:
        return env
    walked = _walk_up_for_marker(PROJECT_ROOT, marker_subdirs=("Datasets",))
    if walked is not None:
        return os.path.join(walked, "Models", "Vbai-2.6AD")
    return os.path.normpath(os.path.join(PROJECT_ROOT, "..", "Models", "Vbai-2.6AD"))


DATASET_ROOT = _resolve_dataset_root()
MODEL_SAVE_ROOT = _resolve_model_save_root()

# Root of the volume files. The visit manifest stores relative paths; this is
# what they are resolved against. SET VBAI_VOLUME_ROOT TO YOUR OWN PATH.
VOLUME_ROOT = os.environ.get("VBAI_VOLUME_ROOT") or os.path.join(DATASET_ROOT, "volumes")

# Kept for backward compatibility with scripts that expect these names.
TBM_ROOT = VOLUME_ROOT
TBM_CSV = os.environ.get("VBAI_VOLUME_MANIFEST") or os.path.join(VOLUME_ROOT, "manifest.csv")

# Which volume modality is in use. Selected by the extraction scripts through
# the --tbm / --t1 flag, which sets this variable before config is imported.
# A checkpoint trained on one modality must never be fed the other.
USE_TBM = bool(int(os.environ.get("VBAI_USE_TBM", "0")))

CACHE_DIR = os.path.join(PROJECT_ROOT, "_cache")
os.makedirs(CACHE_DIR, exist_ok=True)
PAIRED_PARQUET_T1 = os.path.join(CACHE_DIR, "paired_visits.parquet")
PAIRED_PARQUET_TBM = os.path.join(CACHE_DIR, "paired_visits_tbm.parquet")
PAIRED_PARQUET = PAIRED_PARQUET_TBM if USE_TBM else PAIRED_PARQUET_T1

# Tabular feature order — a fixed contract relied on everywhere downstream.
# Your table must provide these columns (missing values are allowed and are
# handled explicitly through a per-feature mask; see NUM_TABULAR_INPUTS).
FEATURE_NAMES: List[str] = [
    "Age",                # demographic
    "Sex",                # 0 = F, 1 = M
    "MMSE",               # cognitive
    "CDRSB",              # cognitive (CDR sum of boxes)
    "APOE4_count",        # genetic, 0/1/2 e4 alleles
    "CSF_ABETA42",        # CSF
    "CSF_TAU",            # CSF
    "CSF_PTAU",           # CSF
    "CSF_AB42_AB40",      # CSF ratio
    "PLASMA_PTAU",        # blood
    "PLASMA_NFL",         # blood
    "PLASMA_AB42_AB40",   # blood ratio
    "PLASMA_GFAP",        # blood
]
NUM_FEATURES = len(FEATURE_NAMES)            # 13
# One value plus one missing-mask bit per feature. The mask is not decoration:
# an unmeasured biomarker must stay distinguishable from a normal one.
NUM_TABULAR_INPUTS = NUM_FEATURES * 2

CLASS_NAMES = ["CN", "MCI", "AD"]
DIAGNOSIS_MAP = {"CN": 0, "MCI": 1, "Dementia": 2, "AD": 2,
                 "EMCI": 1, "LMCI": 1, "SMC": 0}


@dataclass
class ModelConfig:
    mri_input_shape: tuple = (1, 96, 96, 96)
    mri_encoder_channels: List[int] = field(default_factory=lambda: [32, 64, 128, 256])
    mri_bottleneck_channels: int = 512
    mri_feature_dim: int = 512
    mri_dropout: float = 0.4
    use_cbam: bool = True
    use_se_block: bool = True

    num_tabular_inputs: int = NUM_TABULAR_INPUTS
    tabular_hidden_dims: List[int] = field(default_factory=lambda: [128, 256])
    tabular_feature_dim: int = 256
    tabular_dropout: float = 0.3

    fusion_dim: int = 512
    fusion_num_heads: int = 8
    fusion_dropout: float = 0.3

    num_classes: int = 3
    progression_hidden_dim: int = 256
    max_progression_months: int = 120
    num_time_bins: int = 24

    # Modality dropout during training: teaches the model to survive a missing
    # arm at inference instead of collapsing.
    p_drop_mri: float = 0.15
    p_drop_tab: float = 0.15
    # Feature-wise random masking, simulating biomarkers absent at inference.
    p_feature_mask: float = 0.20


@dataclass
class TrainingConfig:
    seed: int = 42
    device: str = "cuda"
    num_workers: int = 4
    pin_memory: bool = True
    mixed_precision: bool = True

    # Phase 1 — MRI encoder pretraining
    phase1_epochs: int = 40
    phase1_batch_size: int = 4
    phase1_lr: float = 3e-4
    phase1_weight_decay: float = 1e-4

    # Phase 2 — tabular encoder pretraining
    phase2_epochs: int = 60
    phase2_batch_size: int = 64
    phase2_lr: float = 1e-3
    phase2_weight_decay: float = 1e-4

    # Phase 3 — joint fusion on paired visits
    phase3_epochs: int = 40
    phase3_batch_size: int = 4
    phase3_lr_backbone: float = 1e-5
    phase3_lr_fusion: float = 5e-4
    phase3_weight_decay: float = 1e-4

    # Loss weights
    w_cls_fused: float = 1.0
    w_cls_mri: float = 0.3
    w_cls_tab: float = 0.3
    w_prog: float = 0.5
    w_contrastive: float = 0.2
    focal_gamma: float = 1.0
    label_smoothing: float = 0.05

    grad_clip: float = 1.0
    val_split: float = 0.15
    test_split: float = 0.15      # subject-level holdout, never visit-level
    early_stopping_patience: int = 20
    min_epochs_before_es: int = 25

    save_dir: str = MODEL_SAVE_ROOT


@dataclass
class DataConfig:
    nifti_target_shape: tuple = (96, 96, 96)
    pair_window_months: int = 6            # MRI ↔ biomarker date tolerance
    progression_horizon_months: int = 60   # 5-year look-ahead for MCI → AD
    aug_rotation_range: float = 8.0
    aug_flip_prob: float = 0.5
    aug_noise_std: float = 0.02
    aug_gamma_range: tuple = (0.85, 1.15)

    # Optional hippocampus-focused crop: the brain bounding box is found, then
    # a centre crop is taken at these ratios and resized to nifti_target_shape.
    # x: left-right (both hemispheres), y: anterior-posterior, z: inferior-superior
    hippocampus_crop_enabled: bool = False
    hippo_x_range: tuple = (0.10, 0.90)
    hippo_y_range: tuple = (0.25, 0.70)
    hippo_z_range: tuple = (0.15, 0.65)