"""Post-processes a TotalSegmentator multi-label output into a single-label
mask containing only the requested structure.

Why this exists: since segmentation_runner.py always runs the full task
(all ~117 CT / ~56 MRI structures, see its docstring for why roi_subset was
dropped), the raw output NIfTI contains every structure TotalSegmentator
found, not just the one the user asked about. Isolating "just this one
label" was first attempted purely client-side in app.js via Niivue's
setColormapLabel, but that did not behave like a sparse index->color lookup
in testing -- every structure rendered in the overlay color, not just the
target. Zeroing out everything but the target label here, server-side,
sidesteps that entirely: the file handed to the browser only ever contains
0 and the target's label id, so however the client colors it, only the
target structure can show up.
"""

import nibabel as nib
import numpy as np


def extract_single_label(segmentation_path, label_id, output_path):
    """Reads segmentation_path, zeroes every voxel not equal to label_id,
    and writes the result to output_path. Returns output_path."""
    img = nib.load(str(segmentation_path))
    data = np.asanyarray(img.dataobj)
    mask = (np.rint(data).astype(np.int32) == int(label_id)).astype(np.uint8) * int(label_id)

    out_img = nib.Nifti1Image(mask, img.affine, header=img.header)
    out_img.set_data_dtype(np.uint8)
    nib.save(out_img, str(output_path))
    return output_path