# Source code for dipy.workflows.tracking

    #!/usr/bin/env python3

    import logging
    import sys

    from dipy.data import SPHERE_FILES, get_sphere
    from dipy.io.image import load_nifti
    from dipy.io.peaks import load_pam
    from dipy.io.stateful_tractogram import Space, StatefulTractogram
    from dipy.io.streamline import save_tractogram
    from dipy.tracking import utils
    from dipy.tracking.stopping_criterion import (
        BinaryStoppingCriterion,
        CmcStoppingCriterion,
        ThresholdStoppingCriterion,
    )
    from dipy.tracking.tracker import (
        closestpeak_tracking,
        deterministic_tracking,
        eudx_tracking,
        pft_tracking,
        probabilistic_tracking,
        ptt_tracking,
    )
    from dipy.workflows.workflow import Workflow

    class LocalFiberTrackingPAMFlow(Workflow):

        @classmethod
        def get_short_name(cls):
            return "track_local"

        def run(
            self,
            pam_files,
            stopping_files,
            seeding_files,
            use_binary_mask=False,
            stopping_thr=0.2,
            seed_density=1,
            minlen=2,
            maxlen=500,
            step_size=0.5,
            tracking_method="deterministic",
            pmf_threshold=0.1,
            max_angle=30.0,
            sphere_name=None,
            save_seeds=False,
            nbr_threads=0,
            random_seed=1,
            seed_buffer_fraction=1.0,
            out_dir="",
            out_tractogram="tractogram.trk",
        ):
            """Workflow for Local Fiber Tracking.

            This workflow use a saved peaks and metrics (PAM) file as input.

            See :footcite:p:`Garyfallidis2012b` and :footcite:p:`Amirbekian2016`
            for further details about the method.

            Parameters
            ----------
            pam_files : string
               Path to the peaks and metrics files. This path may contain
                wildcards to use multiple masks at once.
            stopping_files : string
                Path to images (e.g. FA) used for stopping criterion for tracking.
            seeding_files : string
                A binary image showing where we need to seed for tracking.
            use_binary_mask : bool, optional
                If True, uses a binary stopping criterion. If the provided
                `stopping_files` are not binary, `stopping_thr` will be used to
                binarize the images.
            stopping_thr : float, optional
                Threshold applied to stopping volume's data to identify where
                tracking has to stop.
            seed_density : int, optional
                Number of seeds per dimension inside voxel.
                For example, seed_density of 2 means 8 regularly distributed
                points in the voxel. And seed density of 1 means 1 point at the
                center of the voxel.
            minlen : int, optional
                Minimum length (nb points) of the streamlines.
            maxlen : int, optional
                Maximum length (nb points) of the streamlines.
            step_size : float, optional
                Step size (in mm) used for tracking.
            tracking_method : string, optional
                Select direction getter strategy:
                    - "eudx" (Uses the peaks saved in the pam_files)
                    - "deterministic" or "det" for a deterministic tracking
                    - "probabilistic" or "prob" for a Probabilistic tracking
                    - "closestpeaks" or "cp" for a ClosestPeaks tracking
                    - "ptt" for Parallel Transport Tractography

                By default, the sh coefficients saved in the pam_files are used.
            pmf_threshold : float, optional
                Threshold for ODF functions.
            max_angle : float, optional
                Maximum angle between streamline segments (range [0, 90]).
            sphere_name : string, optional
                The sphere used for tracking. If None, the sphere saved in the
                pam_files is used. For faster tracking, use a smaller
                sphere (e.g. 'repulsion200').
            save_seeds : bool, optional
                If true, save the seeds associated to their streamline
                in the 'data_per_streamline' Tractogram dictionary using
                'seeds' as the key.
            nbr_threads : int, optional
                Number of threads to use for the processing. By default, all available
                threads will be used.
            random_seed : int, optional
                Seed for the random number generator, must be >= 0. A value of greater
                than 0 will all produce the same streamline trajectory for a given seed
                coordinate. A value of 0 may produces various streamline tracjectories
                for a given seed coordinate.
            seed_buffer_fraction : float, optional
                Fraction of the seed buffer to use. A value of 1.0 will use the entire seed
                buffer. A value of 0.5 will use half of the seed buffer then the other half.
                a way to reduce memory usage.
            out_dir : string, optional
               Output directory.
            out_tractogram : string, optional
               Name of the tractogram file to be saved.

            References
            ----------
            .. footbibliography::

            """
            io_it = self.get_io_iterator()

            tracking_method = tracking_method.lower()
            sphere_name = SPHERE_FILES.get(sphere_name, None)

            for pams_path, stopping_path, seeding_path, out_tract in io_it:
                if tracking_method == "eudx":
                    logging.info(f"EuDX Deterministic tracking on {pams_path}")
                else:
                    logging.info(f"{tracking_method.title()} tracking on {pams_path}")

                pam = load_pam(pams_path, verbose=False)
                sphere = pam.sphere if sphere_name is None else get_sphere(name=sphere_name)

                logging.info("Loading stopping criterion")
                stop, affine = load_nifti(stopping_path)
                if use_binary_mask:
                    stopping_criterion = BinaryStoppingCriterion(stop > stopping_thr)
                else:
                    stopping_criterion = ThresholdStoppingCriterion(stop, stopping_thr)

                logging.info("Loading seeds")
                seed_mask, _ = load_nifti(seeding_path)
                seeds = utils.seeds_from_mask(
                    seed_mask, affine, density=[seed_density, seed_density, seed_density]
                )

                logging.info("Starting to track")
                if tracking_method in ["closestpeaks", "cp"]:
                    tracking_result = closestpeak_tracking(
                        seeds,
                        stopping_criterion,
                        affine,
                        sh=pam.shm_coeff,
                        random_seed=random_seed,
                        sphere=sphere,
                        max_angle=max_angle,
                        min_len=minlen,
                        max_len=maxlen,
                        step_size=step_size,
                        pmf_threshold=pmf_threshold,
                        save_seeds=save_seeds,
                        nbr_threads=nbr_threads,
                        seed_buffer_fraction=seed_buffer_fraction,
                    )
                elif tracking_method in [
                    "eudx",
                ]:
                    tracking_result = eudx_tracking(
                        seeds,
                        stopping_criterion,
                        affine,
                        sh=pam.shm_coeff,
                        pam=pam,
                        random_seed=random_seed,
                        sphere=sphere,
                        max_angle=max_angle,
                        min_len=minlen,
                        max_len=maxlen,
                        step_size=step_size,
                        pmf_threshold=pmf_threshold,
                        save_seeds=save_seeds,
                        nbr_threads=nbr_threads,
                        seed_buffer_fraction=seed_buffer_fraction,
                    )
                elif tracking_method in ["deterministic", "det"]:
                    tracking_result = deterministic_tracking(
                        seeds,
                        stopping_criterion,
                        affine,
                        sh=pam.shm_coeff,
                        random_seed=random_seed,
                        sphere=sphere,
                        max_angle=max_angle,
                        min_len=minlen,
                        max_len=maxlen,
                        step_size=step_size,
                        pmf_threshold=pmf_threshold,
                        save_seeds=save_seeds,
                        nbr_threads=nbr_threads,
                        seed_buffer_fraction=seed_buffer_fraction,
                    )
                elif tracking_method in ["probabilistic", "prob"]:
                    tracking_result = probabilistic_tracking(
                        seeds,
                        stopping_criterion,
                        affine,
                        sh=pam.shm_coeff,
                        random_seed=random_seed,
                        sphere=sphere,
                        max_angle=max_angle,
                        min_len=minlen,
                        max_len=maxlen,
                        step_size=step_size,
                        pmf_threshold=pmf_threshold,
                        save_seeds=save_seeds,
                        nbr_threads=nbr_threads,
                        seed_buffer_fraction=seed_buffer_fraction,
                    )
                elif tracking_method in ["ptt"]:
                    tracking_result = ptt_tracking(
                        seeds,
                        stopping_criterion,
                        affine,
                        sh=pam.shm_coeff,
                        random_seed=random_seed,
                        sphere=sphere,
                        max_angle=max_angle,
                        min_len=minlen,
                        max_len=maxlen,
                        step_size=step_size,
                        pmf_threshold=pmf_threshold,
                        save_seeds=save_seeds,
                        nbr_threads=nbr_threads,
                        seed_buffer_fraction=seed_buffer_fraction,
                    )
                else:
                    logging.error(
                        f"Unknown tracking method: {tracking_method}. "
                        f"Please use one of the following: "
                        f"'eudx', 'deterministic', 'probabilistic', 'closestpeaks', 'ptt'"
                    )
                    sys.exit(1)
                if save_seeds:
                    streamlines, seeds = zip(*tracking_result)
                    seeds = {"seeds": seeds}
                else:
                    streamlines = list(tracking_result)
                    seeds = {}

                sft = StatefulTractogram(
                    streamlines, seeding_path, Space.RASMM, data_per_streamline=seeds
                )
                save_tractogram(sft, out_tract, bbox_valid_check=False)
                logging.info(f"Saved {out_tract}")

    class PFTrackingPAMFlow(Workflow):

        @classmethod
        def get_short_name(cls):
            return "track_pft"

        def run(
            self,
            pam_files,
            wm_files,
            gm_files,
            csf_files,
            seeding_files,
            step_size=0.2,
            seed_density=1,
            pmf_threshold=0.1,
            max_angle=20.0,
            sphere_name=None,
            pft_back=2,
            pft_front=1,
            pft_count=15,
            pft_max_trial=20,
            save_seeds=False,
            min_wm_pve_before_stopping=0,
            nbr_threads=0,
            random_seed=1,
            seed_buffer_fraction=1.0,
            out_dir="",
            out_tractogram="tractogram.trk",
        ):
            """Workflow for Particle Filtering Tracking.

            This workflow uses a saved peaks and metrics (PAM) file as input.

            See :footcite:p:`Girard2014` for further details about the method.

            Parameters
            ----------
            pam_files : string
               Path to the peaks and metrics files. This path may contain
                wildcards to use multiple masks at once.
            wm_files : string
                Path to white matter partial volume estimate for tracking (CMC).
            gm_files : string
                Path to grey matter partial volume estimate for tracking (CMC).
            csf_files : string
                Path to cerebrospinal fluid partial volume estimate for tracking
                (CMC).
            seeding_files : string
                A binary image showing where we need to seed for tracking.
            step_size : float, optional
                Step size (in mm) used for tracking.
            seed_density : int, optional
                Number of seeds per dimension inside voxel.
                For example, seed_density of 2 means 8 regularly distributed
                points in the voxel. And seed density of 1 means 1 point at the
                center of the voxel.
            pmf_threshold : float, optional
                Threshold for ODF functions.
            max_angle : float, optional
                Maximum angle between streamline segments (range [0, 90]).
            sphere_name : string, optional
                The sphere used for tracking. If None, the sphere saved in the
                pam_files is used. For faster tracking, use a smaller
                sphere (e.g. 'repulsion200').
            pft_back : float, optional
                Distance in mm to back track before starting the particle filtering
                tractography. The total particle filtering
                tractography distance is equal to back_tracking_dist +
                front_tracking_dist.
            pft_front : float, optional
                Distance in mm to run the particle filtering tractography after the
                the back track distance. The total particle filtering
                tractography distance is equal to back_tracking_dist +
                front_tracking_dist.
            pft_count : int, optional
                Number of particles to use in the particle filter.
            pft_max_trial : int, optional
                Maximum number of trials to run the particle filtering tractography.
            save_seeds : bool, optional
                If true, save the seeds associated to their streamline
                in the 'data_per_streamline' Tractogram dictionary using
                'seeds' as the key.
            min_wm_pve_before_stopping : int, optional
                Minimum white matter pve (1 - stopping_criterion.include_map -
                stopping_criterion.exclude_map) to reach before allowing the
                tractography to stop.
            nbr_threads : int, optional
                Number of threads to use for the processing. By default, all available
                threads will be used.
            random_seed : int, optional
                Seed for the random number generator, must be >= 0. A value of greater
                than 0 will all produce the same streamline trajectory for a given seed
                coordinate. A value of 0 may produces various streamline tracjectories
                for a given seed coordinate.
            seed_buffer_fraction : float, optional
                Fraction of the seed buffer to use. A value of 1.0 will use the entire seed
                buffer. A value of 0.5 will use half of the seed buffer then the other half.
                a way to reduce memory usage.
            out_dir : string, optional
               Output directory.
            out_tractogram : string, optional
               Name of the tractogram file to be saved.

            References
            ----------
            .. footbibliography::

            """
            io_it = self.get_io_iterator()
            sphere_name = SPHERE_FILES.get(sphere_name, None)

            for pams_path, wm_path, gm_path, csf_path, seeding_path, out_tract in io_it:
                logging.info(f"Particle Filtering tracking on {pams_path}")

                pam = load_pam(pams_path, verbose=False)
                sphere = pam.sphere if sphere_name is None else get_sphere(name=sphere_name)

                wm, affine, voxel_size = load_nifti(wm_path, return_voxsize=True)
                gm, _ = load_nifti(gm_path)
                csf, _ = load_nifti(csf_path)
                avs = sum(voxel_size) / len(voxel_size)  # average_voxel_size

                logging.info("Preparing stopping criterion")
                stopping_criterion = CmcStoppingCriterion.from_pve(
                    wm, gm, csf, step_size=step_size, average_voxel_size=avs
                )

                logging.info("Seeding in mask")
                seed_mask, _ = load_nifti(seeding_path)
                seeds = utils.seeds_from_mask(
                    seed_mask, affine, density=[seed_density, seed_density, seed_density]
                )

                logging.info("Start tracking")
                tracking_result = pft_tracking(
                    seeds,
                    stopping_criterion,
                    affine,
                    sh=pam.shm_coeff,
                    sphere=sphere,
                    step_size=step_size,
                    pft_back_tracking_dist=pft_back,
                    pft_front_tracking_dist=pft_front,
                    pft_max_trial=pft_max_trial,
                    particle_count=pft_count,
                    save_seeds=save_seeds,
                    min_wm_pve_before_stopping=min_wm_pve_before_stopping,
                    random_seed=random_seed,
                    max_angle=max_angle,
                    pmf_threshold=pmf_threshold,
                    nbr_threads=nbr_threads,
                    seed_buffer_fraction=seed_buffer_fraction,
                )

                if save_seeds:
                    streamlines, seeds = zip(*tracking_result)
                    seeds = {"seeds": seeds}
                else:
                    streamlines = list(tracking_result)
                    seeds = {}

                sft = StatefulTractogram(
                    streamlines, seeding_path, Space.RASMM, data_per_streamline=seeds
                )
                save_tractogram(sft, out_tract, bbox_valid_check=False)
                logging.info(f"Saved {out_tract}")