# Source code for dipy.align.streamwarp

    import warnings

    import numpy as np
    from scipy.optimize import linear_sum_assignment

    from dipy.align.bundlemin import distance_matrix_mdf
    from dipy.align.cpd import DeformableRegistration
    from dipy.align.streamlinear import slr_with_qbx
    from dipy.segment.clustering import QuickBundles
    from dipy.segment.metricspeed import AveragePointwiseEuclideanMetric
    from dipy.stats.analysis import assignment_map
    from dipy.testing.decorators import warning_for_keywords
    from dipy.tracking.streamline import Streamlines, length, unlist_streamlines
    from dipy.utils.optpkg import optional_package
    from dipy.viz.plotting import bundle_shape_profile

    pd, have_pd, _ = optional_package("pandas")

    def average_bundle_length(bundle):
        """Find average Euclidean length of the bundle in mm.

        Parameters
        ----------
        bundle : Streamlines
            Bundle who's average length is to be calculated.

        Returns
        -------
        int
            Average Euclidean length of bundle in mm.

        """
        metric = AveragePointwiseEuclideanMetric()
        qb = QuickBundles(threshold=85.0, metric=metric)
        clusters = qb.cluster(bundle)
        centroids = Streamlines(clusters.centroids)

        return length(centroids)[0]

    def find_missing(lst, cb):
        """Find unmatched streamline indices in moving bundle.

        Parameters
        ----------
        lst : List
            List of integers containing all the streamlines indices in moving
            bundle.
        cb : List
            List of integers containing streamline indices of the moving bundle
            that were not matched to any streamline in static bundle.

        Returns
        -------
        list
            List containing unmatched streamlines from moving bundle

        """
        return [x for x in range(0, len(cb)) if x not in lst]

    @warning_for_keywords()
    def bundlewarp(
        static, moving, *, dist=None, alpha=0.5, beta=20, max_iter=15, affine=True
    ):
        """Register two bundles using nonlinear method.

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

        Parameters
        ----------
        static : Streamlines
            Reference/fixed bundle.

        moving : Streamlines
            Target bundle that will be moved/registered to match the static bundle.

        dist : float, optional
            Precomputed distance matrix.

        alpha : float, optional
            Represents the trade-off between regularizing the deformation and
            having points match very closely. Lower value of alpha means high
            deformations.

        beta : int, optional
            Represents the strength of the interaction between points
            Gaussian kernel size.

        max_iter : int, optional
            Maximum number of iterations for deformation process in ml-CPD method.

        affine : boolean, optional
            If False, use rigid registration as starting point.

        Returns
        -------
        deformed_bundle : Streamlines
            Nonlinearly moved bundle (warped bundle)

        moving_aligned : Streamlines
            Linearly moved bundle (affinely moved)

        dist : np.ndarray
            Float array containing distance between moving and static bundle

        matched_pairs : np.ndarray
            Int array containing streamline correspondences between two bundles

        warp : np.ndarray
            Nonlinear warp map generated by BundleWarp

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

        """
        if alpha <= 0.01:
            warnings.warn(
                "Using alpha<=0.01 will result in extreme deformations", stacklevel=2
            )

        if average_bundle_length(static) <= 50:
            beta = 10

        x0 = "affine" if affine else "rigid"
        moving_aligned, _, _, _ = slr_with_qbx(static, moving, x0=x0, rm_small_clusters=0)

        if dist is not None:
            print("using pre-computed distances")
        else:
            dist = distance_matrix_mdf(static, moving_aligned).T

        matched_pairs = np.zeros((len(moving), 2))
        matched_pairs1 = np.asarray(linear_sum_assignment(dist)).T

        for mt in matched_pairs1:
            matched_pairs[mt[0]] = mt

        num = len(matched_pairs1)

        all_pairs = list(matched_pairs1[:, 0])
        all_matched = False

        while all_matched is False:
            num = len(all_pairs)

            if num < len(moving):
                ml = find_missing(all_pairs, moving)

                dist2 = dist[:][ml]

                # dist2 has distance among unmatched streamlines of moving bundle
                # and all static bundle's streamlines

                matched_pairs2 = np.asarray(linear_sum_assignment(dist2)).T

                for i in range(matched_pairs2.shape[0]):
                    matched_pairs2[i][0] = ml[matched_pairs2[i][0]]

                for mt in matched_pairs2:
                    matched_pairs[mt[0]] = mt

                all_pairs.extend(matched_pairs2[:, 0])

                num2 = num + len(matched_pairs2)
                if num2 == len(moving):
                    all_matched = True
                    num = num2
            else:
                all_matched = True

        deformed_bundle = Streamlines([])
        warp = []

        # Iterate over each pair of streamlines and deform them
        # Append deformed streamlines in deformed_bundle

        for _, pairs in enumerate(matched_pairs):
            s1 = static[int(pairs[1])]
            s2 = moving_aligned[int(pairs[0])]

            static_s = s1
            moving_s = s2

            reg = DeformableRegistration(
                X=static_s, Y=moving_s, alpha=alpha, beta=beta, max_iterations=max_iter
            )
            ty, pr = reg.register()
            ty = ty.astype(float)
            deformed_bundle.append(ty)
            warp.append(pr)

        warp = pd.DataFrame(warp, columns=["gaussian_kernel", "transforms"])

        # Returns deformed bundle, affinely moved bundle, distance matrix,
        # streamline correspondences, and warp field
        return deformed_bundle, moving_aligned, dist, matched_pairs, warp

    def bundlewarp_vector_filed(moving_aligned, deformed_bundle):
        """Calculate vector fields.

        Vector field computation as the difference between each streamline point
        in the deformed and linearly aligned bundles

        Parameters
        ----------
        moving_aligned : Streamlines
            Linearly (affinely) moved bundle
        deformed_bundle : Streamlines
            Nonlinearly (warped) bundle

        Returns
        -------
        offsets : List
            Vector field modules
        directions : List
            Unitary vector directions
        colors : List
            Colors for bundle warping field vectors. Colors follow the convention
            used in DTI-derived maps (e.g. color FA) :footcite:p:`Pajevic1999`.

        References
        ----------
        .. footbibliography::
        """
        points_aligned, _ = unlist_streamlines(moving_aligned)
        points_deformed, _ = unlist_streamlines(deformed_bundle)
        vector_field = points_deformed - points_aligned

        offsets = np.sqrt(np.sum((vector_field) ** 2, 1))  # vector field modules

        # Normalize vectors to be unitary (directions)
        directions = vector_field / np.array([offsets]).T

        # Define colors mapping the direction vectors to RGB.
        # Absolute value generates DTI-like colors
        colors = directions

        return offsets, directions, colors

    @warning_for_keywords()
    def bundlewarp_shape_analysis(
        moving_aligned, deformed_bundle, *, no_disks=10, plotting=False
    ):
        """Calculate bundle shape difference profile.

        Bundle shape difference analysis using magnitude from BundleWarp
        displacements and BUAN.

        Depending on the number of points of a streamline, and the number of
        segments requested, multiple points may be considered for the computation
        of a given segment; a segment may contain information from a single point;
        or some segments may not contain information from any points. In the latter
        case, the segment will contain an ``np.nan`` value. The point-to-segment
        mapping is defined by the :func:`assignment_map`: for each segment index,
        the point information of the matching index positions, as returned by
        :func:`assignment_map`, are considered for the computation.

        Parameters
        ----------
        moving_aligned : Streamlines
            Linearly (affinely) moved bundle
        deformed_bundle : Streamlines
            Nonlinearly (warped) moved bundle
        no_disks : int, optional
            Number of segments to be created along the length of the bundle
        plotting : Boolean, optional
            Plot bundle shape profile

        Returns
        -------
        shape_profile : np.ndarray
            Float array containing bundlewarp displacement magnitudes along the
            length of the bundle
        stdv : np.ndarray
            Float array containing standard deviations
        """
        n = no_disks
        offsets, directions, colors = bundlewarp_vector_filed(
            moving_aligned, deformed_bundle
        )

        indx = assignment_map(deformed_bundle, deformed_bundle, n)
        indx = np.array(indx)

        rng = np.random.default_rng()

        colors = rng.random((n, 3))

        disks_color = []
        for _, ind in enumerate(indx):
            disks_color.append(tuple(colors[ind]))

        x = np.array(range(1, n + 1))
        shape_profile = np.zeros(n)
        stdv = np.zeros(n)

        for i in range(n):
            mask = indx == i
            if sum(mask):
                shape_profile[i] = np.mean(offsets[mask])
                stdv[i] = np.std(offsets[mask])
            else:
                shape_profile[i] = np.nan
                stdv[i] = np.nan

        if plotting:
            bundle_shape_profile(x, shape_profile, stdv)

        return shape_profile, stdv