"""White-matter tract avoidance (DTI tractography).

Eloquent white-matter tracts — corticospinal tract, arcuate fasciculus, optic
radiations — are avoided the same way as vessels/cortex: as a risk structure the
optimizer must clear. This module turns tractography (streamlines from an
external DTI pipeline, or a precomputed tract mask) into a Sulcora risk
structure, so a planned corridor keeps its distance from the tracts.

Streamline tractography itself is done upstream (MRtrix/DSI Studio/DIPY); here we
rasterize the result and score clearance. Research software — not a medical device.
"""
from __future__ import annotations
import numpy as np
from .volume import Volume


def streamlines_to_mask(streamlines, affine, shape, radius_vox=1) -> Volume:
    """Rasterize world-space streamlines (list of N×3 mm arrays) onto a voxel
    grid, optionally dilating by `radius_vox` to give the tract thickness."""
    from scipy import ndimage
    inv = np.linalg.inv(np.asarray(affine, float))
    mask = np.zeros(tuple(int(s) for s in shape), np.uint8)
    for sl in streamlines:
        pts = np.asarray(sl, float)
        if pts.ndim != 2 or pts.shape[0] == 0:
            continue
        homog = np.hstack([pts, np.ones((pts.shape[0], 1))])
        vox = np.round((inv @ homog.T).T[:, :3]).astype(int)
        ok = np.all((vox >= 0) & (vox < np.array(mask.shape)), axis=1)
        v = vox[ok]
        if v.size:
            mask[v[:, 0], v[:, 1], v[:, 2]] = 1
    if radius_vox > 0 and mask.any():
        mask = ndimage.binary_dilation(mask, iterations=int(radius_vox)).astype(np.uint8)
    return Volume(mask, np.asarray(affine, float))


def load_streamlines(path):
    """Load .trk/.tck streamlines (world mm) + affine via nibabel."""
    import nibabel as nib
    trk = nib.streamlines.load(path)
    return list(trk.streamlines), trk.affine


def tract_mask_from_file(path, affine, shape, radius_vox=1) -> Volume:
    sl, _ = load_streamlines(path)
    return streamlines_to_mask(sl, affine, shape, radius_vox)


def tract_risk_structure(mask, name="tract", min_clearance_mm=2.0, weight=2.0):
    """Wrap a tract mask as a RiskStructure for the optimizer / risk readout."""
    from .optimizer import RiskStructure
    return RiskStructure(name, mask, min_clearance_mm, weight)


def tract_clearance(traj, mask, n=200) -> float:
    """Minimum distance (mm) from a trajectory to the tract mask."""
    from scipy.spatial import cKDTree
    from .vessels import _mask_world_points
    pts = _mask_world_points(mask)
    if len(pts) == 0:
        return float("inf")
    return float(np.min(cKDTree(pts).query(traj.sample(n))[0]))
