src/data_prep/prep_utils.py

"""Prep utilities to standardize T1 + lesion masks into MNI and normalized form.

Defaults are self-contained under the project root:
- ANTs binaries are resolved from `<project>/tools/ants/(bin)/` first.
- TemplateFlow cache defaults to `<project>/data/templateflow`.

Environment overrides still work:
- `ANTS_REG` / `ANTS_APPLY`
- `TEMPLATEFLOW_HOME`
"""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
import os, shutil, subprocess, re, json, csv, hashlib, platform, stat, tempfile, time, zipfile
import urllib.request
import numpy as np
import nibabel as nib
import scipy.ndimage as ndi
from nibabel.processing import resample_from_to
from typing import Iterable
from collections import defaultdict

try:
    import resource
except Exception:  # pragma: no cover - resource is Unix-only
    resource = None


def _project_root() -> Path:
    # .../ARC_ATLAS_Train_v4/src/data_prep/prep_utils.py -> .../ARC_ATLAS_Train_v4
    return Path(__file__).resolve().parents[2]


def _first_existing(paths: list[Path]) -> Path | None:
    for p in paths:
        if p and p.exists():
            return p
    return None


def _ensure_templateflow_home() -> Path:
    tf_home = os.environ.get("TEMPLATEFLOW_HOME")
    if tf_home:
        return Path(tf_home).expanduser()
    local_tf_home = _project_root() / "data" / "templateflow"
    local_tf_home.mkdir(parents=True, exist_ok=True)
    os.environ["TEMPLATEFLOW_HOME"] = str(local_tf_home)
    return local_tf_home


def _ants_candidates(binary_name: str) -> list[Path]:
    root = _project_root()
    return [
        root / "tools" / "ants" / "bin" / binary_name,
        root / "tools" / "ants" / binary_name,
        root / "tools" / binary_name,
    ]


_ANTS_RELEASE_API = "https://api.github.com/repos/ANTsX/ANTs/releases/latest"
_ANTS_FALLBACK_URLS = [
    "https://github.com/ANTsX/ANTs/releases/download/v2.6.5/ants-2.6.5-ubuntu-22.04-X64-gcc.zip",
    "https://github.com/ANTsX/ANTs/releases/download/v2.6.5/ants-2.6.5-ubuntu20.04-X64-gcc.zip",
    "https://github.com/ANTsX/ANTs/releases/download/v2.6.5/ants-2.6.5-ubuntu18.04-X64-gcc.zip",
]
_TF_TEMPLATE = "tpl-MNI152NLin2009cAsym"


def _template_relpath(resolution: int) -> Path:
    res_tag = f"{int(resolution):02d}"
    return Path(_TF_TEMPLATE) / f"{_TF_TEMPLATE}_res-{res_tag}_desc-brain_T1w.nii.gz"


def _template_url(resolution: int) -> str:
    return f"https://templateflow.s3.amazonaws.com/{_template_relpath(resolution).as_posix()}"


def _ants_asset_preferences() -> list[str]:
    system = platform.system().lower()
    machine = platform.machine().lower()
    if system == "linux" and machine in {"x86_64", "amd64"}:
        return [
            "ubuntu-22.04-X64-gcc.zip",
            "ubuntu20.04-X64-gcc.zip",
            "ubuntu18.04-X64-gcc.zip",
            "ubuntu-24.04-X64-gcc.zip",
            "almalinux9-X64-gcc.zip",
            "almalinux8-X64-gcc.zip",
            "centos7-X64-gcc.zip",
        ]
    if system == "darwin" and machine in {"arm64", "aarch64"}:
        return ["macos-14-ARM64-clang.zip"]
    if system == "darwin" and machine in {"x86_64", "amd64"}:
        return ["macos-15-intel-X64-clang.zip"]
    return []


def _resolve_latest_ants_zip_url() -> str | None:
    req = urllib.request.Request(_ANTS_RELEASE_API, headers={"User-Agent": "chronic-stroke-segmentation"})
    with urllib.request.urlopen(req, timeout=30) as resp:
        payload = json.load(resp)
    assets = payload.get("assets", [])
    if not assets:
        return None

    prefs = _ants_asset_preferences()
    for suffix in prefs:
        for asset in assets:
            name = str(asset.get("name", ""))
            if name.endswith(suffix):
                url = asset.get("browser_download_url")
                if url:
                    return str(url)

    for asset in assets:
        name = str(asset.get("name", ""))
        if name.endswith(".zip"):
            url = asset.get("browser_download_url")
            if url:
                return str(url)
    return None


def _download_file(url: str, destination: Path):
    req = urllib.request.Request(url, headers={"User-Agent": "chronic-stroke-segmentation"})
    with urllib.request.urlopen(req, timeout=180) as resp, open(destination, "wb") as out_f:
        shutil.copyfileobj(resp, out_f)


