"""CT support: Hounsfield-unit windowing and tissue segmentation.

Head CT is the primary modality for haemorrhage/trauma/hydrocephalus, and it is
quantitative (Hounsfield Units), so display uses clinical window/level presets
and tissue is segmented by HU thresholds rather than atlas registration. This
module provides:

* HU window presets (brain, stroke, subdural, blood, bone) and ``apply_window``;
* ``is_ct`` modality detection (from metadata or the HU range);
* HU-threshold segmentation: skull/bone, brain extraction, CSF/ventricles, and a
  generic ``segment_hu`` (acute clot lives in ``sulcora.hematoma``).

Pure NumPy + scipy; unit-tested headlessly.
"""
from __future__ import annotations
import numpy as np
from .volume import Volume

# Clinical head-CT window presets as (level, width) in HU.
WINDOWS = {
    "brain": (40.0, 80.0),
    "stroke": (35.0, 35.0),
    "subdural": (75.0, 215.0),
    "blood": (60.0, 130.0),
    "bone": (500.0, 2500.0),
    "soft": (40.0, 350.0),
}


def apply_window(volume: Volume, level: float, width: float) -> Volume:
    """Window/level a volume to a uint8 display image (lo=level-width/2)."""
    lo, hi = level - width / 2.0, level + width / 2.0
    d = np.asarray(volume.data, dtype=np.float32)
    d = np.clip((d - lo) / max(hi - lo, 1e-6), 0.0, 1.0)
    return Volume((d * 255).astype(np.uint8), volume.affine.copy())


def window_preset(volume: Volume, name="brain") -> Volume:
    if name not in WINDOWS:
        raise ValueError(f"Unknown CT window {name!r}; choose from {sorted(WINDOWS)}.")
    return apply_window(volume, *WINDOWS[name])


def is_ct(volume: Volume) -> bool:
    """True if the volume looks like CT — by modality metadata if present, else
    by the HU range (air below ~-200 and bone above ~300)."""
    mod = getattr(volume, "meta", {}).get("modality") if hasattr(volume, "meta") else None
    if mod:
        return str(mod).upper() == "CT"
    d = np.asarray(volume.data)
    return float(d.min()) < -200.0 and float(d.max()) > 300.0


def segment_hu(volume: Volume, lo: float, hi: float, keep_largest=False,
               min_volume_cc=0.0, close=False) -> Volume:
    """Binary mask of voxels with HU in [lo, hi], optionally cleaned."""
    from scipy import ndimage
    d = np.asarray(volume.data)
    m = (d >= lo) & (d <= hi)
    if close:
        m = ndimage.binary_closing(m, iterations=1)
    if keep_largest or min_volume_cc > 0:
        lab, n = ndimage.label(m)
        if n:
            sizes = np.bincount(lab.ravel()); sizes[0] = 0
            voxel_cc = float(np.prod(volume.spacing)) / 1000.0
            if keep_largest:
                keep = {int(np.argmax(sizes))}
            else:
                keep = {i for i in range(1, n + 1) if sizes[i] * voxel_cc >= min_volume_cc}
            m = np.isin(lab, list(keep)) if keep else np.zeros_like(m)
    return Volume(m.astype(np.uint8), volume.affine.copy())


def skull_mask(volume: Volume, hu_lo=300.0) -> Volume:
    """Bone/skull mask (HU above ~300)."""
    return segment_hu(volume, hu_lo, 1e6)


def brain_mask(volume: Volume, hu_lo=-20.0, hu_hi=100.0) -> Volume:
    """Intracranial soft-tissue (brain) mask: soft-tissue HU, largest component,
    holes filled — useful to restrict planning to the brain and drop skull/air."""
    from scipy import ndimage
    d = np.asarray(volume.data)
    m = (d >= hu_lo) & (d <= hu_hi)
    lab, n = ndimage.label(m)
    if n:
        sizes = np.bincount(lab.ravel()); sizes[0] = 0
        m = lab == int(np.argmax(sizes))
    m = ndimage.binary_fill_holes(m)
    return Volume(m.astype(np.uint8), volume.affine.copy())


def csf_mask(volume: Volume, hu_lo=-5.0, hu_hi=18.0) -> Volume:
    """CSF / ventricle mask (HU ~0-15), restricted to inside the brain."""
    brain = np.asarray(brain_mask(volume).data)
    d = np.asarray(volume.data)
    m = (d >= hu_lo) & (d <= hu_hi) & (brain > 0)
    return Volume(m.astype(np.uint8), volume.affine.copy())
