"""Patient-specific segmentation of eloquent structures and ventricles.

Instead of registering a population atlas onto the patient (see ``atlas.py``),
this module labels critical structures **directly from the patient's own T1 MRI**
using a pretrained deep-learning model. That removes the registration step that
is the weakest link in atlas-based risk: a SynthSeg/TotalSegmentator model
produces a per-voxel label map in the patient's native space, from which we build
the same eloquent-cortex / ventricle / deep-structure masks the planner already
consumes.

Two backends are supported, both optional and lazily imported:

* ``"synthseg"``       — FreeSurfer's contrast-agnostic SynthSeg (``mri_synthseg``,
                         run with ``--parc`` for cortical parcellation). Produces
                         standard FreeSurfer aseg/aparc labels. Requires the
                         ``mri_synthseg`` command on ``PATH`` (ships with
                         FreeSurfer; no pip dependency).
* ``"totalsegmentator"`` — TotalSegmentator's ``brain_structures`` MR task.
                         Requires the ``[seg]`` install extra
                         (``pip install -e ".[seg]"``).

Everything *downstream* of the model — turning a label map into structure masks,
scoring trajectory clearance, and comparing against the atlas labels — is pure
NumPy and is unit-tested without any model or network access.

IMPORTANT: Research software only. Not a medical device. Not for clinical,
diagnostic, or treatment use. Automated segmentations are unverified and must
never be relied on for patient care. See README and docs/INTENDED_USE.md.
"""
from __future__ import annotations
from dataclasses import dataclass, field
import numpy as np
from .volume import Volume
from .planning import Trajectory, risk_band


class BackendUnavailable(RuntimeError):
    """Raised when a requested segmentation backend is not installed/available."""


# --------------------------------------------------------------------------- #
# Canonical structure taxonomy
#
# Each structure is matched against a backend's label names (and against the
# Harvard-Oxford atlas labels, for comparison) by case-insensitive substring.
# The patterns are chosen to match both FreeSurfer/SynthSeg names
# (e.g. "Left-Lateral-Ventricle", "ctx-lh-precentral") and Harvard-Oxford names
# (e.g. "Left Lateral Ventricle", "Precentral Gyrus"), so one matcher serves the
# segmentation map and the atlas alike.
# --------------------------------------------------------------------------- #
@dataclass(frozen=True)
class StructureSpec:
    key: str            # canonical identifier
    group: str          # "eloquent" | "ventricle" | "deep"
    patterns: tuple     # case-insensitive substrings identifying the structure


STRUCTURES = [
    # ventricles. Patterns match FreeSurfer/SynthSeg ("Left-Lateral-Ventricle"),
    # Harvard-Oxford ("Lateral Ventricle"), and TotalSegmentator ventricle_parts
    # ("ventricle_body_left", "ventricle_frontal_horn_right", ...).
    StructureSpec("lateral_ventricle", "ventricle",
                  ("lateral-ventricle", "lateral ventricle", "inf-lat-vent",
                   "ventricle_frontal_horn", "ventricle_occipital_horn", "ventricle_body",
                   "ventricle_temporal_horn", "ventricle_trigone")),
    StructureSpec("third_ventricle", "ventricle",
                  ("3rd-ventricle", "3rd ventricle", "third ventricle", "third_ventricle", "ventricle_3")),
    StructureSpec("fourth_ventricle", "ventricle",
                  ("4th-ventricle", "4th ventricle", "fourth ventricle", "fourth_ventricle", "ventricle_4")),
    # deep / subcortical
    StructureSpec("thalamus", "deep", ("thalamus",)),
    StructureSpec("caudate", "deep", ("caudate",)),               # also "caudate_nucleus" (TS)
    StructureSpec("putamen", "deep", ("putamen",)),
    StructureSpec("pallidum", "deep", ("pallidum",)),
    # TotalSegmentator labels putamen+pallidum together as the lentiform nucleus.
    StructureSpec("lentiform_nucleus", "deep", ("lentiform",)),
    StructureSpec("hippocampus", "deep", ("hippocampus",)),
    StructureSpec("amygdala", "deep", ("amygdala",)),
    StructureSpec("accumbens", "deep", ("accumbens",)),
    StructureSpec("brainstem", "deep", ("brain-stem", "brainstem", "brain_stem")),
    # eloquent cortex (requires a cortical parcellation: SynthSeg --parc,
    # FreeSurfer aparc, or TotalSegmentator brain_structures)
    StructureSpec("precentral_gyrus", "eloquent", ("precentral",)),            # primary motor
    StructureSpec("postcentral_gyrus", "eloquent", ("postcentral",)),          # primary sensory
    StructureSpec("pars_opercularis", "eloquent", ("parsopercularis", "pars opercularis")),   # Broca
    StructureSpec("pars_triangularis", "eloquent", ("parstriangularis", "pars triangularis")),  # Broca
    StructureSpec("superior_temporal", "eloquent", ("superiortemporal", "superior temporal")),  # Wernicke region
    StructureSpec("supramarginal", "eloquent", ("supramarginal",)),            # language
    StructureSpec("angular_gyrus", "eloquent", ("angular",)),                  # language
]