def _install_ants_from_zip(zip_path: Path, install_root: Path):
    with tempfile.TemporaryDirectory(prefix="ants_extract_") as td:
        extract_root = Path(td)
        with zipfile.ZipFile(zip_path) as zf:
            zf.extractall(extract_root)

        candidates = []
        for p in extract_root.rglob("*"):
            if not p.is_dir():
                continue
            reg = p / "bin" / "antsRegistration"
            app = p / "bin" / "antsApplyTransforms"
            if reg.exists() and app.exists():
                candidates.append(p)
        if not candidates:
            raise RuntimeError("Downloaded archive did not contain ANTs binaries.")

        src_root = min(candidates, key=lambda p: len(str(p)))
        tmp_install = install_root.parent / f"{install_root.name}.tmp"
        if tmp_install.exists():
            shutil.rmtree(tmp_install, ignore_errors=True)
        shutil.copytree(src_root, tmp_install)
        if install_root.exists():
            shutil.rmtree(install_root, ignore_errors=True)
        tmp_install.rename(install_root)

    reg = install_root / "bin" / "antsRegistration"
    app = install_root / "bin" / "antsApplyTransforms"
    for binary in (reg, app):
        if binary.exists():
            binary.chmod(binary.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)


def _auto_install_ants() -> tuple[Path, Path]:
    install_root = _project_root() / "tools" / "ants"
    install_root.parent.mkdir(parents=True, exist_ok=True)

    urls: list[str] = []
    env_url = os.environ.get("ANTS_ZIP_URL")
    if env_url:
        urls.append(env_url)
    try:
        latest = _resolve_latest_ants_zip_url()
        if latest:
            urls.append(latest)
    except Exception as exc:
        print(f"[setup] unable to resolve ANTs latest release automatically: {exc}")
    urls.extend(_ANTS_FALLBACK_URLS)

    # preserve order but remove duplicates
    deduped: list[str] = []
    seen = set()
    for u in urls:
        if u and u not in seen:
            deduped.append(u)
            seen.add(u)

    last_exc: Exception | None = None
    with tempfile.TemporaryDirectory(prefix="ants_download_") as td:
        archive = Path(td) / "ants.zip"
        for url in deduped:
            print(f"[setup] downloading ANTs binaries from: {url}")
            try:
                _download_file(url, archive)
                _install_ants_from_zip(archive, install_root)
                reg = install_root / "bin" / "antsRegistration"
                app = install_root / "bin" / "antsApplyTransforms"
                if reg.exists() and app.exists():
                    print(f"[setup] ANTs installed under {install_root}")
                    return reg, app
            except Exception as exc:
                last_exc = exc
                print(f"[setup] failed from {url}: {exc}")
                continue

    if last_exc is not None:
        raise RuntimeError(f"Failed to install ANTs binaries automatically: {last_exc}")
    raise RuntimeError("Failed to install ANTs binaries automatically.")


def _ants_bins():
    reg_env = Path(os.environ["ANTS_REG"]).expanduser() if os.environ.get("ANTS_REG") else None
    app_env = Path(os.environ["ANTS_APPLY"]).expanduser() if os.environ.get("ANTS_APPLY") else None

    reg = reg_env or _first_existing(_ants_candidates("antsRegistration"))
    app = app_env or _first_existing(_ants_candidates("antsApplyTransforms"))

    if (reg is None or not reg.exists()) and shutil.which("antsRegistration"):
        reg = Path(shutil.which("antsRegistration"))
    if (app is None or not app.exists()) and shutil.which("antsApplyTransforms"):
        app = Path(shutil.which("antsApplyTransforms"))

    if reg is None or app is None or not reg.exists() or not app.exists():
        auto_install = os.environ.get("ANTS_AUTO_INSTALL", "1").lower() not in {"0", "false", "no"}
        if auto_install:
            print("[setup] ANTs binaries missing; attempting automatic install to tools/ants")
            try:
                reg, app = _auto_install_ants()
            except Exception as exc:
                print(f"[setup] automatic ANTs install failed: {exc}")

    if reg is None or app is None or not reg.exists() or not app.exists():
        exp_reg = _ants_candidates("antsRegistration")[0]
        exp_app = _ants_candidates("antsApplyTransforms")[0]
        raise FileNotFoundError(
            "ANTs binaries not found. Expected local binaries under "
            f"'{exp_reg.parent}' (e.g. {exp_reg.name}, {exp_app.name}), "
            "or set ANTS_REG / ANTS_APPLY. You can also call ensure_prep_runtime() "
            "or set ANTS_ZIP_URL to an ANTs release zip."
        )
    return reg, app


def ensure_prep_runtime(prefetch_template: bool = True) -> dict[str, str]:
    """Ensure ANTs binaries and TemplateFlow MNI template are available locally."""
    reg, app = _ants_bins()
    info = {
        "ants_registration": str(reg),
        "ants_apply_transforms": str(app),
    }
    if prefetch_template:
        tpl = _tpl_path(resolution=1)
        info["template_mni152_1mm_t1w"] = str(tpl)
    return info


def _command_timeout_seconds() -> int | None:
    """Timeout for ANTs subprocesses; <=0 disables timeout."""
    raw = os.environ.get("PREP_CMD_TIMEOUT_SEC", "").strip()
    if not raw:
        raw = os.environ.get("ANTS_CMD_TIMEOUT_SEC", "").strip()
    if not raw:
        return 900
    try:
        seconds = int(float(raw))
    except ValueError:
        print(f"[setup] invalid PREP_CMD_TIMEOUT_SEC={raw!r}; falling back to 900s")
        return 900
    return None if seconds <= 0 else seconds


