# Slurm jobs reuse this kernel's Python interpreter for their Python steps,
# so job results are produced with the same package versions as the notebook.
import os, sys
os.environ["CALMAR_PYTHON"] = sys.executable
print("Notebook and job Python:", sys.executable)

# Install any missing notebook packages, then check them.
# They go to ~/.local (pip --user), which is on the persistent home volume, so a
# Neurodesk restart does not remove them. Pins follow src/constraints-neurodesk.txt;
# packages that are already installed are left as they are.
import importlib, importlib.metadata as _setup_metadata
import site as _setup_site
import subprocess as _setup_subprocess
import sys as _setup_sys

_SETUP_PACKAGES = (
    'numpy', 'pandas', 'scipy', 'matplotlib', 'nibabel', 'nilearn',
    'ipywidgets', 'anywidget', 'ipylab', 'ipyniivue', 'SimpleITK', 'pymupdf', 'pydicom', 'PyYAML',
)
_SETUP_PINS = {'nilearn': 'nilearn==0.13.1', 'ipylab': 'ipylab==1.0.0', 'anywidget': 'anywidget>=0.11,<1',
               'SimpleITK': 'SimpleITK>=2.3', 'pymupdf': 'pymupdf>=1.23', 'ipyniivue': 'ipyniivue>=2.0'}
# Already imported by Jupyter: if pip changed one of these, the kernel must restart.
_SETUP_LOADED = ('numpy', 'pandas', 'packaging', 'ipywidgets')

def _setup_versions(names):
    found = {}
    for name in names:
        try:
            found[name] = _setup_metadata.version(name)
        except _setup_metadata.PackageNotFoundError:
            found[name] = None
    return found

_setup_missing = [p for p, v in _setup_versions(_SETUP_PACKAGES).items() if v is None]
if _setup_missing:
    _CALMAR_SETUP_READY = False
    _setup_before = _setup_versions(_SETUP_LOADED)
    _setup_command = [_setup_sys.executable, '-m', 'pip', 'install', '--user', '--no-input',
                      *[_SETUP_PINS.get(p, p) for p in _setup_missing]]
    print('Installing missing packages:', ', '.join(_setup_missing))
    _setup_result = _setup_subprocess.run(_setup_command, capture_output=True, text=True)
    if _setup_result.returncode != 0:
        raise RuntimeError('Installing ' + ', '.join(_setup_missing) + ' failed:\n'
                           + _setup_result.stdout[-2000:] + _setup_result.stderr[-2000:])
    # ~/.local is only on sys.path at start-up if it existed then.
    if _setup_site.getusersitepackages() not in _setup_sys.path:
        _setup_site.addsitedir(_setup_site.getusersitepackages())
    importlib.invalidate_caches()
    _setup_still = [p for p, v in _setup_versions(_setup_missing).items() if v is None]
    if _setup_still:
        raise RuntimeError('Still missing after installation: ' + ', '.join(_setup_still))
    _setup_changed = [p for p, v in _setup_versions(_SETUP_LOADED).items() if v != _setup_before[p]]
    if _setup_changed:
        raise RuntimeError('Installation changed ' + ', '.join(_setup_changed)
                           + ', which this kernel already loaded. Use Kernel → Restart Kernel and rerun Setup.')
    print('Installed:', ', '.join(_setup_missing))
else:
    print('Setup packages are installed; no installation needed.')

# Imports used throughout the notebook.
_CALMAR_SETUP_READY = False
import sys, importlib, importlib.metadata
import packaging
# pip may have upgraded packages while this kernel kept the old versions in memory.
if packaging.__version__ != importlib.metadata.version('packaging'):
    raise RuntimeError('Python packages changed during this session. Restart the kernel and rerun Setup.')

import os, re, json, shutil, subprocess, warnings
from pathlib import Path
from datetime import datetime

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import nibabel as nib
import ipywidgets as widgets
from IPython.display import display, HTML, clear_output

sys.path.insert(0, str(Path.cwd()))
from calmar import qc as q            # QC records, mask edits; tool helpers run only via in_slurm
from calmar import widgets as cw      # persistent NiiVue viewers (never importlib.reload this)
from calmar import masks as cm        # mask discovery
from calmar.scheduled import submit_step, in_slurm
from calmar.step_summary import show_step

# Aphasia knowledge base: explicit override, sibling working copy, then CALMaR submodule.
from calmar.kb_location import kb_selection_text, resolve_kb_root, update_kb_to_latest
from calmar.kb_paths import load_kb_modules
APHASIA_KB_ROOT, APHASIA_KB_REASON = resolve_kb_root(Path.cwd())
print(update_kb_to_latest(APHASIA_KB_ROOT, APHASIA_KB_REASON))
print(kb_selection_text(APHASIA_KB_ROOT, APHASIA_KB_REASON))
aphasia_kb, = load_kb_modules(APHASIA_KB_ROOT, 'aphasia_kb')

_CALMAR_SETUP_READY = True
print('CALMaR setup imports completed.')

from calmar.timing import install_timing
# Times every cell; BATCH_TIMINGS collects per-participant step times; show_timings() prints a summary.
BATCH_TIMINGS = install_timing(globals(), get_ipython())
print("Step timing active.")

# CONFIGURATION — the single place to change how the workflow runs.
# Every later cell reads CONFIG; edit here, then rerun from this cell.
# Sections: (1) which data, (2) method parameters (scientific choices, printed
# below as a methods table), (3) where outputs go, (4) display and re-run options.
from calmar import colors as cc
from calmar.methods import methods_table

