<commit_before>import numpy as np
import logging
from scipy.spatial.distance import mahalanobis, cdist
from scipy.stats import zscore

import nibabel as nib
from tqdm.auto import tqdm

import dipy.data as dpd
import dipy.tracking.streamline as dts
import dipy.tracking.streamlinespeed as dps
from dipy.segment.bundles import RecoBundles
from dipy.align.streamlinear import whole_brain_slr
from dipy.stats.analysis import gaussian_weights
import dipy.core.gradients as dpg
from dipy.io.stateful_tractogram import StatefulTractogram, Space

import AFQ.registration as reg
import AFQ.utils.models as ut
import AFQ.utils.volume as auv
import AFQ.data as afd

__all__ = ["Segmentation"]


def _resample_tg(tg, n_points):
    # reformat for dipy's set_number_of_points
    if isinstance(tg, np.ndarray):
        if len(tg.shape) > 2:
            streamlines = tg.tolist()
            streamlines = [np.asarray(item) for item in streamlines]
    else:
        streamlines = tg.streamlines

    return dps.set_number_of_points(streamlines, n_points)


class Segmentation:
    def __init__(self,
                 nb_points=False,
                 seg_algo='AFQ',
                 progressive=True,
                 greater_than=50,
                 rm_small_clusters=50,
                 model_clust_thr=40,
                 reduction_thr=40,
                 refine=False,
                 pruning_thr=6,
                 b0_threshold=0,
                 prob_threshold=0,
                 rng=None,
                 return_idx=False,
                 filter_by_endpoints=True,
                 dist_to_aal=4):
        """
        Segment streamlines into bundles.
        Parameters
        ----------
        nb_points : int, boolean
            Resample streamlines to nb_points number of points.
            If False, no resampling is done. Default: False
        seg_algo : string
            Algorithm for segmentation (case-insensitive):
            'AFQ': Segment streamlines into bundles,
                based on inclusion/exclusion ROIs.
            'Reco': Segment streamlines using the RecoBundles algorithm
            [Garyfallidis2017].
            Default: 'AFQ'
        rm_small_clusters : int
            Using RecoBundles Algorithm.
            Remove clusters that have less than this value
                during whole brain SLR.
            Default: 50
        progressive : boolean, optional
            Using RecoBundles Algorithm.
            Whether or not to use progressive technique
                during whole brain SLR.
            Default: True.
        greater_than : int, optional
            Using RecoBundles Algorithm.
            Keep streamlines that have length greater than this value
                during whole brain SLR.
            Default: 50.
        b0_theshold : float.
            Using AFQ Algorithm.
            All b-values with values less than or equal to `bo_threshold` are
            considered as b0s i.e. without diffusion weighting.
            Default: 0.
        prob_threshold : float.
            Using AFQ Algorithm.
            Initial cleaning of fiber groups is done using probability maps
            from [Hua2008]_. Here, we choose an average probability that
            needs to be exceeded for an individual streamline to be retained.
            Default: 0.
        rng : RandomState
            If None, creates RandomState. Used in RecoBundles Algorithm.
            Default: None.
        return_idx : bool
            Whether to return the indices in the original streamlines as part
            of the output of segmentation.
        filter_by_endpoints: bool
            Whether to filter the bundles based on their endpoints relative
            to regions defined in the AAL atlas. Applies only to the waypoint
            approach (XXX for now). Default: True.
        dist_to_aal : float
            If filter_by_endpoints is True, this is the distance from the
            endpoints to the AAL atlas ROIs that is required.
        References
        ----------
        .. [Hua2008] <NAME>, <NAME>, <NAME>, <NAME>, <NAME>, et al. (2008)
        Tract probability maps in stereotaxic spaces: analyses of white
        matter anatomy and tract-specific quantification. Neuroimage 39:
        336-347
        """
        self.logger = logging.getLogger('AFQ.Segmentation')
        self.nb_points = nb_points

        if rng is None:
            self.rng = np.random.RandomState()
        else:
            self.rng = rng

        self.seg_algo = seg_algo.lower()
        self.prob_threshold = prob_threshold
        self.b0_threshold = b0_threshold
        self.progressive = progressive
        self.greater_than = greater_than
        self.rm_small_clusters = rm_small_clusters
        self.model_clust_thr = model_clust_thr
        self.reduction_thr = reduction_thr
        self.refine = refine
        self.pruning_thr = pruning_thr
        self.return_idx = return_idx
        self.filter_by_endpoints = filter_by_endpoints
        self.dist_to_aal = dist_to_aal

    def segment(self, bundle_dict, tg, fdata=None, fbval=None,
                fbvec=None, mapping=None, reg_prealign=None,
                reg_template=None, b0_threshold=0, img_affine=None):

        """
        Segment streamlines into bundles based on either waypoint ROIs
        [Yeatman2012]_ or RecoBundles [Garyfallidis2017]_.
        Parameters
        ----------
        bundle_dict: dict
            Meta-data for the segmentation. The format is something like::
                {'name': {'ROIs':[img1, img2],
                'rules':[True, True]},
                'prob_map': img3,
                'cross_midline': False}
        tg : StatefulTractogram
            Bundles to segment
        fdata, fbval, fbvec : str
            Full path to data, bvals, bvecs
        mapping : DiffeomorphicMap object, str or nib.Nifti1Image, optional.
            A mapping between DWI space and a template. If None, mapping will
            be registered from data used in prepare_img. Default: None.
        reg_prealign : array, optional.
            The linear transformation to be applied to align input images to
            the reference space before warping under the deformation field.
            Default: None.
        reg_template : str or nib.Nifti1Image, optional.
            Template to use for registration. Default: MNI T2.
        img_affine : array, optional.
            The spatial transformation from the measurement to the scanner
            space.
        References
        ----------
        .. [Yeatman2012] Yeatman, <NAME>., <NAME>, <NAME>.
        Myall, <NAME>, and <NAME>. 2012. "Tract Profiles of
        White Matter Properties: Automating Fiber-Tract Quantification"
        PloS One 7 (11): e49790.
        .. [Garyfallidis17] Garyfallidis et al. Recognition of white matter
        bundles using local and global streamline-based registration and
        clustering, Neuroimage, 2017.
        """
        if img_affine is not None:
            if (mapping is None
                or fdata is not None
                or fbval is not None
                    or fbvec is not None):

                self.logger.error(
                    "Provide either the full path to data, bvals, bvecs,"
                    + "or provide the affine of the image and the mapping")

        self.logger.info("Preparing Segmentation Parameters")
        self.img_affine = img_affine
        self.prepare_img(fdata, fbval, fbvec)
        self.logger.info("Preprocessing Streamlines")
        self.tg = tg

        # If resampling over-write the sft:
        if self.nb_points:
            self.tg = StatefulTractogram(
                dps.set_number_of_points(self.tg.streamlines, self.nb_points),
                self.tg, self.tg.space)

            self.resample_streamlines(self.nb_points)

        self.prepare_map(mapping, reg_prealign, reg_template)
        self.bundle_dict = bundle_dict
        self.cross_streamlines()

        if self.seg_algo == "afq":
            return self.segment_afq()
        elif self.seg_algo == "reco":
            return self.segment_reco()

    def prepare_img(self, fdata, fbval, fbvec):
        """
        Prepare image data from DWI data.
        Parameters
        ----------
        fdata, fbval, fbvec : str
            Full path to data, bvals, bvecs
        """
        if self.img_affine is None:
            self.img, _, _, _ = \
                ut.prepare_data(fdata, fbval, fbvec,
                                b0_threshold=self.b0_threshold)
            self.img_affine = self.img.affine

        self.fdata = fdata
        self.fbval = fbval
        self.fbvec = fbvec

    def prepare_map(self, mapping=None, reg_prealign=None, reg_template=None):
        """
        Set mapping between DWI space and a template.
        Parameters
        ----------
        mapping : DiffeomorphicMap object, str or nib.Nifti1Image, optional.
            A mapping between DWI space and a template.
            If None, mapping will be registered from data used in prepare_img.
            Default: None.
        reg_template : str or nib.Nifti1Image, optional.
            Template to use for registration (defaults to the MNI T2)
            Default: None.
        reg_prealign : array, optional.
            The linear transformation to be applied to align input images to
            the reference space before warping under the deformation field.
            Default: None.
        """
        if reg_template is None:
            reg_template = dpd.read_mni_template()

        if mapping is None:
            gtab = dpg.gradient_table(self.fbval, self.fbvec)
            self.mapping = reg.syn_register_dwi(self.fdata, gtab)[1]
        elif isinstance(mapping, str) or isinstance(mapping, nib.Nifti1Image):
            if reg_prealign is None:
                reg_prealign = np.eye(4)
            if self.img is None:
                self.img, _, _, _ = \
                    ut.prepare_data(self.fdata,
                                    self.fbval,
                                    self.fbvec,
                                    b0_threshold=self.b0_threshold)
            self.mapping = reg.read_mapping(mapping, self.img, reg_template,
                                            prealign=np.linalg.inv(reg_prealign))
            # replaced the pre_align with its inverse _aNNe
            # that fixed the warped rois and the warped prob maps which were wrong otherwise
        else:
            self.mapping = mapping

    def cross_streamlines(self, tg=None, template=None, low_coord=10):
        """
        Classify the streamlines by whether they cross the midline.
        Creates a crosses attribute which is an array of booleans. Each boolean
        corresponds to a streamline, and is whether or not that streamline
        crosses the midline.
        Parameters
        ----------
        tg : StatefulTractogram class instance.
        template : nibabel.Nifti1Image class instance
            An affine transformation into a template space.
        """
        if tg is None:
            tg = self.tg
        if template is None:
            template_affine = self.img_affine
        else:
            template_affine = template.affine

        # What is the x,y,z coordinate of 0,0,0 in the template space?
        zero_coord = np.dot(np.linalg.inv(template_affine),
                            np.array([0, 0, 0, 1]))

        self.crosses = np.zeros(len(tg), dtype=bool)
        # already_split = 0
        for sl_idx, sl in enumerate(tg.streamlines):
            if np.any(sl[:, 0] > zero_coord[0]) and \
                    np.any(sl[:, 0] < zero_coord[0]):
                self.crosses[sl_idx] = True
            else:
                self.crosses[sl_idx] = False

    def _get_bundle_info(self, bundle_idx, bundle):
        """
        Get fiber probabilites and ROIs for a given bundle.
        """
        rules = self.bundle_dict[bundle]['rules']
        include_rois = []
        exclude_rois = []
        for rule_idx, rule in enumerate(rules):
            roi = self.bundle_dict[bundle]['ROIs'][rule_idx]
            if not isinstance(roi, np.ndarray):
                roi = roi.get_fdata()
            warped_roi = auv.patch_up_roi(self.mapping.transform_inverse(
                roi.astype(np.float32), interpolation='linear'))
 
            if rule:
                # include ROI:
                include_rois.append(np.array(np.where(warped_roi)).T)
            else:
                # Exclude ROI:
                exclude_rois.append(np.array(np.where(warped_roi)).T)