def _affinity_cpu_count() -> int:
    try:
        return max(1, len(os.sched_getaffinity(0)))
    except Exception:
        return max(1, os.cpu_count() or 1)


def _cpu_quota_count() -> int | None:
    cpu_max = Path("/sys/fs/cgroup/cpu.max")
    if cpu_max.exists():
        try:
            quota_raw, period_raw = cpu_max.read_text().strip().split()[:2]
            if quota_raw != "max":
                quota = int(quota_raw)
                period = int(period_raw)
                if quota > 0 and period > 0:
                    return max(1, int(quota / period))
        except Exception:
            pass

    quota_path = Path("/sys/fs/cgroup/cpu/cpu.cfs_quota_us")
    period_path = Path("/sys/fs/cgroup/cpu/cpu.cfs_period_us")
    if quota_path.exists() and period_path.exists():
        try:
            quota = int(quota_path.read_text().strip())
            period = int(period_path.read_text().strip())
            if quota > 0 and period > 0:
                return max(1, int(quota / period))
        except Exception:
            pass
    return None


def _available_cpu_count() -> int:
    counts = [_affinity_cpu_count()]
    quota_cpus = _cpu_quota_count()
    if quota_cpus is not None:
        counts.append(quota_cpus)
    return max(1, min(counts))


def _ants_thread_count() -> int:
    raw = os.environ.get("PREP_ANTS_THREADS", "").strip()
    if raw:
        try:
            return max(1, int(float(raw)))
        except ValueError:
            print(f"[setup] invalid PREP_ANTS_THREADS={raw!r}; using automatic thread count")

    cap_raw = os.environ.get("PREP_ANTS_THREAD_CAP", "8").strip()
    try:
        cap = max(1, int(float(cap_raw)))
    except ValueError:
        print(f"[setup] invalid PREP_ANTS_THREAD_CAP={cap_raw!r}; using 8")
        cap = 8
    return max(1, min(_available_cpu_count(), cap))


def _subprocess_env() -> dict[str, str]:
    env = os.environ.copy()
    threads = str(_ants_thread_count())
    env["ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS"] = threads
    env["OMP_NUM_THREADS"] = threads
    env["OMP_THREAD_LIMIT"] = threads
    return env


def _child_cpu_snapshot() -> tuple[float, float] | None:
    if resource is None:
        return None
    usage = resource.getrusage(resource.RUSAGE_CHILDREN)
    return float(usage.ru_utime), float(usage.ru_stime)


def _child_cpu_summary(before: tuple[float, float] | None, elapsed: float) -> str:
    after = _child_cpu_snapshot()
    if before is None or after is None:
        return "child_cpu=unknown"
    user = max(0.0, after[0] - before[0])
    system = max(0.0, after[1] - before[1])
    total = user + system
    effective = total / elapsed if elapsed > 0 else 0.0
    return f"child_cpu={total:.1f}s user={user:.1f}s sys={system:.1f}s effective_cores={effective:.2f}"


def _run(cmd: list[str]):
    cmd_str = " ".join(map(str, cmd))
    timeout_sec = _command_timeout_seconds()
    timeout_label = f"{timeout_sec}s" if timeout_sec is not None else "disabled"
    env = _subprocess_env()
    threads = env["ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS"]
    print(">>", cmd_str)
    print(f"[cmd] timeout={timeout_label}")
    print(
        f"[cmd] ants_threads={threads} "
        f"available_cpus={_available_cpu_count()} "
        f"affinity_cpus={_affinity_cpu_count()} "
        f"cpu_quota_cpus={_cpu_quota_count() or 'unknown'} "
        f"thread_cap={os.environ.get('PREP_ANTS_THREAD_CAP', '8')}"
    )
    t0 = time.monotonic()
    child_cpu_before = _child_cpu_snapshot()
    try:
        res = subprocess.run(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            timeout=timeout_sec,
            env=env,
        )
    except subprocess.TimeoutExpired as exc:
        elapsed = time.monotonic() - t0
        partial = exc.stdout or ""
        if isinstance(partial, bytes):
            partial = partial.decode(errors="replace")
        tail = "\n".join(partial.splitlines()[-120:])
        cpu_summary = _child_cpu_summary(child_cpu_before, elapsed)
        raise RuntimeError(
            f"Command timed out after {elapsed:.1f}s (limit={timeout_sec}s): {cmd_str}\n"
            f"[cmd] {cpu_summary}\n"
            f"---- output tail ----\n{tail}"
        ) from exc

    elapsed = time.monotonic() - t0
    print(f"[cmd] completed in {elapsed:.1f}s rc={res.returncode} {_child_cpu_summary(child_cpu_before, elapsed)}")
    if res.returncode != 0:
        output = res.stdout or ""
        tail = "\n".join(output.splitlines()[-200:])
        raise RuntimeError(
            f"Command failed with exit code {res.returncode}: {cmd_str}\n"
            f"---- output tail ----\n{tail}"
        )
    return res.stdout