CONFIG = {
    # ── (1) Which data ─────────────────────────────────────────────────────
    # DATASET_SOURCE: "openneuro" (datalad clone of OPENNEURO_ID), "local" (BIDS
    # folder in LOCAL_DATASET), "isles2022" (an extracted local ISLES 2022 root),
    # "flat" (folder of *_T1w.nii.gz), or "single" (SINGLE_T1W).
    "DATASET_SOURCE":  "openneuro",
    # Imaging route: "t1" = chronic stroke on the T1w (LINDA / SynthStroke);
    # "dwi" = acute stroke on diffusion-weighted images + ADC (OpenADS), e.g. ISLES-2022
    # (DATASET_SOURCE "isles2022", LOCAL_DATASET "data/ISLES-2022"; see "Acute route" below).
    # Whichever route is selected, participant discovery also inventories ordinary DWI
    # series in the same or other sessions. Only an exact DWI+ADC pair enters OpenADS.
    "MODALITY":        "t1",
    "ANALYSE_DWI":     None,              # DWI + matching ADC lesion analysis (OpenADS): None = ask; True = first pair; False = skip
    "ANALYSE_TRACT_DWI": None,            # gradient-bearing 4-D diffusion MRI: None = ask; True = best valid session; False = skip
    "DMRI_MASK_SOURCE": "manual",       # selected native-T1 lesion mask supplied to lesion-aware QSIPrep registration
    "DMRI_OUTPUT_RESOLUTION": 2.0,         # isotropic mm; avoid upsampling the current 2 mm acquisitions
    "DMRI_HMC_METHOD": "tortoise",       # QSIPrep HMC; eddy metadata is unreadable in the current Neurodesk image
    "DMRI_RECON_SPEC": "ss3t_fod_autotrack", # experimental patient tractometry; absence of a bundle is not disconnection
    "DMRI_OVERWRITE": False,              # preserve prior QSIPrep/QSIRecon derivatives unless replacement is explicit
    "ANALYSE_BRAIN_HEALTH": None,          # None = ask; True = all eligible; False = skip
    "BRAIN_HEALTH_MASK_SOURCE": "manual", # native lesion excluded from cavity/FLAIR burden estimates
    "BRAIN_HEALTH_SYNTHSEG_MODE": "robust", # robust | standard; robust is slower and preferred for pathology
    "BRAIN_HEALTH_LESION_EXCLUSION_MM": 3.0, # margin around registered stroke lesion on FLAIR
    "BRAIN_HEALTH_REGISTRATION_SEED": 1729, # reproducible native-image to FLAIR affine registration
    "BRAIN_HEALTH_OVERWRITE": False,       # preserve prior SynthSeg/DeepWMH outputs
    "STROKE_PHASE":    None,              # None = from days since stroke (DWI route: acute); or "acute" | "subacute" | "chronic"
    "OPENNEURO_ID":    "ds004884",
    "OPENNEURO_TAG":   "1.0.2",           # dataset version to check out (None = latest)
    "DATA_SUBDIR":     "data",            # datasets live under PROJECT_DIR/data/
    "LOCAL_DATASET":   None,
    "SINGLE_T1W":      None,
    # Participants: a list, an exact ID, a regex, or None for everyone.
    # Demonstration cohort with existing complete HD-BET, LINDA, SynthStroke, manual-mask,
    # BCBToolkit and DeepDisco outputs in this workspace. The route-aware additional
    # cohort keeps ISLES inputs/outputs in its own dataset and uses OpenADS.
    "SUBJECT_FILTER":  ["sub-M2002", "sub-M2051", "sub-M2061", "sub-M2075", "sub-M2103"],
    "SESSION_FILTER":  ["ses-1440", "ses-284", "ses-252", "ses-602", "ses-222"],
    "T1W_SELECTION":   "first",           # "first" T1w per session, or "all"
    "MAX_SUBJECTS":    5,                 # None = all matching participants
    "ADDITIONAL_COHORTS": [
        {"NAME": "ISLES 2022 acute", "DATASET_SOURCE": "isles2022",
         "LOCAL_DATASET": "data/ISLES-2022", "MODALITY": "dwi", "STROKE_PHASE": "acute",
         "SEGMENTATION_METHODS": ["openads"],
         "SUBJECT_FILTER": ["sub-strokecase0004", "sub-strokecase0009"],
         "SESSION_FILTER": ["ses-0001"], "MAX_SUBJECTS": 2},
    ],
    "RANDOM_SAMPLE":   False,             # True = random MAX_SUBJECTS participants
    "RANDOM_SEED":     67,
    "RUN_TEST":        True,              # run one participant step by step before the batch
    "TEST_SUBJECT_IDX": 0,                # set by the participant picker below

    # Manual (reference) lesion mask. None = use the masks the dataset publisher
    # shipped (derivatives/lesion_masks), preferring the T1w's own session.
    # Your own masks: a template with {sub}/{ses}, e.g. "/data/masks/{sub}_{ses}_mask.nii.gz".
    "MANUAL_MASK_PATH":  None,
    "MANUAL_MASK_SPACE": "t1",            # "t1" | "other" (drawn on MANUAL_MASK_REF) | "mni"
    "MANUAL_MASK_REF":   None,            # scan the mask was drawn on, when space is "other"

    # ── (2) Method parameters ──────────────────────────────────────────────
    # Automatic lesion segmentations to run (at least one). LINDA is slow (1.5-3.7 h
    # per participant, mostly its registration to the template); without it, ANTs
    # registers each brain to the template instead (registration step). Also tick boxes
    # at the single-subject participant picker and the batch checkpoint.
    "SEGMENTATION_METHODS":    ["linda", "synthstroke"],
    "RUN_SEGMENTATION":        True,      # False = only review existing masks
    "USE_HDBET_PREPROCESSING": True,      # our brain mask replaces LINDA's skull strip
    "BRAIN_MASK":              "hdbet",   # HD-BET; "hdbet+synthstrip" adds SynthStrip's CSF (per participant: skull-strip repair)
    "HDBET_MODE":              "fast",    # "fast" (1 model) or "accurate" (5-model ensemble, slow on CPU)
    "HDBET_DEVICE":            None,      # None/"cpu", or a GPU id such as "0"
    "SYNTHSTROKE_TTA":         True,      # test-time augmentation: more robust, ~2x slower
    "SYNTHSTROKE_TIMEOUT":     5400,      # seconds per participant (CPU-only TTA can exceed 30 min)
    "SYNTHSTROKE_WEIGHTS_DIR": "data/models/synthstroke-synth-plus",  # under PROJECT_DIR; /opt is wiped on restart
    "SYNTHSTROKE_ALLOW_BASELINE": False,  # never mix in the 2-class baseline model
    "REGISTRATION_SEED":       1729,      # random seed of every registration: LINDA's, the registration step's and the manual mask's (reproducibility)
    "LINDA_THREADS":           1,         # 1 = identical LINDA result on every run (1.5-3.7 h per participant);
                                          # 4 = ~2x faster, but reruns differ slightly (Dice ~0.96)
    "REGISTRATION_THREADS":    4,         # registration step (LINDA not run): CPU threads; the result does not depend on it
    "LESION_THRESHOLD":        0.5,       # binarise lesion maps at this value
    "BCB_MASKS":               "all",     # linda | synthstroke | manual | both (linda+manual) | all
    "BCB_INPUT_GRID":          "mni152",  # "mni152": place masks on BCBToolkit's MNI152 grid by world
                                          # coordinates (fixes a 1 mm shift); "native": original behaviour
    "BCB_RUN_TRACTOTRON":      True,     # optional tract-probability tables
    "DEEPDISCO_MASK_SOURCE":   "all",     # linda | synthstroke | manual | all
    "DEEPDISCO_RUN_ALL":       True,      # all four fibre-class models
    "DEEPDISCO_MODELS":        ["whole_brain"],  # used when RUN_ALL is False
    "DEEPDISCO_DOWNSAMPLING":  "coverage", # 1 mm -> 2 mm: "coverage" (>= half the voxel) or "nearest"
    "ATLASES":       ["HarvardOxford", "AAL", "Destrieux", "Schaefer400", "JHU", "XTRACT"],
    "DEFAULT_ATLAS": "HarvardOxford",

    # ── (3) Where outputs go ───────────────────────────────────────────────
    # Image derivatives go inside the dataset (BIDS derivatives/); reports,
    # QC figures and job logs go under PROJECT_DIR/reports.
    "PROJECT_DIR":         Path.cwd(),
    "DERIVATIVES_SUBDIR":  "derivatives/linda",
    "HDBET_SUBDIR":        "derivatives/hdbet",
    "SYNTHSTROKE_SUBDIR":  "derivatives/synthstroke",
    "BCB_SUBDIR":          "derivatives/bcbtoolkit",
    "DEEPDISCO_SUBDIR":    "derivatives/deepdisco",
    "REGISTRATION_SUBDIR": "derivatives/registration",   # template registration when LINDA is not run
    "OPENADS_SUBDIR":      "derivatives/openads",        # acute (DWI) route: OpenADS lesion and its registration
    "REPORTS_SUBDIR":      "reports",
    "ATLAS_SUBDIR":        "atlases",

    # ── (4) Display and re-run options ─────────────────────────────────────
    # Outputs are never silently replaced. When a method option above changes,
    # existing outputs are reported as stale; set the matching *_OVERWRITE to
    # recompute (the old output is kept as *.previous.*).
    "OVERWRITE":           False,         # HD-BET, LINDA, SynthStroke, manual registration
    "BCB_OVERWRITE":       False,
    "DEEPDISCO_OVERWRITE": False,
    "MASK_COLORS": {"linda": "red", "manual": "blue", "synthstroke": "green",
                    "hdbet": "yellow", "bcbtoolkit": "magenta", "deepdisco": "cyan"},
    "NIIVUE_URL_BASE": False,             # viewer loading: False = works everywhere (see README)
    "VERBOSE": True,
}

