src/inference/runner.py

from __future__ import annotations

import os
import time
from datetime import datetime
from pathlib import Path

import nibabel as nib
import numpy as np
import scipy.ndimage as ndi
from nibabel.processing import resample_from_to

from .io_manifest import case_key_from_name, write_manifest
from .model_loader import LoadedModelBundle, load_model_bundle
from .volume_preprocess import prepare_image_for_model, restore_prediction_to_original

_MNI_BRAIN_TEMPLATE_REL = (
    Path("tpl-MNI152NLin2009cAsym") / "tpl-MNI152NLin2009cAsym_res-01_desc-brain_T1w.nii.gz"
)
_BRAIN_INTERIOR_EROSION_ITERS = 1
_EDGE_COMPONENT_MIN_INTERIOR_VOXELS = 16
_EDGE_COMPONENT_MIN_INTERIOR_FRACTION = 0.05
_EDGE_COMPONENT_LARGE_SIZE_VOXELS = 1024
_EDGE_SHELL_SMALL_COMPONENT_VOXELS = 32
_EDGE_SHELL_MIN_CONTACT_VOXELS = 8
_EDGE_SHELL_MIN_CONTACT_FRACTION = 0.5


def _templateflow_home() -> Path:
    tf_home = os.environ.get("TEMPLATEFLOW_HOME")
    if tf_home:
        return Path(tf_home).expanduser()
    return Path(__file__).resolve().parents[2] / "data" / "templateflow"


def _load_mni_brain_mask(ref_img: nib.Nifti1Image) -> np.ndarray | None:
    tpl_path = _templateflow_home() / _MNI_BRAIN_TEMPLATE_REL
    if not tpl_path.exists():
        return None

    tpl_img = nib.load(str(tpl_path))
    if tpl_img.shape[:3] != ref_img.shape[:3] or not np.allclose(tpl_img.affine, ref_img.affine, atol=1e-4):
        tpl_img = resample_from_to(tpl_img, ref_img, order=0)

    return (tpl_img.get_fdata() > 0).astype(np.uint8)


def _filter_edge_components(mask: np.ndarray, brain_mask: np.ndarray) -> tuple[np.ndarray, int]:
    mask_bool = np.asarray(mask, dtype=bool)
    brain_bool = np.asarray(brain_mask, dtype=bool)
    if not np.any(mask_bool) or not np.any(brain_bool):
        return mask_bool.astype(np.uint8), 0

    brain_interior = ndi.binary_erosion(
        brain_bool,
        iterations=_BRAIN_INTERIOR_EROSION_ITERS,
        border_value=0,
    )
    if not np.any(brain_interior):
        return mask_bool.astype(np.uint8), 0
    brain_shell = np.logical_and(brain_bool, np.logical_not(brain_interior))

    labeled, nlab = ndi.label(mask_bool)
    if nlab <= 0:
        return mask_bool.astype(np.uint8), 0

    keep = np.zeros_like(mask_bool, dtype=bool)
    removed = 0
    for label_idx in range(1, nlab + 1):
        component = labeled == label_idx
        component_size = int(np.count_nonzero(component))
        if component_size == 0:
            continue

        interior_overlap = int(np.count_nonzero(component & brain_interior))
        interior_fraction = float(interior_overlap) / float(component_size)
        has_strong_interior_support = (
            interior_overlap >= _EDGE_COMPONENT_MIN_INTERIOR_VOXELS
            or interior_fraction >= _EDGE_COMPONENT_MIN_INTERIOR_FRACTION
            or (component_size >= _EDGE_COMPONENT_LARGE_SIZE_VOXELS and interior_overlap > 0)
        )

        if has_strong_interior_support:
            interior_component = np.logical_and(component, brain_interior)
            keep |= interior_component

            shell_component = np.logical_and(component, brain_shell)
            if np.any(shell_component):
                shell_labels, shell_n = ndi.label(shell_component)
                interior_touch_zone = ndi.binary_dilation(interior_component, iterations=1, border_value=0)
                for shell_idx in range(1, shell_n + 1):
                    shell_piece = shell_labels == shell_idx
                    shell_size = int(np.count_nonzero(shell_piece))
                    if shell_size == 0:
                        continue

                    contact_voxels = int(np.count_nonzero(shell_piece & interior_touch_zone))
                    contact_fraction = float(contact_voxels) / float(shell_size)
                    has_strong_contact = (
                        contact_fraction >= _EDGE_SHELL_MIN_CONTACT_FRACTION
                        or (
                            shell_size <= _EDGE_SHELL_SMALL_COMPONENT_VOXELS
                            and contact_voxels >= _EDGE_SHELL_MIN_CONTACT_VOXELS
                        )
                    )

                    if has_strong_contact:
                        keep |= shell_piece
                    else:
                        removed += 1
        else:
            removed += 1

    return keep.astype(np.uint8), removed


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