########for debugging: save the warped ROI that is actually used in segment_afq() ##########
#       path_for_debugging = '/debugpath/'
#            nib.save(nib.Nifti1Image(warped_roi.astype(np.float32),
#                                     self.img_affine),
#                     debugpath+'warpedROI_'+bundle+'as_used.nii.gz')
############################################################################################

        # The probability map if doesn't exist is all ones with the same
        # shape as the ROIs:
        prob_map = self.bundle_dict[bundle].get(
            'prob_map', np.ones(roi.shape))

        if not isinstance(prob_map, np.ndarray):
            prob_map = prob_map.get_fdata()
        warped_prob_map = \
            self.mapping.transform_inverse(prob_map,
                                           interpolation='nearest')
        return warped_prob_map, include_rois, exclude_rois

    def _check_sl_with_inclusion(self, sl, include_rois, tol):
        """
        Helper function to check that a streamline is close to a list of
        inclusion ROIS.
        """
        dist = []
        for roi in include_rois:
            # Use squared Euclidean distance, because it's faster:
            dist.append(cdist(sl, roi, 'sqeuclidean'))
            if np.min(dist[-1]) > tol:
                # Too far from one of them:
                return False, []
        # Apparently you checked all the ROIs and it was close to all of them
        return True, dist

    def _check_sl_with_exclusion(self, sl, exclude_rois, tol):
        """ Helper function to check that a streamline is not too close to a
        list of exclusion ROIs.
        """
        for roi in exclude_rois:
            # Use squared Euclidean distance, because it's faster:
            if np.min(cdist(sl, roi, 'sqeuclidean')) < tol:
                return False
        # Either there are no exclusion ROIs, or you are not close to any:
        return True

    def _return_empty(self, bundle):
        """
        Helper function for segment_afq, to return an empty dict under
        some conditions.
        """

        if self.return_idx:
            self.fiber_groups[bundle] = {}
            self.fiber_groups[bundle]['sl'] = StatefulTractogram(
                [], self.img, Space.VOX)
            self.fiber_groups[bundle]['idx'] = np.array([])
        else:
            self.fiber_groups[bundle] = StatefulTractogram(
                [], self.img, Space.VOX)

    def segment_afq(self, tg=None):
        """
        Assign streamlines to bundles using the waypoint ROI approach
        Parameters
        ----------
        tg : StatefulTractogram class instance
        """
        if tg is None:
            tg = self.tg
        else:
            self.tg = tg

        self.tg.to_vox()
        # For expedience, we approximate each streamline as a 100 point curve:
        fgarray = np.array(_resample_tg(tg, 100))
        
        # comment _aNNe
        # in general, this might cause errors: 
        # if rois were traversed by streamlines in just a few voxels
        # and if streamlines were so long or resolution so high 
        # that one hundredth of a streamline length was more than a voxel,
        # then the contact check below (closest distance streamline to ROI < voxel width) can fail when resampling to 100 points
        # To be cartain that the resampling does not cause problems, 
        # the number of resamplign points has to be larger than the length of the streamline in voxels in native space! 
        # end comment

        n_streamlines = fgarray.shape[0]

        streamlines_in_bundles = np.zeros(
            (n_streamlines, len(self.bundle_dict)))
        min_dist_coords = np.zeros(
            (n_streamlines, len(self.bundle_dict), 2), dtype=int)
        self.fiber_groups = {}

        if self.return_idx:
            out_idx = np.arange(n_streamlines, dtype=int)

        if self.filter_by_endpoints:
            aal_atlas = afd.read_aal_atlas()['atlas'].get_fdata()
            # We need to calculate the size of a voxel, so we can transform
            # from mm to voxel units:
            R = self.img_affine[0:3, 0:3]
            vox_dim = np.mean(np.diag(np.linalg.cholesky(R.T.dot(R))))
            dist_to_aal = self.dist_to_aal / vox_dim

        self.logger.info("Assigning Streamlines to Bundles")
        # Tolerance is set to the square of the distance to the corner
        # because we are using the squared Euclidean distance in calls to
        # `cdist` to make those calls faster.
        tol = dts.dist_to_corner(self.img_affine)**2
        for bundle_idx, bundle in enumerate(self.bundle_dict):
            self.logger.info(f"Finding Streamlines for {bundle}")
            warped_prob_map, include_roi, exclude_roi = \
                self._get_bundle_info(bundle_idx, bundle)