# The guided notebook passes its form values here.
CONFIG.update(globals().get("_CALMAR_CONFIG_OVERRIDES", {}))
if CONFIG["DATASET_SOURCE"] == "isles2022":
    CONFIG["MODALITY"] = "dwi"
    CONFIG["SEGMENTATION_METHODS"] = ["openads"]
from calmar.study import segmentation_methods
segmentation_methods(CONFIG)   # an invalid SEGMENTATION_METHODS fails here, not in a Slurm job

# ── Derived paths (no need to edit) ──────────────────────────────────────────
RUN_TEST = CONFIG["RUN_TEST"]
TEST_SUBJECT_IDX = CONFIG["TEST_SUBJECT_IDX"]
_SINGLE_OPTIONAL_COMPLETED = {}  # reset when configuration is rerun
PROJECT_DIR = Path(CONFIG["PROJECT_DIR"]).resolve()
if CONFIG["DATASET_SOURCE"] in ("local", "flat", "isles2022"):
    RAW_DATASET_DIR = Path(CONFIG["LOCAL_DATASET"]).expanduser().resolve()
elif CONFIG["DATASET_SOURCE"] == "single":
    RAW_DATASET_DIR = Path(CONFIG["SINGLE_T1W"]).expanduser().resolve().parent
else:
    RAW_DATASET_DIR = PROJECT_DIR / CONFIG["DATA_SUBDIR"] / CONFIG["OPENNEURO_ID"]
DERIV_DIR   = RAW_DATASET_DIR / CONFIG["DERIVATIVES_SUBDIR"]   # LINDA outputs
REPORTS_DIR = PROJECT_DIR / CONFIG["REPORTS_SUBDIR"]
ATLAS_DIR   = PROJECT_DIR / CONFIG["ATLAS_SUBDIR"]
# Derivative folders inside the dataset are created by the jobs, after `datalad install`.
for d in (ATLAS_DIR, REPORTS_DIR):
    d.mkdir(parents=True, exist_ok=True)

cc.configure(CONFIG["MASK_COLORS"])
cw.configure(url_base=CONFIG["NIIVUE_URL_BASE"])
# The %%bash fetch cell reads these.
os.environ.update(RAW_DATASET_DIR=str(RAW_DATASET_DIR), DATASET_SOURCE=CONFIG["DATASET_SOURCE"],
                  OPENNEURO_ID=CONFIG.get("OPENNEURO_ID") or "", OPENNEURO_TAG=CONFIG.get("OPENNEURO_TAG") or "",
                  LOCAL_DATASET=str(CONFIG["LOCAL_DATASET"] or ""), SINGLE_T1W=str(CONFIG["SINGLE_T1W"] or ""),
                  NILEARN_DATA=str(ATLAS_DIR))

print(f"Dataset:  {RAW_DATASET_DIR}\nReports:  {REPORTS_DIR}\n")
display(methods_table(CONFIG).style.hide(axis="index")
        .set_properties(**{"text-align": "left", "white-space": "normal", "vertical-align": "top"})
        .set_table_styles([{"selector": "th", "props": [("text-align", "left"), ("vertical-align", "top")]}]))

# Raised by checkpoint cells to stop "Run all" cleanly (no traceback) until you click a button.
class StopExecution(Exception):
    def _render_traceback_(self):
        return []

# Prepare atlases in a retained Slurm job; the kernel loads only published descriptors.
from calmar.reporting import load_atlases
FORCE_RELOAD_ATLASES = False
ALL_ATLASES = load_atlases(globals(), force=FORCE_RELOAD_ATLASES)
print("Atlases ready:", list(ALL_ATLASES))

%%bash
# Fetch the source dataset (metadata + git-annex pointers only)
# ------------------------------------------------------------
# Lightweight: `datalad install` clones the dataset's structure but
# does NOT download large files. Subject T1w content is fetched on
# demand by the discovery cell — only for subjects we actually
# sample (so we don't waste disk on subjects we will not process).
# Manual lesion masks + T2w are similarly fetched on demand by the
# comparison cell.
# ============================================================
set -uo pipefail
shopt -s nullglob