def _key_from_name(name: str) -> str | None:
    sub = re.search(r"(sub-[^_]+)", name)
    ses = re.search(r"(ses-[^_]+)", name)
    parts = [m.group(1) for m in (sub, ses) if m]
    return "_".join(parts) if parts else None


def _slug_from_raw(raw_root: Path, name: str) -> str:
    stem = raw_root.name
    safe = re.sub(r"[^a-zA-Z0-9]+", "-", f"{name}-{stem}").strip("-")
    digest = hashlib.md5(str(raw_root.resolve()).encode()).hexdigest()[:8]
    return f"{safe}-{digest}"


def _tpl_path(resolution: int = 1) -> Path:
    """Return MNI T1 path for given resolution, using local cache with robust fallback."""
    tf_home = _ensure_templateflow_home()
    local_tpl = tf_home / _template_relpath(resolution)
    if local_tpl.exists():
        return local_tpl

    # Prefer TemplateFlow client when available.
    try:
        from templateflow.api import get as tf_get
    except Exception as exc:
        print(f"[setup] templateflow package unavailable ({exc}); using direct template download.")
    else:
        try:
            tpl = tf_get(
                "MNI152NLin2009cAsym",
                resolution=resolution,
                suffix="T1w",
                desc="brain",
                extension="nii.gz",
            )
            tpl_path = Path(tpl[0]) if isinstance(tpl, (list, tuple)) else Path(tpl)
            if tpl_path.exists():
                return tpl_path
        except Exception as exc:
            print(f"[setup] templateflow API lookup failed ({exc}); using direct template download.")

    # Direct download fallback avoids optional templateflow dependency issues.
    local_tpl.parent.mkdir(parents=True, exist_ok=True)
    _download_file(_template_url(resolution), local_tpl)
    if not local_tpl.exists():
        raise FileNotFoundError(f"Template download failed: expected {local_tpl}")
    return local_tpl

def _find_existing_output(out_root: Path, raw_root: Path) -> Path | None:
    """Look for an existing prep folder whose marker matches raw_root."""
    raw_root = raw_root.resolve()
    for m in out_root.glob("*/source.json"):
        try:
            info = json.loads(m.read_text())
            if Path(info.get("raw_root", "")).resolve() == raw_root:
                return m.parent
        except Exception:
            continue
    return None


def list_pairs(raw_root: Path, t1_glob: str, mask_glob: str):
    t1s = list(raw_root.glob(t1_glob))
    masks = list(raw_root.glob(mask_glob))
    return _match_pairs(t1s, masks)


def _norm_key_from_name(name: str) -> str:
    base = name
    if base.endswith(".nii.gz"):
        base = base[:-7]
    elif base.endswith(".nii"):
        base = base[:-4]
    drop = [
        "_T1w_MNI_norm", "_T1w_MNI", "_T1w_brain", "_T1w", "_T1",
        "_lesion_mask_MNI_clean", "_lesion_mask_MNI", "_lesion_mask",
        "_desc-lesion_mask", "_mask", "mask"
    ]
    for sfx in drop:
        if base.endswith(sfx):
            base = base[: -len(sfx)]
    return base


def _match_pairs(t1s: list[Path], masks: list[Path]):
    """Match T1 and mask by the most specific key available.
    Priority: exact normalized basename -> unique; else sub/ses key unique. Ambiguous cases are skipped.
    """
    mask_by_base = defaultdict(list)
    mask_by_subses = defaultdict(list)
    for m in masks:
        base = _norm_key_from_name(m.name)
        mask_by_base[base].append(m)
        key = _key_from_name(m.name) or base
        mask_by_subses[key].append(m)

    pairs = []
    for t1 in t1s:
        base = _norm_key_from_name(t1.name)
        key = _key_from_name(t1.name) or base
        chosen = None
        if mask_by_base.get(base):
            if len(mask_by_base[base]) == 1:
                chosen = mask_by_base[base][0]
            else:
                print(f"[warn] multiple masks share base {base}; skipping")
                continue
        elif mask_by_subses.get(key):
            if len(mask_by_subses[key]) == 1:
                chosen = mask_by_subses[key][0]
            else:
                print(f"[warn] ambiguous masks for {t1.name} (key {key}): {len(mask_by_subses[key])}; skipping")
                continue
        if not chosen:
            print(f"[warn] no mask for {t1.name} (key {key}); skipping")
            continue
        pairs.append((t1, chosen, base))
    return pairs


def resample_mask_to_t1(mask: Path, t1: Path, out_path: Path):
    mi = nib.load(str(mask))
    ti = nib.load(str(t1))
    if mi.shape[:3] != ti.shape[:3] or not np.allclose(mi.affine, ti.affine, atol=1e-4):
        rs = resample_from_to(mi, (ti.shape, ti.affine), order=0)
        data = (rs.get_fdata() > 0.5).astype(np.uint8)
    else:
        data = (mi.get_fdata() > 0.5).astype(np.uint8)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    _save_nifti_like(out_path, data, ti, np.uint8)