def _copy_input_t1(path: Path, t1_path: Path) -> Path:
    path.parent.mkdir(parents=True, exist_ok=True)
    img = nib.load(str(t1_path))
    data = np.asanyarray(img.dataobj)
    header = img.header.copy()
    if np.issubdtype(data.dtype, np.floating):
        data = np.asarray(data, dtype=np.float32)
        header.set_data_dtype(np.float32)
    else:
        data = np.asarray(data)
        header.set_data_dtype(data.dtype)
    nib.save(nib.Nifti1Image(data, img.affine, header), str(path))
    return path


def _predict_one(
    t1_path: Path,
    bundle: LoadedModelBundle,
    threshold: float,
    input_dir: Path,
    hard_dir: Path,
) -> dict:
    t0 = time.monotonic()

    model = bundle.model
    target_shape = tuple(int(x) for x in bundle.input_shape[:3])

    prepped, ctx = prepare_image_for_model(t1_path, target_shape=target_shape)

    x = np.zeros((1, *bundle.input_shape), dtype=np.float32)
    x[0, ..., 0] = prepped

    pred_target = model.predict(x, verbose=0)[0, ..., 0].astype(np.float32)
    pred_hard_target = (pred_target >= float(threshold)).astype(np.uint8)

    pred_hard_orig = restore_prediction_to_original(pred_hard_target.astype(np.float32), ctx)
    pred_hard_orig = (pred_hard_orig >= 0.5).astype(np.uint8)

    brain_mask = _load_mni_brain_mask(ctx.original_img)
    removed_components = 0
    if brain_mask is not None:
        pred_hard_orig = (pred_hard_orig * brain_mask.astype(np.uint8)).astype(np.uint8)
        pred_hard_orig, removed_components = _filter_edge_components(pred_hard_orig, brain_mask)

    case_key = case_key_from_name(t1_path.name)
    input_name = t1_path.name
    hard_name = f"{case_key}_lesion_pred_hard_th{int(round(threshold * 100)):03d}.nii.gz"

    input_path = _copy_input_t1(input_dir / input_name, t1_path)
    hard_path = _save_hard_prediction(hard_dir / hard_name, pred_hard_orig, ctx.original_img)

    elapsed = time.monotonic() - t0
    hard_voxels = int(np.count_nonzero(pred_hard_orig))

    return {
        "case": case_key,
        "input_t1": str(t1_path.resolve()),
        "saved_input_t1": str(input_path.resolve()),
        "pred_hard": str(hard_path.resolve()),
        "threshold": float(threshold),
        "hard_voxels": hard_voxels,
        "input_shape": "x".join(str(v) for v in bundle.input_shape),
        "edge_components_removed": removed_components,
        "run_seconds": round(float(elapsed), 4),
    }


def run_inference_on_prepared_t1(
    model_dir: Path,
    t1_paths: list[Path],
    output_root: Path,
    threshold: float = 0.50,
) -> dict:
    if not t1_paths:
        raise ValueError("No input T1 files were provided for inference.")

    threshold = float(threshold)
    if threshold < 0.0 or threshold > 1.0:
        raise ValueError(f"Threshold must be in [0, 1], got {threshold}")

    t1_paths = [Path(p).expanduser().resolve() for p in t1_paths]
    for t1 in t1_paths:
        if not t1.exists():
            raise FileNotFoundError(f"Missing input file: {t1}")

    output_root = Path(output_root).expanduser().resolve()
    run_tag = datetime.now().strftime("run_%Y%m%d_%H%M%S")
    run_dir = output_root / run_tag
    input_dir = run_dir / "input_t1"
    hard_dir = run_dir / "hard"
    manifest_path = run_dir / "manifest.csv"

    bundle = load_model_bundle(model_dir)

    rows: list[dict] = []
    errors: list[dict] = []
    for t1 in t1_paths:
        try:
            rows.append(_predict_one(t1, bundle, threshold, input_dir, hard_dir))
        except Exception as exc:
            errors.append({"input_t1": str(t1), "error": str(exc)})

    write_manifest(rows, manifest_path)

    if not rows:
        raise RuntimeError(
            "Inference failed for all cases. "
            f"First error: {errors[0]['error'] if errors else 'unknown error'}"
        )

    return {
        "model_dir": str(bundle.model_dir),
        "config_path": str(bundle.config_path),
        "weights_path": str(bundle.weights_path),
        "input_shape": bundle.input_shape,
        "threshold": threshold,
        "run_dir": str(run_dir),
        "input_dir": str(input_dir),
        "hard_dir": str(hard_dir),
        "manifest": str(manifest_path),
        "rows": rows,
        "errors": errors,
    }