"""Atlas-based structure labelling and structure-aware risk scoring.

Uses the Harvard-Oxford cortical atlas (real, peer-reviewed anatomical labels)
to identify eloquent structures and compute true distance-to-structure risk for
a planned trajectory. For patient images not already in atlas space, the atlas
is resampled onto the patient grid (nearest-neighbour) — a registration step
that becomes patient-specific deformable registration in the product roadmap.
"""
from __future__ import annotations
import numpy as np
from .volume import Volume
from .planning import Trajectory, risk_band

# Eloquent / functionally critical cortex (Harvard-Oxford label names)
ELOQUENT_DEFAULT = [
    "Precentral Gyrus",                                  # primary motor
    "Postcentral Gyrus",                                 # primary sensory
    "Inferior Frontal Gyrus, pars opercularis",          # Broca / speech
    "Inferior Frontal Gyrus, pars triangularis",         # Broca / speech
    "Superior Temporal Gyrus, posterior division",       # language (Wernicke region)
    "Supramarginal Gyrus, posterior division",           # language
    "Angular Gyrus",                                     # language
    "Intracalcarine Cortex",                             # primary visual
    "Occipital Pole",                                    # visual
]

# Deep / subcortical critical structures (Harvard-Oxford subcortical atlas;
# load with atlas_name="sub-maxprob-thr25-2mm"). Matched by substring.
SUBCORTICAL_DEFAULT = [
    "Thalamus", "Caudate", "Putamen", "Pallidum",
    "Hippocampus", "Amygdala", "Accumbens", "Brain-Stem",
]


# White-matter fiber tracts (Juelich histological atlas label substrings).
TRACTS_DEFAULT = [
    "WM Corticospinal tract",   # motor pathway
    "WM Optic radiation",       # visual pathway
    "WM Cingulum", "WM Fornix", "WM Callosal body", "WM Acoustic radiation",
]


def load_harvard_oxford(atlas_name="cort-maxprob-thr25-2mm"):
    """Return (atlas_volume, labels). Downloads on first use (cached)."""
    from nilearn.datasets import fetch_atlas_harvard_oxford
    a = fetch_atlas_harvard_oxford(atlas_name)
    img = a.maps
    data = np.asanyarray(img.dataobj).astype(np.int16)
    return Volume(data, img.affine), list(a.labels)


def load_juelich(atlas_name="maxprob-thr25-1mm"):
    """Return (atlas_volume, labels) for the Juelich histological atlas, which
    includes real white-matter fiber tracts (corticospinal tract, optic
    radiation, etc.). Downloads on first use (cached)."""
    from nilearn.datasets import fetch_atlas_juelich
    a = fetch_atlas_juelich(atlas_name)
    data = np.asanyarray(a.maps.dataobj).astype(np.int16)
    return Volume(data, a.maps.affine), list(a.labels)


def _to_nibabel(vol: Volume):
    import nibabel as nib
    return nib.Nifti1Image(np.asarray(vol.data), vol.affine)


def resample_atlas_to(atlas_vol: Volume, target_vol: Volume) -> Volume:
    """Resample a label atlas onto a target image grid (nearest-neighbour)."""
    from nilearn.image import resample_to_img
    res = resample_to_img(_to_nibabel(atlas_vol), _to_nibabel(target_vol),
                          interpolation="nearest", force_resample=True, copy_header=True)
    return Volume(np.asanyarray(res.dataobj).astype(np.int16), res.affine)


def structure_mask(atlas_vol: Volume, labels, names) -> Volume:
    """Binary mask Volume covering the named structures."""
    idx = [labels.index(n) for n in names if n in labels]
    m = np.isin(atlas_vol.data, idx).astype(np.uint8)
    return Volume(m, atlas_vol.affine.copy())


def structure_risk(traj: Trajectory, atlas_vol: Volume, labels, names=None, n=240) -> dict:
    """Minimum clearance (mm) from the trajectory to any selected eloquent
    structure, the nearest structure name, and a risk band — single pass."""
    names = names or [l for l in ELOQUENT_DEFAULT if l in labels]
    idx = [labels.index(nm) for nm in names if nm in labels]
    A = np.asarray(atlas_vol.data)
    sel = np.isin(A, idx)
    coords = np.argwhere(sel)
    if coords.size == 0:
        return {"clearance_mm": float("inf"), "nearest": None, "band": "LOW", "per_structure": {}}
    lab_at = A[sel]
    homog = np.hstack([coords.astype(np.float64), np.ones((coords.shape[0], 1))])
    world = (atlas_vol.affine @ homog.T).T[:, :3]
    best_d, best_lab = float("inf"), None
    per = {}
    for p in traj.sample(n):
        dd = np.linalg.norm(world - p[None, :], axis=1)
        j = int(np.argmin(dd))
        if dd[j] < best_d:
            best_d, best_lab = float(dd[j]), int(lab_at[j])
    # per-structure minimum (for a richer report)
    for nm in names:
        if nm not in labels:
            continue
        li = labels.index(nm)
        w = world[lab_at == li]
        if w.size == 0:
            continue
        dmin = min(np.min(np.linalg.norm(w - p[None, :], axis=1)) for p in traj.sample(n))
        per[nm] = float(dmin)
    return {"clearance_mm": best_d,
            "nearest": labels[best_lab] if best_lab is not None else None,
            "band": risk_band(best_d),
            "per_structure": per}