def _save_nifti_like(path: Path, data: np.ndarray, ref_img: nib.Nifti1Image, dtype) -> Path:
    path.parent.mkdir(parents=True, exist_ok=True)
    header = ref_img.header.copy()
    header.set_data_dtype(dtype)
    nib.save(nib.Nifti1Image(np.asarray(data, dtype=dtype), ref_img.affine, header), str(path))
    return path


def _ants_registration_profile() -> dict[str, str]:
    """Return ANTs registration schedule based on PREP_ANTS_PROFILE."""
    profile = os.environ.get("PREP_ANTS_PROFILE", "balanced").strip().lower()
    presets: dict[str, dict[str, str]] = {
        "quick": {
            "metric_sampling": "Random,0.1",
            "rigid_affine_convergence": "120x60x20",
            "rigid_affine_smoothing": "3x2x1vox",
            "rigid_affine_shrink": "8x4x2",
            "syn_convergence": "8x4x2",
            "syn_smoothing": "2x1x0vox",
            "syn_shrink": "4x2x1",
            "syn_transform": "SyN[0.05,3,0]",
        },
        "fast": {
            "metric_sampling": "Random,0.2",
            "rigid_affine_convergence": "300x120x40",
            "rigid_affine_smoothing": "3x2x1vox",
            "rigid_affine_shrink": "6x4x2",
            "syn_convergence": "24x12x6",
            "syn_smoothing": "2x1x0vox",
            "syn_shrink": "4x2x1",
            "syn_transform": "SyN[0.08,3,0]",
        },
        "balanced": {
            "metric_sampling": "Regular,0.2",
            "rigid_affine_convergence": "600x250x100",
            "rigid_affine_smoothing": "3x2x1vox",
            "rigid_affine_shrink": "4x2x1",
            "syn_convergence": "40x20x10",
            "syn_smoothing": "2x1x0vox",
            "syn_shrink": "4x2x1",
            "syn_transform": "SyN[0.1,3,0]",
        },
        "accurate": {
            "metric_sampling": "Regular,0.25",
            "rigid_affine_convergence": "1000x500x250",
            "rigid_affine_smoothing": "3x2x1vox",
            "rigid_affine_shrink": "4x2x1",
            "syn_convergence": "60x40x20",
            "syn_smoothing": "2x1x0vox",
            "syn_shrink": "4x2x1",
            "syn_transform": "SyN[0.1,3,0]",
        },
    }
    if profile not in presets:
        print(f"[setup] unknown PREP_ANTS_PROFILE={profile!r}; using 'balanced'")
        profile = "balanced"
    cfg = dict(presets[profile])
    cfg["profile"] = profile
    return cfg


def ants_register(t1: Path, prefix: Path, tpl: Path, reg_bin: Path, use_2mm: bool = True):
    tpl_reg = tpl
    if use_2mm:
        try:
            from templateflow.api import get as tf_get
            tpl2 = tf_get("MNI152NLin2009cAsym", resolution=2, suffix="T1w", desc="brain", extension="nii.gz")
            tpl_reg = Path(tpl2[0]) if isinstance(tpl2, (list, tuple)) else Path(tpl2)
        except Exception:
            tpl_reg = tpl
    cfg = _ants_registration_profile()
    print(
        "[ants] profile="
        f"{cfg['profile']} rigid_affine={cfg['rigid_affine_convergence']} "
        f"syn={cfg['syn_convergence']}"
    )
    _run([
        str(reg_bin), '-d','3',
        '-r', f'[{tpl_reg},{t1},1]',
        '-m', f"Mattes[{tpl_reg},{t1},1,32,{cfg['metric_sampling']}]",
        '-t','Rigid[0.1]',
        '-c', cfg['rigid_affine_convergence'],
        '-s', cfg['rigid_affine_smoothing'],
        '-f', cfg['rigid_affine_shrink'],
        '-m', f"Mattes[{tpl_reg},{t1},1,32,{cfg['metric_sampling']}]",
        '-t','Affine[0.1]',
        '-c', cfg['rigid_affine_convergence'],
        '-s', cfg['rigid_affine_smoothing'],
        '-f', cfg['rigid_affine_shrink'],
        '-m', f'CC[{tpl_reg},{t1},1,4]',
        '-t', cfg['syn_transform'],
        '-c', cfg['syn_convergence'],
        '-s', cfg['syn_smoothing'],
        '-f', cfg['syn_shrink'],
        '-o', f'[{prefix},{prefix}warped.nii.gz,{prefix}invwarped.nii.gz]'
    ])


def ants_apply(img_in: Path, ref: Path, xfm_prefix: Path, out_path: Path, apply_bin: Path, nn: bool=False):
    args = [str(apply_bin), '-d','3', '-i', str(img_in), '-r', str(ref), '-o', str(out_path)]
    if nn:
        args += ['-n','NearestNeighbor']
    args += ['-t', str(xfm_prefix)+'1Warp.nii.gz', '-t', str(xfm_prefix)+'0GenericAffine.mat']
    _run(args)