STRUCTURE_BY_KEY = {s.key: s for s in STRUCTURES}
GROUPS = ("eloquent", "ventricle", "deep")


def _selected_specs(keys=None, groups=None):
    specs = STRUCTURES
    if groups is not None:
        groups = {groups} if isinstance(groups, str) else set(groups)
        specs = [s for s in specs if s.group in groups]
    if keys is not None:
        keys = {keys} if isinstance(keys, str) else set(keys)
        specs = [s for s in specs if s.key in keys]
    return specs


# --------------------------------------------------------------------------- #
# Label-map -> mask (works for any int label Volume + {id: name} or [names])
# --------------------------------------------------------------------------- #
def _iter_labels(labels):
    """Yield (int_id, name) for either a {id: name} dict or a [name, ...] list."""
    if isinstance(labels, dict):
        return [(int(i), n) for i, n in labels.items()]
    return list(enumerate(labels))


def matching_ids(labels, patterns):
    """Integer label ids whose name contains any of the (case-insensitive) patterns."""
    pats = [p.lower() for p in patterns]
    return [i for i, name in _iter_labels(labels)
            if name and any(p in str(name).lower() for p in pats)]


def mask_from_labelmap(label_vol: Volume, labels, patterns) -> Volume:
    """Binary mask Volume over the labels matching `patterns` in `label_vol`."""
    ids = matching_ids(labels, patterns)
    if ids:
        m = np.isin(np.asarray(label_vol.data), ids).astype(np.uint8)
    else:
        m = np.zeros(label_vol.shape, np.uint8)
    return Volume(m, label_vol.affine.copy())


def dice(a: Volume, b: Volume) -> float:
    """Dice overlap of two binary masks on a shared grid. NaN if both empty."""
    A = np.asarray(a.data) > 0
    B = np.asarray(b.data) > 0
    denom = int(A.sum() + B.sum())
    if denom == 0:
        return float("nan")
    return float(2 * int(np.logical_and(A, B).sum()) / denom)


def overlap_metrics(seg_mask: Volume, ref_mask: Volume) -> dict:
    """Dice / Jaccard / voxel counts comparing a segmentation mask to a reference."""
    A = np.asarray(seg_mask.data) > 0
    B = np.asarray(ref_mask.data) > 0
    inter = int(np.logical_and(A, B).sum())
    union = int(np.logical_or(A, B).sum())
    return {
        "seg_voxels": int(A.sum()),
        "ref_voxels": int(B.sum()),
        "dice": dice(seg_mask, ref_mask),
        "jaccard": float(inter / union) if union else float("nan"),
    }


# --------------------------------------------------------------------------- #
# Segmentation result
# --------------------------------------------------------------------------- #
@dataclass
class SegmentationResult:
    """A patient-space label map plus its {id: name} lookup.

    label_volume : Volume of integer structure labels in the patient's grid.
    labels       : {int_id: structure_name} (e.g. FreeSurfer/TotalSegmentator).
    backend      : which engine produced it.
    """
    label_volume: Volume
    labels: dict = field(default_factory=dict)
    backend: str = "unknown"

    def structure_mask(self, keys=None, groups=None) -> Volume:
        """Binary mask over the union of the selected structures."""
        out = np.zeros(self.label_volume.shape, np.uint8)
        for s in _selected_specs(keys, groups):
            out |= mask_from_labelmap(self.label_volume, self.labels, s.patterns).data
        return Volume(out, self.label_volume.affine.copy())

    def present_structures(self, keys=None, groups=None) -> list:
        """Canonical keys actually found (non-empty) in this segmentation."""
        found = []
        for s in _selected_specs(keys, groups):
            if mask_from_labelmap(self.label_volume, self.labels, s.patterns).data.any():
                found.append(s.key)
        return found