########for debugging: save the warped probability map that is actually used in segment_afq() ##########
#            path_for_debugging = '/debugpath/'
#            nib.save(nib.Nifti1Image(warped_prob_map.astype(np.float32),
#                                     self.img_affine),
#                      debugpath+'warpedprobmap_'+bundle+'as_used.nii.gz')
############################################################################################
            fiber_probabilities = dts.values_from_volume(
                warped_prob_map,<commit_msg>Adds a flag to save intermediate variables within the segmentation object<commit_after>import os.path as op
import logging

import numpy as np
from scipy.spatial.distance import cdist
from scipy.stats import zscore

import nibabel as nib
from tqdm.auto import tqdm

import dipy.data as dpd
import dipy.tracking.streamline as dts
import dipy.tracking.streamlinespeed as dps
from dipy.segment.bundles import RecoBundles
from dipy.align.streamlinear import whole_brain_slr
from dipy.stats.analysis import gaussian_weights
import dipy.core.gradients as dpg
from dipy.io.stateful_tractogram import StatefulTractogram, Space

import AFQ.registration as reg
import AFQ.utils.models as ut
import AFQ.utils.volume as auv
import AFQ.data as afd

__all__ = ["Segmentation"]


def _resample_tg(tg, n_points):
    # reformat for dipy's set_number_of_points
    if isinstance(tg, np.ndarray):
        if len(tg.shape) > 2:
            streamlines = tg.tolist()
            streamlines = [np.asarray(item) for item in streamlines]
    else:
        streamlines = tg.streamlines

    return dps.set_number_of_points(streamlines, n_points)


class Segmentation:
    def __init__(self,
                 nb_points=False,
                 seg_algo='AFQ',
                 progressive=True,
                 greater_than=50,
                 rm_small_clusters=50,
                 model_clust_thr=40,
                 reduction_thr=40,
                 refine=False,
                 pruning_thr=6,
                 b0_threshold=0,
                 prob_threshold=0,
                 rng=None,
                 return_idx=False,
                 filter_by_endpoints=True,
                 dist_to_aal=4,
                 save_intermediates=False):
        """
        Segment streamlines into bundles.
        Parameters
        ----------
        nb_points : int, boolean
            Resample streamlines to nb_points number of points.
            If False, no resampling is done. Default: False
        seg_algo : string
            Algorithm for segmentation (case-insensitive):
            'AFQ': Segment streamlines into bundles,
                based on inclusion/exclusion ROIs.
            'Reco': Segment streamlines using the RecoBundles algorithm
            [Garyfallidis2017].
            Default: 'AFQ'
        rm_small_clusters : int
            Using RecoBundles Algorithm.
            Remove clusters that have less than this value
                during whole brain SLR.
            Default: 50
        progressive : boolean, optional
            Using RecoBundles Algorithm.
            Whether or not to use progressive technique
                during whole brain SLR.
            Default: True.
        greater_than : int, optional
            Using RecoBundles Algorithm.
            Keep streamlines that have length greater than this value
                during whole brain SLR.
            Default: 50.
        b0_theshold : float.
            Using AFQ Algorithm.
            All b-values with values less than or equal to `bo_threshold` are
            considered as b0s i.e. without diffusion weighting.
            Default: 0.
        prob_threshold : float.
            Using AFQ Algorithm.
            Initial cleaning of fiber groups is done using probability maps
            from [Hua2008]_. Here, we choose an average probability that
            needs to be exceeded for an individual streamline to be retained.
            Default: 0.
        rng : RandomState
            If None, creates RandomState. Used in RecoBundles Algorithm.
            Default: None.
        return_idx : bool
            Whether to return the indices in the original streamlines as part
            of the output of segmentation.
        filter_by_endpoints: bool
            Whether to filter the bundles based on their endpoints relative
            to regions defined in the AAL atlas. Applies only to the waypoint
            approach (XXX for now). Default: True.
        dist_to_aal : float
            If filter_by_endpoints is True, this is the distance from the
            endpoints to the AAL atlas ROIs that is required.
        save_intermediates : str, optional
            The full path to a folder into which intermediate products
            are saved. Default: None, means no saving of intermediates.

        References
        ----------
        .. [Hua2008] <NAME>, <NAME>, <NAME>, <NAME>, <NAME>, et al. (2008)
        Tract probability maps in stereotaxic spaces: analyses of white
        matter anatomy and tract-specific quantification. Neuroimage 39:
        336-347
        """
        self.logger = logging.getLogger('AFQ.Segmentation')
        self.nb_points = nb_points

        if rng is None:
            self.rng = np.random.RandomState()
        else:
            self.rng = rng

        self.seg_algo = seg_algo.lower()
        self.prob_threshold = prob_threshold
        self.b0_threshold = b0_threshold
        self.progressive = progressive
        self.greater_than = greater_than
        self.rm_small_clusters = rm_small_clusters
        self.model_clust_thr = model_clust_thr
        self.reduction_thr = reduction_thr
        self.refine = refine
        self.pruning_thr = pruning_thr
        self.return_idx = return_idx
        self.filter_by_endpoints = filter_by_endpoints
        self.dist_to_aal = dist_to_aal
        self.save_intermediates = save_intermediates

    def segment(self, bundle_dict, tg, fdata=None, fbval=None,
                fbvec=None, mapping=None, reg_prealign=None,
                reg_template=None, b0_threshold=0, img_affine=None):

        """
        Segment streamlines into bundles based on either waypoint ROIs
        [Yeatman2012]_ or RecoBundles [Garyfallidis2017]_.
        Parameters
        ----------
        bundle_dict: dict
            Meta-data for the segmentation. The format is something like::
                {'name': {'ROIs':[img1, img2],
                'rules':[True, True]},
                'prob_map': img3,
                'cross_midline': False}
        tg : StatefulTractogram
            Bundles to segment
        fdata, fbval, fbvec : str
            Full path to data, bvals, bvecs
        mapping : DiffeomorphicMap object, str or nib.Nifti1Image, optional.
            A mapping between DWI space and a template. If None, mapping will
            be registered from data used in prepare_img. Default: None.
        reg_prealign : array, optional.
            The linear transformation to be applied to align input images to
            the reference space before warping under the deformation field.
            Default: None.
        reg_template : str or nib.Nifti1Image, optional.
            Template to use for registration. Default: MNI T2.
        img_affine : array, optional.
            The spatial transformation from the measurement to the scanner
            space.
        References
        ----------
        .. [Yeatman2012] <NAME>., <NAME>, N<NAME>.
        Myall, <NAME>, and <NAME>. 2012. "Tract Profiles of
        White Matter Properties: Automating Fiber-Tract Quantification"
        PloS One 7 (11): e49790.
        .. [Garyfallidis17] Garyfallidis et al. Recognition of white matter
        bundles using local and global streamline-based registration and
        clustering, Neuroimage, 2017.
        """
        if img_affine is not None:
            if (mapping is None
                or fdata is not None
                or fbval is not None
                    or fbvec is not None):

                self.logger.error(
                    "Provide either the full path to data, bvals, bvecs,"
                    + "or provide the affine of the image and the mapping")

        self.logger.info("Preparing Segmentation Parameters")
        self.img_affine = img_affine
        self.prepare_img(fdata, fbval, fbvec)
        self.logger.info("Preprocessing Streamlines")
        self.tg = tg

        # If resampling over-write the sft:
        if self.nb_points:
            self.tg = StatefulTractogram(
                dps.set_number_of_points(self.tg.streamlines, self.nb_points),
                self.tg, self.tg.space)

            self.resample_streamlines(self.nb_points)

        self.prepare_map(mapping, reg_prealign, reg_template)
        self.bundle_dict = bundle_dict
        self.cross_streamlines()

        if self.seg_algo == "afq":
            return self.segment_afq()
        elif self.seg_algo == "reco":
            return self.segment_reco()

    def prepare_img(self, fdata, fbval, fbvec):
        """
        Prepare image data from DWI data.
        Parameters
        ----------
        fdata, fbval, fbvec : str
            Full path to data, bvals, bvecs
        """
        if self.img_affine is None:
            self.img, _, _, _ = \
                ut.prepare_data(fdata, fbval, fbvec,
                                b0_threshold=self.b0_threshold)
            self.img_affine = self.img.affine

        self.fdata = fdata
        self.fbval = fbval
        self.fbvec = fbvec

    def prepare_map(self, mapping=None, reg_prealign=None, reg_template=None):
        """
        Set mapping between DWI space and a template.
        Parameters
        ----------
        mapping : DiffeomorphicMap object, str or nib.Nifti1Image, optional.
            A mapping between DWI space and a template.
            If None, mapping will be registered from data used in prepare_img.
            Default: None.
        reg_template : str or nib.Nifti1Image, optional.
            Template to use for registration (defaults to the MNI T2)
            Default: None.
        reg_prealign : array, optional.
            The linear transformation to be applied to align input images to
            the reference space before warping under the deformation field.
            Default: None.
        """
        if reg_template is None:
            reg_template = dpd.read_mni_template()

        if mapping is None:
            gtab = dpg.gradient_table(self.fbval, self.fbvec)
            self.mapping = reg.syn_register_dwi(self.fdata, gtab)[1]
        elif isinstance(mapping, str) or isinstance(mapping, nib.Nifti1Image):
            if reg_prealign is None:
                reg_prealign = np.eye(4)
            if self.img is None:
                self.img, _, _, _ = \
                    ut.prepare_data(self.fdata,
                                    self.fbval,
                                    self.fbvec,
                                    b0_threshold=self.b0_threshold)
            self.mapping = reg.read_mapping(
                mapping,
                self.img,
                reg_template,
                prealign=np.linalg.inv(reg_prealign))
        else:
            self.mapping = mapping

    def cross_streamlines(self, tg=None, template=None, low_coord=10):
        """
        Classify the streamlines by whether they cross the midline.
        Creates a crosses attribute which is an array of booleans. Each boolean
        corresponds to a streamline, and is whether or not that streamline
        crosses the midline.
        Parameters
        ----------
        tg : StatefulTractogram class instance.
        template : nibabel.Nifti1Image class instance
            An affine transformation into a template space.
        """
        if tg is None:
            tg = self.tg
        if template is None:
            template_affine = self.img_affine
        else:
            template_affine = template.affine

        # What is the x,y,z coordinate of 0,0,0 in the template space?
        zero_coord = np.dot(np.linalg.inv(template_affine),
                            np.array([0, 0, 0, 1]))

        self.crosses = np.zeros(len(tg), dtype=bool)
        # already_split = 0
        for sl_idx, sl in enumerate(tg.streamlines):
            if np.any(sl[:, 0] > zero_coord[0]) and \
                    np.any(sl[:, 0] < zero_coord[0]):
                self.crosses[sl_idx] = True
            else:
                self.crosses[sl_idx] = False

    def _get_bundle_info(self, bundle_idx, bundle):
        """
        Get fiber probabilites and ROIs for a given bundle.
        """
        rules = self.bundle_dict[bundle]['rules']
        include_rois = []
        exclude_rois = []
        for rule_idx, rule in enumerate(rules):
            roi = self.bundle_dict[bundle]['ROIs'][rule_idx]
            if not isinstance(roi, np.ndarray):
                roi = roi.get_fdata()
            warped_roi = auv.patch_up_roi(self.mapping.transform_inverse(
                roi.astype(np.float32), interpolation='linear'))

            if rule:
                # include ROI:
                include_rois.append(np.array(np.where(warped_roi)).T)
            else:
                # Exclude ROI:
                exclude_rois.append(np.array(np.where(warped_roi)).T)

            # For debugging purposes, we can save the variable as it is:
            if self.save_intermediates is not None:
                nib.save(
                    nib.Nifti1Image(warped_roi.astype(np.float32),
                                    self.img_affine),
                    op.join(self.save_intermediates,
                            'warpedROI_',
                            bundle,
                            'as_used.nii.gz'))

        # The probability map if doesn't exist is all ones with the same
        # shape as the ROIs:
        prob_map = self.bundle_dict[bundle].get(
            'prob_map', np.ones(roi.shape))

        if not isinstance(prob_map, np.ndarray):
            prob_map = prob_map.get_fdata()
        warped_prob_map = \
            self.mapping.transform_inverse(prob_map,
                                           interpolation='nearest')
        return warped_prob_map, include_rois, exclude_rois

    def _check_sl_with_inclusion(self, sl, include_rois, tol):
        """
        Helper function to check that a streamline is close to a list of
        inclusion ROIS.
        """
        dist = []
        for roi in include_rois:
            # Use squared Euclidean distance, because it's faster:
            dist.append(cdist(sl, roi, 'sqeuclidean'))
            if np.min(dist[-1]) > tol:
                # Too far from one of them:
                return False, []
        # Apparently you checked all the ROIs and it was close to all of them
        return True, dist

    def _check_sl_with_exclusion(self, sl, exclude_rois, tol):
        """ Helper function to check that a streamline is not too close to a
        list of exclusion ROIs.
        """
        for roi in exclude_rois:
            # Use squared Euclidean distance, because it's faster:
            if np.min(cdist(sl, roi, 'sqeuclidean')) < tol:
                return False
        # Either there are no exclusion ROIs, or you are not close to any:
        return True

    def _return_empty(self, bundle):
        """
        Helper function for segment_afq, to return an empty dict under
        some conditions.
        """

        if self.return_idx:
            self.fiber_groups[bundle] = {}
            self.fiber_groups[bundle]['sl'] = StatefulTractogram(
                [], self.img, Space.VOX)
            self.fiber_groups[bundle]['idx'] = np.array([])
        else:
            self.fiber_groups[bundle] = StatefulTractogram(
                [], self.img, Space.VOX)

    def segment_afq(self, tg=None):
        """
        Assign streamlines to bundles using the waypoint ROI approach
        Parameters
        ----------
        tg : StatefulTractogram class instance
        """
        if tg is None:
            tg = self.tg
        else:
            self.tg = tg

        self.tg.to_vox()

        # For expedience, we approximate each streamline as a 100 point curve.
        # This is only used in extracting the values from the probability map,
        # so will not affect measurement of distance from the waypoint ROIs
        fgarray = np.array(_resample_tg(tg, 100))
        n_streamlines = fgarray.shape[0]

        streamlines_in_bundles = np.zeros(
            (n_streamlines, len(self.bundle_dict)))
        min_dist_coords = np.zeros(
            (n_streamlines, len(self.bundle_dict), 2), dtype=int)
        self.fiber_groups = {}

        if self.return_idx:
            out_idx = np.arange(n_streamlines, dtype=int)

        if self.filter_by_endpoints:
            aal_atlas = afd.read_aal_atlas()['atlas'].get_fdata()
            # We need to calculate the size of a voxel, so we can transform
            # from mm to voxel units:
            R = self.img_affine[0:3, 0:3]
            vox_dim = np.mean(np.diag(np.linalg.cholesky(R.T.dot(R))))
            dist_to_aal = self.dist_to_aal / vox_dim

        self.logger.info("Assigning Streamlines to Bundles")
        # Tolerance is set to the square of the distance to the corner
        # because we are using the squared Euclidean distance in calls to
        # `cdist` to make those calls faster.
        tol = dts.dist_to_corner(self.img_affine)**2
        for bundle_idx, bundle in enumerate(self.bundle_dict):
            self.logger.info(f"Finding Streamlines for {bundle}")
            warped_prob_map, include_roi, exclude_roi = \
                self._get_bundle_info(bundle_idx, bundle)
            if self.save_intermediates is not None:
                nib.save(
                    nib.Nifti1Image(warped_prob_map.astype(np.float32),
                                    self.img_affine),
                    op.join(self.save_intermediates,
                            'warpedprobmap',
                            bundle,
                            'as_used.nii.gz'))

            fiber_probabilities = dts.values_from_volume(
                warped_prob_map,
{'fraction_non_alphanumeric': 0.0809723386420788, 'fraction_numerical': 0.007516065940206762, 'mean_word_length': 1.8159716758457907, 'pattern_counts': {'":': 4, '<': 28, '<?xml version=': 0, '>': 29, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '723337', 'n_tokens_mistral': 10286, 'n_tokens_neox': 9738, 'n_words': 2941}