def normalize_t1(vol: np.ndarray) -> np.ndarray:
    nz = vol[vol>0]
    if nz.size == 0:
        return np.zeros_like(vol, np.float32)
    p1,p99 = np.percentile(nz,[1,99])
    vol = np.clip(vol, p1, p99)
    mu, sd = nz.mean(), nz.std()
    vol = (vol - mu)/(sd+1e-8)
    mn, mx = vol.min(), vol.max()
    return ((vol - mn)/(mx - mn + 1e-8)).astype(np.float32)


def largest_component(mask: np.ndarray) -> np.ndarray:
    labeled, nlab = ndi.label(mask)
    if nlab <= 1:
        return mask.astype(np.uint8)
    sizes = np.bincount(labeled.ravel())
    keep = sizes[1:].argmax() + 1
    return (labeled == keep).astype(np.uint8)


def _bbox(mask: np.ndarray):
    coords = np.argwhere(mask > 0)
    if coords.size == 0:
        return None
    mins = coords.min(axis=0)
    maxs = coords.max(axis=0)
    return mins, maxs


def _overlap_score(mask: np.ndarray, brain: np.ndarray) -> float:
    inter = np.logical_and(mask > 0, brain > 0).sum()
    return inter / max((mask > 0).sum(), 1)


def _try_flips(mask: np.ndarray, brain: np.ndarray):
    base_score = _overlap_score(mask, brain)
    best = (mask, base_score, "none")
    flips = [
        (np.flip(mask, axis=0), "flip_x"),
        (np.flip(mask, axis=1), "flip_y"),
        (np.flip(mask, axis=2), "flip_z"),
    ]
    for flipped, tag in flips:
        score = _overlap_score(flipped, brain)
        if score > best[1]:
            best = (flipped, score, tag)
    return best


@dataclass
class DatasetConfig:
    name: str
    raw_root: Path | None = None
    images_dir: Path | None = None
    masks_dir: Path | None = None
    t1_glob: str = "**/*_T1w.nii.gz"
    mask_glob: str = "**/*mask*.nii.gz"
    overwrite: bool = False
    already_mni: bool = False