def structure_risk(traj: Trajectory, result: SegmentationResult,
                   keys=None, groups=None, n=240) -> dict:
    """Trajectory clearance to the nearest segmented critical structure.

    Mirrors ``atlas.structure_risk`` but uses the patient-specific segmentation
    instead of a registered atlas.
    """
    present = result.present_structures(keys, groups)
    if not present:
        return {"clearance_mm": float("inf"), "nearest": None, "band": "LOW", "per_structure": {}}
    best, nearest, per = float("inf"), None, {}
    for k in present:
        m = mask_from_labelmap(result.label_volume, result.labels, STRUCTURE_BY_KEY[k].patterns)
        d, _ = traj.min_clearance_to_mask(m, n=n)
        per[k] = float(d)
        if d < best:
            best, nearest = float(d), k
    return {"clearance_mm": round(best, 2), "nearest": nearest,
            "band": risk_band(best), "per_structure": per}


# --------------------------------------------------------------------------- #
# Comparison against the atlas-based labels (the registration-free vs.
# registration-based experiment). Requires the [sample] extra (nilearn) for the
# atlas; used by scripts/make_segmentation_proof.py, not by the unit tests.
# --------------------------------------------------------------------------- #
def compare_to_atlas(result: SegmentationResult, groups=("eloquent", "ventricle", "deep"),
                     cort_atlas="cort-maxprob-thr25-2mm",
                     sub_atlas="sub-maxprob-thr25-2mm") -> list:
    """Per-structure overlap between the patient segmentation and the Harvard-Oxford
    atlas resampled into the same patient grid. Returns a list of metric dicts.
    """
    from . import atlas as A
    target = result.label_volume
    cort_v, cort_l = A.load_harvard_oxford(cort_atlas)
    cort_v = A.resample_atlas_to(cort_v, target)
    sub_v, sub_l = A.load_harvard_oxford(sub_atlas)
    sub_v = A.resample_atlas_to(sub_v, target)

    rows = []
    for spec in _selected_specs(groups=groups):
        seg_mask = mask_from_labelmap(target, result.labels, spec.patterns)
        cort_mask = mask_from_labelmap(cort_v, cort_l, spec.patterns)
        sub_mask = mask_from_labelmap(sub_v, sub_l, spec.patterns)
        # take whichever HO atlas actually carries the structure
        atlas_mask = cort_mask if cort_mask.data.sum() >= sub_mask.data.sum() else sub_mask
        if seg_mask.data.sum() == 0 and atlas_mask.data.sum() == 0:
            continue  # neither source labels this structure -> nothing to compare
        m = overlap_metrics(seg_mask, atlas_mask)
        m.update({"structure": spec.key, "group": spec.group})
        rows.append(m)
    return rows