case "$DATASET_SOURCE" in
  openneuro)
      if [ -z "$OPENNEURO_ID" ]; then
          echo "❌ DATASET_SOURCE=openneuro but OPENNEURO_ID is empty"; exit 1;
      fi
      # Real datalad checkout? Look for .datalad/ — the directory
      # alone may exist as an empty placeholder if other cells\'
      # mkdir() pre-created it (e.g. derivatives subtree).
      if [ ! -d "$RAW_DATASET_DIR/.datalad" ]; then
          # Empty placeholder? Wipe it so `datalad install` can
          # populate cleanly (it refuses to install into a
          # non-empty directory).
          if [ -d "$RAW_DATASET_DIR" ]; then
              if [ -z "$(ls -A "$RAW_DATASET_DIR/derivatives" 2>/dev/null \
                          | grep -v -E "^(linda|hdbet)$")" ] \
                 && [ -z "$(ls -A "$RAW_DATASET_DIR" 2>/dev/null \
                            | grep -v -E "^(derivatives)$")" ]; then
                  echo "  removing empty placeholder at $RAW_DATASET_DIR"
                  rm -rf "$RAW_DATASET_DIR"
              else
                  echo "❌ $RAW_DATASET_DIR exists, contains files, but is "
                  echo "   NOT a datalad checkout (no .datalad/ folder)."
                  echo "   Either rename it out of the way or delete it"
                  echo "   manually, then re-run this cell."
                  exit 1
              fi
          fi
          echo "📥 datalad install $OPENNEURO_ID → $RAW_DATASET_DIR"
          datalad install -s "https://github.com/OpenNeuroDatasets/${OPENNEURO_ID}.git" \
              "$RAW_DATASET_DIR"
          if [ -n "$OPENNEURO_TAG" ]; then
              ( cd "$RAW_DATASET_DIR" && git checkout "$OPENNEURO_TAG" )
          fi
      else
          echo "✓ Using existing datalad checkout at $RAW_DATASET_DIR"
      fi

      # Quick stats: how many subjects + T1w stubs are visible?
      cd "$RAW_DATASET_DIR"
      n_subj=$(ls -d sub-*/ 2>/dev/null | wc -l | tr -d ' ')
      n_t1=$(find . -name "*_T1w.nii.gz" 2>/dev/null | wc -l | tr -d ' ')
      echo "   $n_subj subject(s) registered, $n_t1 T1w pointer(s) available"
      echo "   (no content downloaded yet — discovery cell handles that)"
      ;;

  local)
      if [ -z "$LOCAL_DATASET" ] || [ ! -d "$LOCAL_DATASET" ]; then
          echo "❌ LOCAL_DATASET is not a directory: $LOCAL_DATASET"; exit 1;
      fi
      echo "✓ Using local BIDS dataset at $LOCAL_DATASET"
      n_subj=$(ls -d "$LOCAL_DATASET"/sub-*/ 2>/dev/null | wc -l | tr -d ' ')
      n_t1=$(find "$LOCAL_DATASET" -name "*_T1w.nii.gz" 2>/dev/null | wc -l | tr -d ' ')
      echo "   $n_subj subject(s), $n_t1 T1w file(s)"
      ;;

  isles2022)
      if [ -z "$LOCAL_DATASET" ] || [ ! -d "$LOCAL_DATASET" ]; then
          echo "❌ LOCAL_DATASET is not a directory: $LOCAL_DATASET"; exit 1;
      fi
      n_pairs=0
      while IFS= read -r dwi; do
          adc="${dwi%_dwi.nii.gz}_adc.nii.gz"
          if [ -f "$adc" ] || [ -L "$adc" ]; then n_pairs=$((n_pairs + 1)); fi
      done < <(find "$LOCAL_DATASET" -type d -name derivatives -prune -o -name "*_dwi.nii.gz" -print 2>/dev/null)
      if [ "$n_pairs" -eq 0 ]; then
          echo "❌ No ISLES-style sibling DWI+ADC pairs found under $LOCAL_DATASET"; exit 1;
      fi
      echo "✓ Using extracted ISLES 2022 dataset at $LOCAL_DATASET"
      echo "   $n_pairs DWI+ADC pair(s); expert *_msk masks are discovered from derivatives/"
      ;;

  flat)
      if [ -z "$LOCAL_DATASET" ] || [ ! -d "$LOCAL_DATASET" ]; then
          echo "❌ LOCAL_DATASET is not a directory: $LOCAL_DATASET"; exit 1;
      fi
      echo "✓ Using flat folder at $LOCAL_DATASET"
      n_t1=$(ls "$LOCAL_DATASET"/*_T1w.nii.gz 2>/dev/null | wc -l | tr -d ' ')
      echo "   $n_t1 T1w file(s)"
      ;;

  single)
      if [ -z "$SINGLE_T1W" ] || [ ! -f "$SINGLE_T1W" ]; then
          echo "❌ SINGLE_T1W is not a file: $SINGLE_T1W"; exit 1;
      fi
      echo "✓ Single T1w: $SINGLE_T1W"
      ;;

  *)
      echo "❌ Unknown DATASET_SOURCE: $DATASET_SOURCE"; exit 1;
      ;;
esac

# Select the processing entries, define their output locations, then fetch selected content.
from calmar.study import Study, config_for_entry, select_cohorts
from calmar.scheduled import submit_step
from calmar.step_summary import show_mask_inventory, show_step

SUBJECTS, _sampling_note = select_cohorts(CONFIG, RAW_DATASET_DIR, PROJECT_DIR)
study = Study(CONFIG, RAW_DATASET_DIR, PROJECT_DIR)
globals().update(study.notebook_helpers())

_routes = {config_for_entry(CONFIG, entry).get('MODALITY', 't1') for entry in SUBJECTS}
_route_label = ("mixed T1w / DWI+ADC entries" if len(_routes) > 1 else
                ("DWI+ADC pairs" if 'dwi' in _routes else "T1w images"))
print(f"Selected {len(SUBJECTS)} {_route_label}. {_sampling_note}")

try:
    acquisition = submit_step("acquisition", SUBJECTS, study, timings=BATCH_TIMINGS)
except Exception:
    print("Acquisition failed for selected entries:")
    for _entry in SUBJECTS:
        print(f"  • {study.label(_entry)}")
    raise
show_step(acquisition, title="Selected inputs and DWI availability")
show_mask_inventory(study, SUBJECTS)

from calmar.dwi_choice import show_dwi_choice

_dwi_choice = show_dwi_choice(globals())
if _dwi_choice is not None:
    raise StopExecution("Choose whether to analyse the available DWI in the single-subject test and batch.")

from calmar.dmri_choice import show_dmri_choice

_dmri_choice = show_dmri_choice(globals())
if _dmri_choice is not None:
    raise StopExecution("Choose whether to run diffusion tract-microstructure analysis in the single-subject test and after lesion QC for the batch.")

from calmar.brain_health_choice import show_brain_health_choice

_brain_health_choice = show_brain_health_choice(globals())
if _brain_health_choice is not None:
    raise StopExecution("Choose whether to add structural brain-health analysis in the single-subject test and after lesion QC for the batch.")

if not CONFIG.get("RUN_TEST", False):
    print("RUN_TEST = False — continuing into batch processing")
else:
    from calmar.single_subject import show_single_subject_inputs
    _single_subject_inputs = show_single_subject_inputs(globals())
    raise StopExecution("Choose a participant and click Run single-subject test.")

from calmar.scheduled import submit_step
from calmar.step_summary import show_step

if not RUN_TEST:
    print("RUN_TEST = False — continuing into batch processing")
else:
    test_entry = SUBJECTS[TEST_SUBJECT_IDX]          # chosen in the cell above
    brain = submit_step("brain", [test_entry], study, timings=BATCH_TIMINGS)
    show_step(brain, study, title="Brain extraction (HD-BET)")

if not RUN_TEST:
    print("RUN_TEST = False — continuing into batch processing")
else:
    test_entry = SUBJECTS[TEST_SUBJECT_IDX]          # chosen with the participant picker
    linda = submit_step("linda", [test_entry], study, timings=BATCH_TIMINGS)
    show_step(linda, study, title="LINDA lesion segmentation")

if not RUN_TEST:
    print("RUN_TEST = False — continuing into batch processing")
else:
    test_entry = SUBJECTS[TEST_SUBJECT_IDX]          # chosen with the participant picker
    synthstroke = submit_step("synthstroke", [test_entry], study, timings=BATCH_TIMINGS)
    show_step(synthstroke, study, title="SynthStroke lesion segmentation")

if not RUN_TEST:
    print("RUN_TEST = False — continuing into batch processing")
else:
    test_entry = SUBJECTS[TEST_SUBJECT_IDX]          # chosen with the participant picker
    registration = submit_step("registration", [test_entry], study, timings=BATCH_TIMINGS)
    show_step(registration, study, title="Template registration (ANTs, without LINDA)")

if not RUN_TEST:
    print("RUN_TEST = False — continuing into batch processing")
else:
    test_entry = SUBJECTS[TEST_SUBJECT_IDX]          # chosen with the participant picker
    openads = submit_step("openads", [test_entry], study, timings=BATCH_TIMINGS)
    show_step(openads, study, title="OpenADS acute lesion segmentation (DWI)")

if not RUN_TEST:
    print("RUN_TEST = False — continuing into batch processing")
else:
    test_entry = SUBJECTS[TEST_SUBJECT_IDX]          # chosen with the participant picker
    manual = submit_step("manual", [test_entry], study, timings=BATCH_TIMINGS)
    show_step(manual, title="Manual mask registration and agreement")
    # Keep native registration and template-space inspection in one tabbed view.
    import ipywidgets as widgets
    from calmar.study import stage_selected
    from calmar.step_viewer import show_step_viewer
    _manual_views = []
    _native_view = show_step_viewer("manual", [test_entry], study, key="step_manual", display_panel=False)
    if _native_view is not None:
        _manual_views.append(("Native mask and reference scan", _native_view["panel"]))
    if (stage_selected("registration", globals().get("CONFIG", getattr(study, "config", {})))
            and cm.present(study.template_dir(test_entry) / "ExpertMask_in_MNI.nii.gz")):
        _mni_view = show_step_viewer("registration", [test_entry], study, key="step_registration_manual", display_panel=False)
        if _mni_view is not None:
            _manual_views.append(("Template-space mask", _mni_view["panel"]))
    if len(_manual_views) > 1:
        _manual_tabs = widgets.Tab(children=[panel for _, panel in _manual_views], selected_index=0)
        for _i, (_title, _) in enumerate(_manual_views): _manual_tabs.set_title(_i, _title)
        display(_manual_tabs)
    elif _manual_views:
        display(_manual_views[0][1])

if not RUN_TEST:
    print("RUN_TEST = False — continuing into batch processing")
else:
    from calmar.single_subject import mark_optional_completed, optional_entries_for_test
    _single_optional_completed = globals().setdefault("_SINGLE_OPTIONAL_COMPLETED", {})
    _single_optional_skipped = []
    _test_entry = SUBJECTS[TEST_SUBJECT_IDX]

    _test_dwi = optional_entries_for_test(globals().get("DWI_SUBJECTS"), _test_entry)
    if CONFIG.get("MODALITY", "t1") == "dwi":
        print("DWI is the primary route; its OpenADS test already ran in the main single-subject cells.")
    elif _test_dwi:
        try:
            test_dwi_acquisition = submit_step("acquisition", _test_dwi, DWI_STUDY, timings=BATCH_TIMINGS)
            show_step(test_dwi_acquisition, title="Test participant DWI+ADC acquisition", images=True)
            test_dwi_openads = submit_step("openads", _test_dwi, DWI_STUDY, timings=BATCH_TIMINGS)
            show_step(test_dwi_openads, DWI_STUDY, title="Test participant OpenADS acute DWI+ADC segmentation", images=True)
            if not mark_optional_completed(_single_optional_completed, "dwi", _test_dwi, DWI_STUDY,
                                           [test_dwi_acquisition, test_dwi_openads]):
                print("DWI completion evidence could not be recorded; the later batch will rerun it safely.")
        except Exception as exc:
            DWI_ERROR = str(exc)
            print(f"Optional DWI test failed; it was not marked complete and the lesion test will continue. {exc}")
    else:
        _single_optional_skipped.append("acute DWI+ADC/OpenADS lesion segmentation")

    _test_dmri = optional_entries_for_test(globals().get("DMRI_SUBJECTS"), _test_entry)
    if _test_dmri:
        try:
            test_dmri_preproc = submit_step("dmri_preproc", _test_dmri, DMRI_STUDY, timings=BATCH_TIMINGS)
            show_step(test_dmri_preproc, title="Test participant lesion-aware QSIPrep", images=True, viewer=False)
            test_dmri_recon = submit_step("dmri_recon", _test_dmri, DMRI_STUDY, timings=BATCH_TIMINGS)
            show_step(test_dmri_recon, title="Test participant experimental tractometry", images=True, viewer=False)
            test_dmri_metrics = submit_step("dmri_metrics", _test_dmri, DMRI_STUDY, timings=BATCH_TIMINGS)
            show_step(test_dmri_metrics, title="Test participant diffusion tract-microstructure metrics", images=True, viewer=False)
            if not mark_optional_completed(_single_optional_completed, "dmri", _test_dmri, DMRI_STUDY,
                                           [test_dmri_preproc, test_dmri_recon, test_dmri_metrics]):
                print("Diffusion tract-microstructure completion evidence could not be recorded; the later batch will rerun it safely.")
            print("Review the QSIPrep/QSIRecon HTML reports before accepting participant tract-microstructure measurements.")
        except Exception as exc:
            DMRI_ERROR = str(exc)
            print(f"Optional diffusion tract-microstructure test failed; it was not marked complete and the lesion test will continue. {exc}")
    else:
        _single_optional_skipped.append("diffusion tract-microstructure analysis")

    _test_brain_health = optional_entries_for_test(globals().get("BRAIN_HEALTH_SUBJECTS"), _test_entry)
    if _test_brain_health:
        try:
            _test_bh_anatomy = [e for e in _test_brain_health if "atrophy" in e["brain_health_components"]]
            _test_bh_wmh = [e for e in _test_brain_health if "wmh" in e["brain_health_components"]]
            _test_brain_health_results = []
            if _test_bh_anatomy:
                test_brain_health_anatomy = submit_step("brain_health_anatomy", _test_bh_anatomy, study, timings=BATCH_TIMINGS)
                show_step(test_brain_health_anatomy, title="Test participant T1 volumetry (SynthSeg)", images=True, viewer=False)
                _test_brain_health_results.append(test_brain_health_anatomy)
            if _test_bh_wmh:
                test_brain_health_wmh = submit_step("brain_health_wmh", _test_bh_wmh, study, timings=BATCH_TIMINGS)
                show_step(test_brain_health_wmh, title="Test participant FLAIR WMH candidates (DeepWMH)", images=True, viewer=False)
                _test_brain_health_results.append(test_brain_health_wmh)
            test_brain_health_metrics = submit_step("brain_health_metrics", _test_brain_health, study, timings=BATCH_TIMINGS)
            show_step(test_brain_health_metrics, title="Test participant stroke-aware brain-health features", images=True, viewer=False)
            _test_brain_health_results.append(test_brain_health_metrics)
            if not mark_optional_completed(_single_optional_completed, "brain_health", _test_brain_health, study,
                                           _test_brain_health_results):
                print("Brain-health completion evidence could not be recorded; the later batch will rerun it safely.")
            print("Review each brain_health_qc.png before treating these provisional features as interpretable.")
        except Exception as exc:
            BRAIN_HEALTH_ERROR = str(exc)
            print(f"Optional brain-health test failed; it was not marked complete and the lesion test will continue. {exc}")
    else:
        _single_optional_skipped.append("structural brain-health analysis")

    if _single_optional_skipped:
        print("Optional analyses not selected or not available for this test participant "
              "(selections for other participants run in the batch section): "
              + "; ".join(_single_optional_skipped) + ".")

if not RUN_TEST:
    print("RUN_TEST = False — continuing into batch processing")
else:
    test_entry = SUBJECTS[TEST_SUBJECT_IDX]          # chosen with the participant picker
    bcb = submit_step("bcb", [test_entry], study, timings=BATCH_TIMINGS)
    show_step(bcb, study, title="BCBToolkit disconnectome")

if not RUN_TEST:
    print("RUN_TEST = False — continuing into batch processing")
else:
    test_entry = SUBJECTS[TEST_SUBJECT_IDX]          # chosen with the participant picker
    deepdisco = submit_step("deepdisco", [test_entry], study, timings=BATCH_TIMINGS)
    show_step(deepdisco, study, title="DeepDisco disconnectome")

# BCBToolkit vs DeepDisco — interactive comparison for the test participant.
COMPARE_MASK_SOURCE = None          # "linda" | "synthstroke" | "manual"; None = first source with maps
COMPARE_MODEL = "whole_brain"       # DeepDisco map shown first: whole_brain | association | projection | commissural

if not RUN_TEST:
    print('RUN_TEST = False — continuing into batch processing')
else:
    from calmar.disconnectome_compare import show_disconnectome_comparison
    entry = SUBJECTS[TEST_SUBJECT_IDX]
    # One NiiVue viewer; the participant menu also offers the other discovered participants.
    # Agreement metrics (Pearson r, Dice > 0.5) run as a Slurm reporting job when their section is opened.
    _disconnectome_comparison = show_disconnectome_comparison(
        study, entry, mask_source=COMPARE_MASK_SOURCE, model=COMPARE_MODEL, entries=SUBJECTS)

# Pause after the single-subject test; skipped in batch mode (RUN_TEST = False).
if RUN_TEST:
    show_timings(min_seconds=1.0)
    from calmar.execution import pause_before_batch
    _batch_checkpoint = pause_before_batch()
    raise StopExecution("Choose Continue to batch or Jump to reporting.")
else:
    print("RUN_TEST = False — continuing into batch processing")

from calmar.scheduled import submit_step
from calmar.step_summary import show_step

brain = submit_step("brain", SUBJECTS, study, timings=BATCH_TIMINGS)
show_step(brain, study, title="Brain extraction (HD-BET) — all participants")
linda = submit_step("linda", SUBJECTS, study, timings=BATCH_TIMINGS)
show_step(linda, study, title="LINDA lesion segmentation — all participants")

synthstroke = submit_step("synthstroke", SUBJECTS, study, timings=BATCH_TIMINGS)
show_step(synthstroke, study, title="SynthStroke lesion segmentation — all participants")

registration = submit_step("registration", SUBJECTS, study, timings=BATCH_TIMINGS)
show_step(registration, study, title="Template registration (ANTs, without LINDA) — all participants")

openads = submit_step("openads", SUBJECTS, study, timings=BATCH_TIMINGS)
show_step(openads, study, title="OpenADS acute lesion segmentation (DWI) — all participants")

manual = submit_step("manual", SUBJECTS, study, timings=BATCH_TIMINGS)
show_step(manual, title="Manual mask registration and agreement — all participants")
# Native registration and template-space inspection share one tabbed view.
import ipywidgets as widgets
from calmar.study import stage_selected
from calmar.step_viewer import show_step_viewer
_manual_views = []
_native_view = show_step_viewer("manual", SUBJECTS, study, key="step_manual_batch", display_panel=False)
if _native_view is not None:
    _manual_views.append(("Native mask and reference scan", _native_view["panel"]))
_with_manual_mni = [entry for entry in SUBJECTS
                    if cm.present(study.template_dir(entry) / "ExpertMask_in_MNI.nii.gz")]
if (stage_selected("registration", globals().get("CONFIG", getattr(study, "config", {})))
        and _with_manual_mni):
    _mni_view = show_step_viewer("registration", _with_manual_mni, study, key="step_registration_manual", display_panel=False)
    if _mni_view is not None:
        _manual_views.append(("Template-space mask", _mni_view["panel"]))
if len(_manual_views) > 1:
    _manual_tabs = widgets.Tab(children=[panel for _, panel in _manual_views], selected_index=0)
    for _i, (_title, _) in enumerate(_manual_views): _manual_tabs.set_title(_i, _title)
    display(_manual_tabs)
elif _manual_views:
    display(_manual_views[0][1])

%calmar_continue skip-qc-to-report

# QC review panel. Pauses the notebook until you choose Do QC, Use existing QC or Skip QC.
# Display settings only; they do not change what is saved.
from calmar.qc_panel import QCDisplay, open_qc_review

LESION_OPACITY = 0.80              # lesion overlay opacity (0 = invisible, 1 = opaque)
LESION_AS_BOUNDARY = True          # draw only the lesion outline (cached next to the mask as *_boundary_tN.nii.gz)
LESION_BOUNDARY_THICKNESS = 1      # outline width in voxels (1 thin – 3 thick)
BRAIN_MASK_OPACITY = 0.30          # brain-mask overlay when the extracted brain cannot be shown
SKULL_STRIP_SHOW_EXTRACTED = True  # skull strip: show N4 T1 × brain mask, so leftover skull/dura is visible
MANUAL_OVERLAY_OPACITY = 0.50      # manual-mask overlay (native reference and MNI warp)

_qc_panel = open_qc_review(
    SUBJECTS, study,
    QCDisplay(lesion_opacity=LESION_OPACITY, lesion_as_boundary=LESION_AS_BOUNDARY,
              lesion_boundary_thickness=LESION_BOUNDARY_THICKNESS,
              brain_mask_opacity=BRAIN_MASK_OPACITY,
              skull_strip_show_extracted=SKULL_STRIP_SHOW_EXTRACTED,
              manual_overlay_opacity=MANUAL_OVERLAY_OPACITY),
    viewer=cw.fresh_viewer("qc_dashboard", height=400), encode_volumes=cw.vols)
raise StopExecution("Use existing QC, review/redo QC, or choose Skip QC. Use the buttons above to continue.")

# Skull-strip repair: HD-BET again, then LINDA with the new brain mask (both in Slurm).
from calmar.qc_repair import run_skull_strip_repair

REPAIR_MODE = "step_marked"     # "step_marked": next queued participant (rerun for the next)
                                # "batch_marked": every queued participant; "single": REPAIR_SUBJECT_IDX
REPAIR_SUBJECT_IDX = 0          # position in SUBJECTS (0-based); used only by "single"
DRY_RUN = False                 # True = HD-BET only; LINDA, ratings and the queue are left unchanged
HDBET_REPAIR_MODE = "accurate"  # "accurate" (5-model ensemble, slower) or "fast"
HDBET_DEVICE = None             # None = GPU if the Slurm node has one, else CPU; or "cpu" / "0"
HDBET_TTA = False               # HD-BET test-time augmentation (dry run only; much slower)
REPAIR_BRAIN_MASK = "auto"      # "auto": HD-BET OR SynthStrip if the skull strip is tagged lesion_excluded_by_strip,
                                # else CONFIG["BRAIN_MASK"]; or force "hdbet" / "hdbet+synthstrip"
ONLY_SKULL_STRIP_ISSUES = True  # queue only participants whose skull strip is rated 3 or has an issue tag
REVIEWER = ""                   # name recorded in the edit log

repair_results = run_skull_strip_repair(
    SUBJECTS, study, CONFIG, mode=REPAIR_MODE, subject_index=REPAIR_SUBJECT_IDX,
    dry_run=DRY_RUN, hdbet_mode=HDBET_REPAIR_MODE, device=HDBET_DEVICE, tta=HDBET_TTA,
    brain_mask=REPAIR_BRAIN_MASK,
    only_skull_strip_issues=ONLY_SKULL_STRIP_ISSUES, reviewer=REVIEWER or None, viewers=cw)

# Repair 2 — deterministic clean-up of LINDA's native lesion (re-warp to MNI runs in Slurm).
from calmar.qc_repair import MaskCleanupPanel

REVIEWER = ""   # name recorded with each saved edit

MaskCleanupPanel(SUBJECTS, study, cw.fresh_viewer("qc_edit_preview", height=300),
                 encode_volumes=cw.vols, reviewer=REVIEWER)

# Repair 3 — edit LINDA's native lesion by hand in an external editor, then import it.
# Export copies the native T1 + lesion to QC_EDITS_DIR/<sub>_<ses>/ and shows the command to open
# them on the Neurodesktop. Import checks the grid matches, replaces Prediction3_native.nii.gz
# (LINDA's prediction kept in _linda_original/) and re-warps to MNI in Slurm.
from calmar.qc_repair import ExternalEditPanel

QC_EDITS_DIR = PROJECT_DIR / "qc_edits"   # working copies for the editor
EDITOR = "ITK-SNAP"                       # default editor: "ITK-SNAP", "FSLeyes" or "MRIcroGL"
REVIEWER = ""                             # name recorded with the import

ExternalEditPanel(SUBJECTS, study, cw.fresh_viewer("qc_export_preview", height=300),
                  edits_dir=QC_EDITS_DIR, encode_volumes=cw.vols, reviewer=REVIEWER, editor=EDITOR)

# Repair 4 — re-run LINDA for participants marked for repair, using the best existing HD-BET mask
# (accurate > fast; HD-BET runs first if there is none). Runs in Slurm; the notebook waits.
from calmar.qc_repair import RerunQueuePanel

ALLOW_OVERWRITING_EDITS = False   # True = allow a re-run that discards manual edits made since the last LINDA run
REVIEWER = ""                     # name recorded in the edit log

RerunQueuePanel(SUBJECTS, study, CONFIG, allow_overwriting_edits=ALLOW_OVERWRITING_EDITS,
                reviewer=REVIEWER or None)

# QC summary: one row per participant with the saved rating of each stage (1/2/3, blank = unrated,
# n/a = no output to rate), issue tags, repair flag and number of logged edits.
from calmar.qc_repair import show_qc_summary

QC_SUMMARY_CSV = DERIV_DIR / "qc_summary.csv"   # written on every run

qc_summary = show_qc_summary(SUBJECTS, study, QC_SUMMARY_CSV)

# Reset saved QC ratings (dry run by default).
from calmar.qc_repair import reset_qc

RESET_SUBJECT = None   # e.g. "sub-M2040"; None = every participant in SUBJECTS
RESET_SESSION = None   # e.g. "ses-842"; None = every session
RESET_STAGE = None     # "skull_strip", "lesion", "synthstroke_lesion", "manual_lesion",
                       # "expert_mni_warp"; None = all stages + reviewer and repair flag
KEEP_EDITS = True      # False (with RESET_STAGE = None) deletes the sidecar, edit history included
DRY_RUN = True         # True = only list what would change

affected = reset_qc(SUBJECTS, study, subject=RESET_SUBJECT, session=RESET_SESSION,
                    stage=RESET_STAGE, keep_edits=KEEP_EDITS, dry_run=DRY_RUN)
print(("Would reset" if DRY_RUN else "Reset") + f" {len(affected)} sidecar(s):",
      *[f"  {path}" for path in affected], sep="\n")
if DRY_RUN and affected:
    print("Set DRY_RUN = False and rerun to apply.")

if CONFIG.get("MODALITY", "t1") == "dwi":
    print("DWI is the primary route and its OpenADS outputs were already reviewed in the main QC step.")
elif not globals().get("DWI_SUBJECTS"):
    print("Acute DWI+ADC/OpenADS lesion segmentation was unavailable or not selected. The separate diffusion tract-microstructure route is handled below.")
else:
    from calmar.qc_panel import open_qc_review
    from calmar.single_subject import pending_optional_entries
    _dwi_pending = pending_optional_entries(
        DWI_SUBJECTS, globals().get("_SINGLE_OPTIONAL_COMPLETED", {}).get("dwi"),
        route="dwi", study=DWI_STUDY)
    if _dwi_pending:
        dwi_acquisition = submit_step("acquisition", _dwi_pending, DWI_STUDY, timings=BATCH_TIMINGS)
        show_step(dwi_acquisition, title="Selected DWI+ADC acquisition")
        dwi_openads = submit_step("openads", _dwi_pending, DWI_STUDY, timings=BATCH_TIMINGS)
        show_step(dwi_openads, DWI_STUDY, title="Optional OpenADS DWI analysis", viewer=False)
    else:
        print("All selected optional DWI entries completed in the single-subject test; opening their QC.")
    _dwi_qc_panel = open_qc_review(
        DWI_SUBJECTS, DWI_STUDY, viewer=cw.fresh_viewer("dwi_qc_dashboard", height=400),
        encode_volumes=cw.vols, continue_role="after-dwi-qc",
        continue_label="Continue after DWI QC",
        continue_prompt="Continue to disconnectomes and reports using the reviewed DWI outputs?")
    raise StopExecution("Review or explicitly skip the optional DWI QC, then continue.")

if not globals().get("DMRI_SUBJECTS"):
    print("Diffusion tract-microstructure analysis was unavailable or not selected; this is separate from OpenADS acute-lesion segmentation.")
else:
    from calmar.single_subject import pending_optional_entries
    _dmri_pending = pending_optional_entries(
        DMRI_SUBJECTS, globals().get("_SINGLE_OPTIONAL_COMPLETED", {}).get("dmri"),
        route="dmri", study=DMRI_STUDY)
    if not _dmri_pending:
        print("All selected diffusion tract-microstructure entries completed in the single-subject test.")
    else:
        try:
            dmri_preproc = submit_step("dmri_preproc", _dmri_pending, DMRI_STUDY, timings=BATCH_TIMINGS)
            show_step(dmri_preproc, title="Lesion-aware QSIPrep diffusion preprocessing", viewer=False)
            dmri_recon = submit_step("dmri_recon", _dmri_pending, DMRI_STUDY, timings=BATCH_TIMINGS)
            show_step(dmri_recon, title="Experimental participant tractometry (QSIRecon)", viewer=False)
            dmri_metrics = submit_step("dmri_metrics", _dmri_pending, DMRI_STUDY, timings=BATCH_TIMINGS)
            show_step(dmri_metrics, title="Diffusion tract-microstructure metrics for participant reports", viewer=False)
            print("Review the QSIPrep/QSIRecon HTML reports before accepting participant tract-microstructure measurements.")
        except Exception as exc:
            DMRI_ERROR = str(exc)
            print(f"Optional diffusion tract-microstructure analysis failed; standard disconnectomes and reports will continue. {exc}")

if not globals().get("BRAIN_HEALTH_SUBJECTS"):
    print("Brain-health analysis unavailable or not selected.")
else:
    from calmar.single_subject import pending_optional_entries
    _brain_health_pending = pending_optional_entries(
        BRAIN_HEALTH_SUBJECTS, globals().get("_SINGLE_OPTIONAL_COMPLETED", {}).get("brain_health"),
        route="brain_health", study=study)
    if not _brain_health_pending:
        print("All selected brain-health entries completed in the single-subject test.")
    else:
        try:
            _bh_anatomy = [e for e in _brain_health_pending if "atrophy" in e["brain_health_components"]]
            _bh_wmh = [e for e in _brain_health_pending if "wmh" in e["brain_health_components"]]
            if _bh_anatomy:
                brain_health_anatomy = submit_step("brain_health_anatomy", _bh_anatomy, study, timings=BATCH_TIMINGS)
                show_step(brain_health_anatomy, title="Atrophy-sensitive T1 volumetry (SynthSeg)", viewer=False)
            if _bh_wmh:
                brain_health_wmh = submit_step("brain_health_wmh", _bh_wmh, study, timings=BATCH_TIMINGS)
                show_step(brain_health_wmh, title="Raw FLAIR hyperintensity segmentation (DeepWMH)", viewer=False)
            brain_health_metrics = submit_step("brain_health_metrics", _brain_health_pending, study, timings=BATCH_TIMINGS)
            show_step(brain_health_metrics, title="Stroke-aware brain-health features for reports", viewer=False)
            print("Review each brain_health_qc.png before treating these provisional features as interpretable.")
        except Exception as exc:
            BRAIN_HEALTH_ERROR = str(exc)
            print(f"Optional brain-health analysis failed; lesion and disconnectome reports will continue. {exc}")

from calmar.scheduled import submit_step
from calmar.step_summary import show_step

bcb = submit_step("bcb", SUBJECTS, study, timings=BATCH_TIMINGS)
show_step(bcb, study, title="BCBToolkit disconnectome — all participants")

deepdisco = submit_step("deepdisco", SUBJECTS, study, timings=BATCH_TIMINGS)
show_step(deepdisco, study, title="DeepDisco disconnectome — all participants")

# Compute lesion–atlas overlap tables in Slurm (one job; cached results are reused).
# Output: OVERLAP_DFS = {"<atlas>_<source>": table}, also saved as CSV under reports/.
from calmar.reporting import overlap_tables

OVERLAP_MASK_SOURCE = "all"   # "all", or one of "linda", "synthstroke", "openads", "manual"

# "all": selected T1 segmentations, optional OpenADS when chosen, and the manual mask.
from calmar.study import analysis_mask_sources, config_for_entry
_sources = (sorted({source for entry in SUBJECTS
                    for source in analysis_mask_sources(config_for_entry(CONFIG, entry))})
            if OVERLAP_MASK_SOURCE == "all" else [OVERLAP_MASK_SOURCE])
# overlap_tables reads ALL_ATLASES, _mask_path_for, CONFIG and PROJECT_DIR from the namespace.
OVERLAP_DFS = overlap_tables(globals(), SUBJECTS, _sources)
print("Atlas tables ready:", list(OVERLAP_DFS))

# Preview the atlas tables: one selector (participant, atlas, mask) drives the table,
# brain view and region chart. Group summaries are computed only when opened.
from calmar.atlas_results import build_results_panel

RESULTS_DEFAULT_ATLAS = CONFIG.get("DEFAULT_ATLAS")   # atlas shown first

_subject_report_panel = build_results_panel(
    OVERLAP_DFS, SUBJECTS, RESULTS_DEFAULT_ATLAS, ALL_ATLASES, cw,
    deriv_path_for=study.qc_dir, template_path_for=study.template_dir, inventory_for=study.masks,
    deriv_dir=DERIV_DIR, atlas_dir=ATLAS_DIR, reports_dir=REPORTS_DIR,
    threshold=CONFIG["LESION_THRESHOLD"],
    previous=globals().get("_subject_report_panel"))   # closes the panel from an earlier run
if _subject_report_panel is None:
    print("No atlas-overlap tables available. Run the atlas-overlap calculation first.")
else:
    display(_subject_report_panel["widget"])

# Add "Export selected atlas PDF" to the preview above. The PDF holds the participant,
# atlas and mask currently selected there: a summary page (lesion size, QC checkpoint
# status, top-20 bar chart of lesion_in_roi_percent) and the full region table.
from calmar.atlas_pdf import attach_pdf_export

ATLAS_PDF_DIR = REPORTS_DIR / "pdfs"   # where exported PDFs are written

attach_pdf_export(_subject_report_panel, ATLAS_PDF_DIR, atlas_dir=ATLAS_DIR, deriv_path_for=study.qc_dir)

# Settings for the participant interpretation report.
# The form below sets participant, mask, display atlas and decoding; edit the rest here.
from calmar.participant_report import rules_html

INTERP_SUBJECT_IDX = TEST_SUBJECT_IDX       # participant first shown in the form (index into SUBJECTS)
INTERP_MASK_SOURCE = "synthstroke"           # preferred lesion mask: "linda", "synthstroke" or "manual"
INTERP_DISPLAY_ATLAS = "HarvardOxford"       # atlas of the damaged-regions table
INTERP_RUN_DECODING = False                  # also run optional Neurosynth decoding (slower)
INTERP_KB_ATLASES = ["AAL", "HarvardOxford", "Destrieux", "Schaefer400", "JHU", "XTRACT"]  # atlases used to query the KB
INTERP_HIGH_CONF_ONLY = False                # keep only KB findings with confidence "high"
INTERP_EVIDENCE = "all"                      # KB evidence used (also in the form): "approved" entries only,
                                             # "in_review" (approved + in review) or "all" (+ unreviewed drafts)
INTERP_MIN_OVERLAP_PCT = 5.0                 # an ROI counts as damaged if it holds >= this % of the lesion
INTERP_MIN_REGION_PCT = 25.0                 # ... or if >= this % of the ROI is destroyed (small regions, thin tracts)
INTERP_DISCONNECTOME = "bcb"                 # tract disconnection from the "bcb" (BCBToolkit) or "deepdisco" map, or "none"
INTERP_MIN_DISCONNECTED_PCT = 50.0           # a tract counts as disconnected if >= this % of it has probability > 0.5
INTERP_TOP_N = 5                             # rows per predictor / therapy table (not the symptom relevance rule)
INTERP_MAX_ROIS = 15                         # rows in the damaged-regions table
INTERP_MAX_EVIDENCE_REGIONS = 8              # main damaged-ROI evidence blocks before the additional-region fold
INTERP_MIN_SUMMARY_CITATIONS = 2             # primary symptom summary: at least this many citation keys
INTERP_MIN_SUMMARY_ANATOMIES = 2             # ... and this many anatomy mappings after KB synonyms are merged
from calmar.kb_location import kb_selection_text, resolve_kb_root
APHASIA_KB_ROOT, APHASIA_KB_REASON = resolve_kb_root(PROJECT_DIR)
print(kb_selection_text(APHASIA_KB_ROOT, APHASIA_KB_REASON))

# The fixed clinical/display rules the report applies (sources in stroke-aphasia-kb/scripts/aphasia_rules.py).
display(HTML(rules_html(APHASIA_KB_ROOT)))

# Choose participant, mask, display atlas and decoding, then click Generate participant report.
# The form writes INTERP_SUBJECT_IDX / INTERP_MASK_SOURCE / INTERP_DISPLAY_ATLAS /
# INTERP_RUN_DECODING into the notebook namespace (show_report_inputs needs globals() for that).
from calmar.report_ui import show_report_inputs

_report_inputs = show_report_inputs(globals())
_rp_subj, _rp_mask = _report_inputs.subject, _report_inputs.source
raise StopExecution("Choose report inputs and click Generate participant report.")

# Generate the interpretation report for the participant and mask chosen in the form.
from calmar.participant_report import ReportSettings, generate_participant_report, selected_entry

REPORT_SETTINGS = ReportSettings(
    kb_atlases=tuple(INTERP_KB_ATLASES), display_atlas=INTERP_DISPLAY_ATLAS,
    high_conf_only=INTERP_HIGH_CONF_ONLY, evidence=INTERP_EVIDENCE,
    min_overlap_pct=INTERP_MIN_OVERLAP_PCT, min_region_pct=INTERP_MIN_REGION_PCT,
    disconnectome=INTERP_DISCONNECTOME, min_disconnected_pct=INTERP_MIN_DISCONNECTED_PCT, top_n=INTERP_TOP_N,
    max_rois=INTERP_MAX_ROIS, max_evidence_regions=INTERP_MAX_EVIDENCE_REGIONS,
    min_summary_citations=INTERP_MIN_SUMMARY_CITATIONS,
    min_summary_anatomies=INTERP_MIN_SUMMARY_ANATOMIES,
    threshold=CONFIG["LESION_THRESHOLD"])     # mask threshold used for Dice between masks

# namespace=globals() is needed only by the Slurm helpers that prepare overlap tables/overlay.
PARTICIPANT_REPORT = generate_participant_report(
    selected_entry(SUBJECTS, INTERP_SUBJECT_IDX), INTERP_MASK_SOURCE, REPORT_SETTINGS,
    study=study, namespace=globals(), reports_dir=REPORTS_DIR, dataset_dir=RAW_DATASET_DIR,
    atlas_dir=ATLAS_DIR, kb_root=APHASIA_KB_ROOT, kb_reason=APHASIA_KB_REASON,
    helpers=cw, run_decoding=INTERP_RUN_DECODING)
display(PARTICIPANT_REPORT.panel.widget)

# Short names used by the decoding and outputs cells.
_participant_report_panel, _rpt_path = PARTICIPANT_REPORT.panel, PARTICIPANT_REPORT.path
_calmar_generated_report_selection = PARTICIPANT_REPORT.selection
if globals().get("_report_inputs"):
    _report_inputs.message.value = ("Participant report ready. Optional Neurosynth decoding is running…"
                                    if INTERP_RUN_DECODING else "Participant report ready below.")

# Optional Neurosynth decoding of the report's mask (Slurm job; result joins the report above).
# Correlates the lesion with Neurosynth term maps; shows the top 10 terms with r > 0.05.
from calmar.participant_report import decode_into_report, selected_entry

if INTERP_RUN_DECODING:
    if not globals().get("_CALMAR_SETUP_READY", False):
        raise RuntimeError("Setup imports did not complete. Resolve the setup error before decoding.")
    _decode_entry = selected_entry(SUBJECTS, INTERP_SUBJECT_IDX)
    _decode_dataset = (study.dataset_for(_decode_entry)
                       if 'study' in globals() and hasattr(study, 'dataset_for') else RAW_DATASET_DIR)
    decode_into_report(
        globals().get("PARTICIPANT_REPORT"), _decode_entry, INTERP_MASK_SOURCE,
        kb_root=APHASIA_KB_ROOT, dataset_dir=_decode_dataset, project_dir=PROJECT_DIR,
        form=globals().get("_report_inputs"))

# Inventory of available outputs (presence only; does not imply QC was completed).
from calmar.participant_report import outputs_inventory_html

_latest_report = globals().get("PARTICIPANT_REPORT")   # None until a report is generated
display(HTML(outputs_inventory_html(SUBJECTS, study.masks, REPORTS_DIR,
                                    _latest_report.path if _latest_report else None)))
if callable(globals().get("show_timings")):
    show_timings()