def run_prep(datasets: Iterable[DatasetConfig], out_root: Path, force_overwrite: bool = False):
    tpl = None
    reg_bin = apply_bin = None
    out_root.mkdir(parents=True, exist_ok=True)
    qc_rows = []
    outputs = []

    for ds in datasets:
        raw = Path(ds.raw_root).expanduser() if ds.raw_root else None
        img_root = Path(ds.images_dir).expanduser() if ds.images_dir else raw
        msk_root = Path(ds.masks_dir).expanduser() if ds.masks_dir else raw
        needs_ants = not ds.already_mni

        print(f"[{ds.name}] image root: {img_root} | mask root: {msk_root}")
        print(f"[{ds.name}] globs: t1={ds.t1_glob} masks={ds.mask_glob}")

        if not img_root or not img_root.exists():
            print(f"[skip] {ds.name}: images root missing {img_root}")
            continue
        if not msk_root or not msk_root.exists():
            print(f"[skip] {ds.name}: masks root missing {msk_root}")
            continue

        t1s = list(img_root.glob(ds.t1_glob))
        mks = list(msk_root.glob(ds.mask_glob))
        pairs = _match_pairs(t1s, mks)
        print(f"[{ds.name}] images: {len(t1s)} masks: {len(mks)} pairs found: {len(pairs)}")
        slug_base = raw or img_root
        slug = _slug_from_raw(slug_base, ds.name)
        overwrite = force_overwrite or ds.overwrite
        out_ds_existing = None if overwrite else _find_existing_output(out_root, slug_base)
        out_ds = out_ds_existing or (out_root / slug)
        marker = out_ds / "source.json"
        if overwrite and out_ds.exists():
            shutil.rmtree(out_ds, ignore_errors=True)
            print(f"[{ds.name}] overwrite=True -> cleared {out_ds}")
        elif marker.exists() and not overwrite:
            print(f"[{ds.name}] already processed -> {out_ds}, skipping (overwrite=True to redo).")
            outputs.append(out_ds)
            continue
        if not pairs:
            print(f"[warn] {ds.name}: no matched pairs; skipping dataset.")
            continue
        out_nat = out_ds / 'native_resampled_masks'
        out_mni = out_ds / 'mni_1mm_ants_fixed'
        out_norm = out_mni / 't1_norm'
        out_clean = out_mni / 'masks_clean'
        out_xfm = out_ds / 'xfm'
        for d in (out_nat, out_mni, out_norm, out_clean, out_xfm):
            d.mkdir(parents=True, exist_ok=True)

        for t1, mask, key in pairs:
            mask_t1 = out_nat / f"{key}_lesion_mask_T1w_native.nii.gz"
            prefix = out_xfm / f"{key}_t1_to_mni_"
            t1_mni = out_mni / f"{key}_T1w_MNI.nii.gz"
            mask_mni = out_mni / f"{key}_lesion_mask_MNI.nii.gz"
            t1_norm = out_norm / f"{key}_T1w_MNI_norm.nii.gz"
            mask_clean = out_clean / f"{key}_lesion_mask_MNI_clean.nii.gz"

            if ds.already_mni:
                if not mask_t1.exists():
                    resample_mask_to_t1(mask, t1, mask_t1)
                if not t1_mni.exists():
                    shutil.copy2(t1, t1_mni)
                if not mask_mni.exists():
                    shutil.copy2(mask_t1, mask_mni)
            else:
                if tpl is None:
                    tpl = _tpl_path(resolution=1)
                if reg_bin is None or apply_bin is None:
                    reg_bin, apply_bin = _ants_bins()
                if not mask_t1.exists():
                    resample_mask_to_t1(mask, t1, mask_t1)
                if not (prefix.with_name(prefix.name+'0GenericAffine.mat')).exists():
                    ants_register(t1, prefix, tpl, reg_bin, use_2mm=True)
                if not t1_mni.exists():
                    ants_apply(t1, tpl, prefix, t1_mni, apply_bin, nn=False)
                if not mask_mni.exists():
                    ants_apply(mask_t1, tpl, prefix, mask_mni, apply_bin, nn=True)
                    mi = nib.load(str(mask_mni)); data=(mi.get_fdata()>0.5).astype(np.uint8)
                    _save_nifti_like(mask_mni, data, mi, np.uint8)
            if not t1_norm.exists():
                t1_mni_img = nib.load(str(t1_mni))
                vol = t1_mni_img.get_fdata().astype(np.float32)
                norm = normalize_t1(vol)
                _save_nifti_like(t1_norm, norm, t1_mni_img, np.float32)
            if not mask_clean.exists():
                data = (nib.load(str(mask_mni)).get_fdata()>0.5).astype(np.uint8)
                data = largest_component(data)
                # alignment sanity check: ensure mask overlaps brain; try flips if not
                brain = (nib.load(str(t1_mni)).get_fdata()>0).astype(np.uint8)
                best_mask, best_score, tag = _try_flips(data, brain)
                if best_score < 0.1:
                    print(f"[warn] low overlap for {mask_mni.name} (score {best_score:.3f}); keeping original")
                elif tag != "none":
                    print(f"[info] flipped {tag} for {mask_mni.name} (overlap {best_score:.3f})")
                    data = best_mask
                mask_mni_img = nib.load(str(mask_mni))
                _save_nifti_like(mask_clean, data, mask_mni_img, np.uint8)

            ti = nib.load(str(t1_mni)); mi = nib.load(str(mask_clean))
            qc_rows.append(dict(
                dataset=ds.name,
                slug=slug,
                key=key,
                t1_mni=str(t1_mni),
                mask_mni=str(mask_clean),
                t1_shape=str(ti.shape[:3]),
                t1_zooms=str(tuple(round(z,3) for z in ti.header.get_zooms()[:3])),
                mask_nonzero=int(np.count_nonzero(mi.get_fdata()>0)),
            ))
        marker.parent.mkdir(parents=True, exist_ok=True)
        marker.write_text(json.dumps({"raw_root": str(slug_base), "slug": slug}, indent=2))
        outputs.append(out_ds)

    if qc_rows:
        qc_csv = out_root / 'prep_qc.csv'
        with open(qc_csv, 'w', newline='') as f:
            writer = csv.DictWriter(f, fieldnames=qc_rows[0].keys())
            writer.writeheader(); writer.writerows(qc_rows)
        print('QC written', qc_csv)
    else:
        print('No QC rows written (no datasets processed).')
    return outputs


def combine_standardized(dataset_roots: Iterable[Path], dest: Path):
    dest_t1 = dest / "t1"
    dest_mk = dest / "masks"
    dest_t1.mkdir(parents=True, exist_ok=True)
    dest_mk.mkdir(parents=True, exist_ok=True)
    manifest = []
    for ds_root in dataset_roots:
        slug = ds_root.name
        t1_dir = ds_root / "mni_1mm_ants_fixed" / "t1_norm"
        msk_dir = ds_root / "mni_1mm_ants_fixed" / "masks_clean"
        if not t1_dir.exists() or not msk_dir.exists():
            print(f"[combine] missing t1_norm or masks_clean in {ds_root}, skipping")
            continue
        for t1 in sorted(t1_dir.glob("*.nii.gz")):
            base = t1.name
            mask = msk_dir / base.replace("_T1w_MNI_norm", "_lesion_mask_MNI_clean")
            if not mask.exists():
                continue
            out_t1 = dest_t1 / f"{slug}__{base}"
            out_mk = dest_mk / f"{slug}__{mask.name}"
            shutil.copy2(t1, out_t1)
            shutil.copy2(mask, out_mk)
            manifest.append({"slug": slug, "key": base, "t1": str(out_t1.resolve()), "mask": str(out_mk.resolve())})
    if manifest:
        mf = dest / "manifest.csv"
        with open(mf, "w", newline="") as f:
            writer = csv.DictWriter(f, fieldnames=manifest[0].keys())
            writer.writeheader(); writer.writerows(manifest)
        print("Combined manifest ->", mf)
    else:
        print("No combined manifest written (no pairs).")