# --------------------------------------------------------------------------- #
# Backends (lazily imported; not exercised by the unit tests)
# --------------------------------------------------------------------------- #
# FreeSurfer aseg + aparc (DKT) labels emitted by SynthSeg --parc. Subset
# covering the structures Sulcora scores; ids follow the standard FreeSurfer LUT.
FREESURFER_LABELS = {
    0: "Unknown",
    2: "Left-Cerebral-White-Matter", 3: "Left-Cerebral-Cortex",
    4: "Left-Lateral-Ventricle", 5: "Left-Inf-Lat-Vent",
    7: "Left-Cerebellum-White-Matter", 8: "Left-Cerebellum-Cortex",
    10: "Left-Thalamus", 11: "Left-Caudate", 12: "Left-Putamen", 13: "Left-Pallidum",
    14: "3rd-Ventricle", 15: "4th-Ventricle", 16: "Brain-Stem",
    17: "Left-Hippocampus", 18: "Left-Amygdala", 26: "Left-Accumbens-area", 28: "Left-VentralDC",
    41: "Right-Cerebral-White-Matter", 42: "Right-Cerebral-Cortex",
    43: "Right-Lateral-Ventricle", 44: "Right-Inf-Lat-Vent",
    46: "Right-Cerebellum-White-Matter", 47: "Right-Cerebellum-Cortex",
    49: "Right-Thalamus", 50: "Right-Caudate", 51: "Right-Putamen", 52: "Right-Pallidum",
    53: "Right-Hippocampus", 54: "Right-Amygdala", 58: "Right-Accumbens-area", 60: "Right-VentralDC",
    # aparc DKT cortical parcels (eloquent subset; lh 10xx / rh 20xx)
    1018: "ctx-lh-parsopercularis", 1020: "ctx-lh-parstriangularis",
    1022: "ctx-lh-postcentral", 1024: "ctx-lh-precentral",
    1030: "ctx-lh-superiortemporal", 1031: "ctx-lh-supramarginal",
    2018: "ctx-rh-parsopercularis", 2020: "ctx-rh-parstriangularis",
    2022: "ctx-rh-postcentral", 2024: "ctx-rh-precentral",
    2030: "ctx-rh-superiortemporal", 2031: "ctx-rh-supramarginal",
}


def _as_input_path(image, work_dir):
    """Return a filesystem path to a NIfTI for `image` (a path or a Volume)."""
    import os
    if isinstance(image, str):
        return image
    if isinstance(image, Volume):
        import nibabel as nib
        path = os.path.join(work_dir, "sulcora_seg_input.nii.gz")
        nib.save(nib.Nifti1Image(np.asarray(image.data).astype(np.float32), image.affine), path)
        return path
    raise TypeError("image must be a path (str) or a sulcora Volume")


def _load_label_volume(path) -> Volume:
    import nibabel as nib
    img = nib.load(path)
    return Volume(np.asanyarray(img.dataobj).astype(np.int32), img.affine)


def segment_synthseg(image, work_dir=None, command="mri_synthseg",
                     parc=True, robust=True) -> SegmentationResult:
    """Segment a T1 with FreeSurfer SynthSeg (contrast-agnostic).

    Shells out to ``mri_synthseg`` (ships with FreeSurfer). Use ``--parc`` to add
    the cortical parcellation needed for eloquent gyri. Raises BackendUnavailable
    if the command is not on PATH.
    """
    import os, shutil, subprocess, tempfile
    if shutil.which(command) is None:
        raise BackendUnavailable(
            f"'{command}' not found on PATH. SynthSeg ships with FreeSurfer; "
            "install FreeSurfer and ensure mri_synthseg is available.")
    work_dir = work_dir or tempfile.mkdtemp(prefix="sulcora_synthseg_")
    os.makedirs(work_dir, exist_ok=True)
    inp = _as_input_path(image, work_dir)
    out = os.path.join(work_dir, "synthseg_labels.nii.gz")
    cmd = [command, "--i", inp, "--o", out]
    if parc:
        cmd.append("--parc")
    if robust:
        cmd.append("--robust")
    subprocess.run(cmd, check=True)
    return SegmentationResult(_load_label_volume(out), dict(FREESURFER_LABELS), backend="synthseg")


# TotalSegmentator MR brain tasks that are not openly available and require a
# free (non-commercial) license key — https://totalsegmentator.com
LICENSED_TS_TASKS = frozenset({
    "brain_structures", "ventricle_parts", "head_glands_cavities", "head_muscles",
    "headneck_bones_vessels", "headneck_muscles", "face", "face_mr",
    "craniofacial_structures", "brain_aneurysm",
})


def segment_totalsegmentator(image, work_dir=None, fast=False,
                             tasks=("brain_structures", "ventricle_parts"),
                             license_number=None) -> SegmentationResult:
    """Segment a T1 MR with TotalSegmentator.

    Runs one or more MR brain tasks and merges their label maps into a single
    Volume. The default pairing — ``brain_structures`` (thalamus, caudate,
    lentiform nucleus, brainstem, lobes, whole ventricle) plus ``ventricle_parts``
    (lateral horns + 3rd/4th ventricle) — covers ventricles and the deep
    structures Sulcora scores. Later tasks overwrite earlier ones where they
    overlap (so the detailed ventricle parts replace the coarse whole-ventricle
    label).

    Requires the ``[seg]`` extra. These MR brain tasks are not openly available
    and need a free non-commercial license key, supplied via ``license_number=``,
    the ``SULCORA_TS_LICENSE`` env var, or a one-time ``totalseg_set_license -l
    <KEY>``. Raises BackendUnavailable if the package or a required license is
    missing.
    """
    import os, tempfile
    try:
        from totalsegmentator.python_api import totalsegmentator
        from totalsegmentator.map_to_binary import class_map
    except ImportError as e:
        raise BackendUnavailable(
            "TotalSegmentator is not installed. Install the segmentation extra: "
            'pip install -e ".[seg]"') from e
    if isinstance(tasks, str):
        tasks = (tasks,)

    license_number = license_number or os.environ.get("SULCORA_TS_LICENSE")
    if license_number:
        try:
            from totalsegmentator.config import set_license_number
            set_license_number(license_number)
        except Exception:
            pass
    # Proactively check the license for gated tasks so we fail with guidance
    # rather than TotalSegmentator's internal sys.exit().
    need_license = [t for t in tasks if t in LICENSED_TS_TASKS]
    if need_license:
        try:
            from totalsegmentator.config import has_valid_license_offline
            status = has_valid_license_offline()
            status = status[0] if isinstance(status, (tuple, list)) else status
        except Exception:
            status = None
        if status in ("missing_license", "invalid_license"):
            raise BackendUnavailable(
                f"TotalSegmentator task(s) {need_license} require a free, non-commercial "
                "license key (status: " + str(status) + "). Acquire one at "
                "https://backend.totalsegmentator.com/license-academic/ then set it via "
                "`totalseg_set_license -l <KEY>`, the SULCORA_TS_LICENSE env var, or the "
                "license_number= argument.")

    work_dir = work_dir or tempfile.mkdtemp(prefix="sulcora_totalseg_")
    os.makedirs(work_dir, exist_ok=True)
    inp = _as_input_path(image, work_dir)

    merged, affine, labels = None, None, {0: "background"}
    for ti, task in enumerate(tasks):
        out = os.path.join(work_dir, f"totalseg_{task}.nii.gz")
        try:
            totalsegmentator(inp, out, task=task, ml=True, fast=fast)
        except SystemExit as e:  # TotalSegmentator exits hard on a license problem
            raise BackendUnavailable(
                f"TotalSegmentator could not run task '{task}' — most likely a missing or "
                "invalid license for a gated brain task. See "
                "https://backend.totalsegmentator.com/license-academic/") from e
        vol = _load_label_volume(out)
        data = np.asarray(vol.data)
        if merged is None:
            merged, affine = np.zeros(data.shape, np.int32), vol.affine
        base = (ti + 1) * 1000  # offset each task's ids so they never collide
        for i, name in class_map[task].items():
            gid = base + int(i)
            labels[gid] = name
            merged[data == int(i)] = gid  # later tasks overwrite earlier overlaps
    return SegmentationResult(Volume(merged, affine), labels, backend="totalsegmentator")


_BACKENDS = {"synthseg": segment_synthseg, "totalsegmentator": segment_totalsegmentator}


def segment_t1(image, backend="auto", **kwargs) -> SegmentationResult:
    """Segment a patient T1 MRI into labelled structures.

    image   : path to a NIfTI/DICOM-derived T1, or a sulcora Volume.
    backend : "synthseg", "totalsegmentator", or "auto" (try each in turn).
    kwargs  : forwarded to the chosen backend.
    """
    if backend == "auto":
        errors = []
        for name in ("totalsegmentator", "synthseg"):
            try:
                return _BACKENDS[name](image, **kwargs)
            except BackendUnavailable as e:
                errors.append(f"{name}: {e}")
        raise BackendUnavailable(
            "No segmentation backend available.\n  " + "\n  ".join(errors))
    if backend not in _BACKENDS:
        raise ValueError(f"Unknown backend {backend!r}; choose from {sorted(_BACKENDS)} or 'auto'.")
    return _BACKENDS[backend](image, **kwargs)