def run_prep_images_only(
    images_dir: Path,
    out_root: Path,
    name: str = "TEST",
    t1_glob: str = "**/*.nii.gz",
    already_mni: bool = False,
    overwrite: bool = False,
) -> Path:
    """Prep images to MNI/normalized space without requiring lesion masks."""
    images_dir = Path(images_dir).expanduser()
    out_root = Path(out_root).expanduser()
    if not images_dir.exists():
        raise FileNotFoundError(f"Images root missing: {images_dir}")
    out_root.mkdir(parents=True, exist_ok=True)

    t1s = sorted(images_dir.glob(t1_glob))
    if not t1s:
        raise RuntimeError(f"No T1 images found in {images_dir} using glob {t1_glob}")

    slug = _slug_from_raw(images_dir, name)
    out_ds = out_root / slug
    marker = out_ds / "source.json"
    if overwrite and out_ds.exists():
        shutil.rmtree(out_ds, ignore_errors=True)
    elif marker.exists() and not overwrite:
        print(f"[{name}] already processed -> {out_ds}, skipping (overwrite=True to redo).")
        return out_ds

    out_mni = out_ds / "mni_1mm_ants_fixed"
    out_norm = out_mni / "t1_norm"
    out_xfm = out_ds / "xfm"
    for d in (out_mni, out_norm, out_xfm):
        d.mkdir(parents=True, exist_ok=True)

    tpl = None
    reg_bin = apply_bin = None
    seen_keys = set()
    qc_rows = []

    def _unique_key(base: str) -> str:
        key = base or "case"
        if key not in seen_keys:
            seen_keys.add(key)
            return key
        i = 2
        while f"{key}_{i}" in seen_keys:
            i += 1
        new_key = f"{key}_{i}"
        seen_keys.add(new_key)
        return new_key

    for t1 in t1s:
        key = _unique_key(_norm_key_from_name(t1.name))
        prefix = out_xfm / f"{key}_t1_to_mni_"
        t1_mni = out_mni / f"{key}_T1w_MNI.nii.gz"
        t1_norm = out_norm / f"{key}_T1w_MNI_norm.nii.gz"

        if already_mni:
            if not t1_mni.exists():
                shutil.copy2(t1, t1_mni)
        else:
            if tpl is None:
                tpl = _tpl_path(resolution=1)
            if reg_bin is None or apply_bin is None:
                reg_bin, apply_bin = _ants_bins()
            if not (prefix.with_name(prefix.name + "0GenericAffine.mat")).exists():
                ants_register(t1, prefix, tpl, reg_bin, use_2mm=True)
            if not t1_mni.exists():
                ants_apply(t1, tpl, prefix, t1_mni, apply_bin, nn=False)

        if not t1_norm.exists():
            img_mni = nib.load(str(t1_mni))
            norm = normalize_t1(img_mni.get_fdata().astype(np.float32))
            _save_nifti_like(t1_norm, norm, img_mni, np.float32)

        ti = nib.load(str(t1_mni))
        qc_rows.append(
            dict(
                dataset=name,
                slug=slug,
                key=key,
                t1_mni=str(t1_mni),
                t1_shape=str(ti.shape[:3]),
                t1_zooms=str(tuple(round(z, 3) for z in ti.header.get_zooms()[:3])),
            )
        )

    marker.parent.mkdir(parents=True, exist_ok=True)
    marker.write_text(json.dumps({"raw_root": str(images_dir), "slug": slug}, indent=2))

    if qc_rows:
        qc_csv = out_root / "prep_qc_images_only.csv"
        with open(qc_csv, "w", newline="") as f:
            writer = csv.DictWriter(f, fieldnames=qc_rows[0].keys())
            writer.writeheader()
            writer.writerows(qc_rows)
        print("QC written", qc_csv)
    return out_ds


def combine_standardized_images_only(dataset_roots: Iterable[Path], dest: Path):
    """Combine standardized image-only outputs into a single test_input root."""
    dest_t1 = dest / "t1"
    dest_t1.mkdir(parents=True, exist_ok=True)
    manifest = []
    for ds_root in dataset_roots:
        slug = ds_root.name
        t1_dir = ds_root / "mni_1mm_ants_fixed" / "t1_norm"
        if not t1_dir.exists():
            print(f"[combine] missing t1_norm in {ds_root}, skipping")
            continue
        for t1 in sorted(t1_dir.glob("*.nii.gz")):
            out_t1 = dest_t1 / f"{slug}__{t1.name}"
            shutil.copy2(t1, out_t1)
            manifest.append({"slug": slug, "key": t1.name, "t1": str(out_t1.resolve())})
    if manifest:
        mf = dest / "manifest.csv"
        with open(mf, "w", newline="") as f:
            writer = csv.DictWriter(f, fieldnames=manifest[0].keys())
            writer.writeheader()
            writer.writerows(manifest)
        print("Combined image-only manifest ->", mf)
    else:
        print("No image-only manifest written (no images).")


__all__ = [
    'DatasetConfig',
    'ensure_prep_runtime',
    'run_prep',
    'combine_standardized',
    'run_prep_images_only',
    'combine_standardized_images_only',
]