"""Riemannian and pseudo-Riemannian metrics."""

import math
import warnings

import autograd

import geomstats.backend as gs
from geomstats.geometry.connection import Connection


EPSILON = 1e-4
N_CENTERS = 10
TOLERANCE = 1e-5
N_REPETITIONS = 20
N_MAX_ITERATIONS = 50000
N_STEPS = 10


def loss(y_pred, y_true, metric):
    """Compute loss function between prediction and ground truth.

    Loss function given by a Riemannian metric,
    expressed as the squared geodesic distance between the prediction
    and the ground truth.

    Parameters
    ----------
    y_pred
    y_true
    metric

    Returns
    -------
    loss

    """
    loss = metric.squared_dist(y_pred, y_true)
    return loss


def grad(y_pred, y_true, metric):
    """Closed-form for the gradient of the loss function."""
    tangent_vec = metric.log(base_point=y_pred, point=y_true)
    grad_vec = - 2. * tangent_vec

    inner_prod_mat = metric.inner_product_matrix(base_point=y_pred)

    grad = gs.einsum('ni,nij->ni',
                     grad_vec,
                     gs.transpose(inner_prod_mat, axes=(0, 2, 1)))

    return grad


class RiemannianMetric(Connection):
    """Class for Riemannian and pseudo-Riemannian metrics."""

    def __init__(self, dimension, signature=None):
        assert isinstance(dimension, int) or dimension == math.inf
        assert dimension > 0
        super().__init__(dimension=dimension)
        self.signature = signature

    def inner_product_matrix(self, base_point=None):
        """Inner product matrix at the tangent space at a base point.

        Parameters
        ----------
        base_point : array-like, shape=[n_samples, dimension], optional
        """
        raise NotImplementedError(
            'The computation of the inner product matrix'
            ' is not implemented.')

    def inner_product_inverse_matrix(self, base_point=None):
        """Inner product matrix at the tangent space at a base point.

        Parameters
        ----------
        base_point : array-like, shape=[n_samples, dimension], optional
        """
        metric_matrix = self.inner_product_matrix(base_point)
        cometric_matrix = gs.linalg.inv(metric_matrix)
        return cometric_matrix

    def inner_product_derivative_matrix(self, base_point=None):
        """Compute derivative of the inner prod matrix at base point.

        Parameters
        ----------
        base_point : array-like, shape=[n_samples, dimension], optional
        """
        metric_derivative = autograd.jacobian(self.inner_product_matrix)
        return metric_derivative(base_point)

    def christoffels(self, base_point):
        """Compute Christoffel symbols associated with the connection.

        Parameters
        ----------
        base_point: array-like, shape=[n_samples, dimension]

        Returns
        -------
        christoffels: array-like,
                             shape=[n_samples, dimension, dimension, dimension]
        """
        cometric_mat_at_point = self.inner_product_inverse_matrix(base_point)
        metric_derivative_at_point = self.inner_product_derivative_matrix(
            base_point)
        term_1 = gs.einsum('nim,nmkl->nikl',
                           cometric_mat_at_point,
                           metric_derivative_at_point)
        term_2 = gs.einsum('nim,nmlk->nilk',
                           cometric_mat_at_point,
                           metric_derivative_at_point)
        term_3 = - gs.einsum('nim,nklm->nikl',
                             cometric_mat_at_point,
                             metric_derivative_at_point)

        christoffels = 0.5 * (term_1 + term_2 + term_3)
        return christoffels

    def inner_product(self, tangent_vec_a, tangent_vec_b, base_point=None):
        """Inner product between two tangent vectors at a base point.

        Parameters
        ----------
        tangent_vec_a: array-like, shape=[n_samples, dimension]
                                   or shape=[1, dimension]

        tangent_vec_b: array-like, shape=[n_samples, dimension]
                                   or shape=[1, dimension]

        base_point: array-like, shape=[n_samples, dimension]
                                or shape=[1, dimension]

        Returns
        -------
        inner_product : array-like, shape=[n_samples,]
        """
        tangent_vec_a = gs.to_ndarray(tangent_vec_a, to_ndim=2)
        tangent_vec_b = gs.to_ndarray(tangent_vec_b, to_ndim=2)
        n_tangent_vec_a = gs.shape(tangent_vec_a)[0]
        n_tangent_vec_b = gs.shape(tangent_vec_b)[0]

        inner_prod_mat = self.inner_product_matrix(base_point)
        inner_prod_mat = gs.to_ndarray(inner_prod_mat, to_ndim=3)
        n_mats = gs.shape(inner_prod_mat)[0]

        if n_tangent_vec_a != n_mats:
            if n_tangent_vec_a == 1:
                tangent_vec_a = gs.squeeze(tangent_vec_a, axis=0)
                einsum_str_a = 'j,njk->nk'
            elif n_mats == 1:
                inner_prod_mat = gs.squeeze(inner_prod_mat, axis=0)
                einsum_str_a = 'nj,jk->nk'
            else:
                raise ValueError('Shape mismatch for einsum.')
        else:
            einsum_str_a = 'nj,njk->nk'

        aux = gs.einsum(einsum_str_a, tangent_vec_a, inner_prod_mat)
        n_auxs, _ = gs.shape(aux)

        if n_tangent_vec_b != n_auxs:
            if n_auxs == 1:
                aux = gs.squeeze(aux, axis=0)
                einsum_str_b = 'k,nk->n'
            elif n_tangent_vec_b == 1:
                tangent_vec_b = gs.squeeze(tangent_vec_b, axis=0)
                einsum_str_b = 'nk,k->n'
            else:
                raise ValueError('Shape mismatch for einsum.')
        else:
            einsum_str_b = 'nk,nk->n'

        inner_prod = gs.einsum(einsum_str_b, aux, tangent_vec_b)
        inner_prod = gs.to_ndarray(inner_prod, to_ndim=2, axis=1)

        assert gs.ndim(inner_prod) == 2, inner_prod.shape
        return inner_prod

    def squared_norm(self, vector, base_point=None):
        """Compute the square of the norm of a vector.

        Squared norm of a vector associated to the inner product
        at the tangent space at a base point.

        Parameters
        ----------
        vector : array-like, shape=[n_samples, dimension]
        base_point : array-like, shape=[n_samples, dimension]

        Returns
        -------
        sq_norm : array-like, shape=[n_samples,]
        """
        sq_norm = self.inner_product(vector, vector, base_point)
        return sq_norm

    def norm(self, vector, base_point=None):
        """Compute norm of a vector.

        Norm of a vector associated to the inner product
        at the tangent space at a base point.

        Note: This only works for positive-definite
        Riemannian metrics and inner products.

        Parameters
        ----------
        vector : array-like, shape=[n_samples, dimension]
        base_point : array-like, shape=[n_samples, dimension]

        Returns
        -------
        norm : array-like, shape=[n_samples,]
        """
        sq_norm = self.squared_norm(vector, base_point)
        norm = gs.sqrt(sq_norm)
        return norm

    def geodesic(self, initial_point,
                 end_point=None, initial_tangent_vec=None,
                 point_type='vector'):
        """Return the geodesic as function of t.

        Geodesic curve defined by either:
        - an initial point and an initial tangent vector, or
        - an initial point and an end point.

        The geodesic is returned as a function parameterized by t.

        Parameters
        ----------
        initial_point : array-like, shape=[n_samples, dimension]
        end_point : array-like, shape=[n_samples, dimension], optional
        initial_tangent_vec : array-like, shape=[n_samples, dimension],
            optional
        point_type : str, optional

        Returns
        -------
        path : callable
        """
        point_ndim = 1
        if point_type == 'matrix':
            point_ndim = 2

        initial_point = gs.to_ndarray(initial_point,
                                      to_ndim=point_ndim + 1)

        if end_point is None and initial_tangent_vec is None:
            raise ValueError('Specify an end point or an initial tangent '
                             'vector to define the geodesic.')
        if end_point is not None:
            end_point = gs.to_ndarray(end_point,
                                      to_ndim=point_ndim + 1)
            shooting_tangent_vec = self.log(point=end_point,
                                            base_point=initial_point)
            if initial_tangent_vec is not None:
                assert gs.allclose(shooting_tangent_vec, initial_tangent_vec)
            initial_tangent_vec = shooting_tangent_vec
        initial_tangent_vec = gs.array(initial_tangent_vec)
        initial_tangent_vec = gs.to_ndarray(initial_tangent_vec,
                                            to_ndim=point_ndim + 1)

        def path(t):
            """Generate a function parameterizing the geodesic.

            Parameters
            ----------
            t :  parameter value of the geodesic

            Returns
            -------
            point_at_time_t : callable
            """
            t = gs.cast(t, gs.float32)
            t = gs.to_ndarray(t, to_ndim=1)
            t = gs.to_ndarray(t, to_ndim=2, axis=1)
            new_initial_point = gs.to_ndarray(
                initial_point,
                to_ndim=point_ndim + 1)
            new_initial_tangent_vec = gs.to_ndarray(
                initial_tangent_vec,
                to_ndim=point_ndim + 1)

            if point_type == 'vector':
                tangent_vecs = gs.einsum('il,nk->ik',
                                         t,
                                         new_initial_tangent_vec)
            elif point_type == 'matrix':
                tangent_vecs = gs.einsum('il,nkm->ikm',
                                         t,
                                         new_initial_tangent_vec)

            point_at_time_t = self.exp(tangent_vec=tangent_vecs,
                                       base_point=new_initial_point)
            return point_at_time_t

        return path

    def squared_dist(self, point_a, point_b):
        """Squared geodesic distance between two points.

        Parameters
        ----------
        point_a : array-like, shape=[n_samples, dimension]
        point_b : array-like, shape=[n_samples, dimension]

        Returns
        -------
        sq_dist : array-like, shape=[n_samples,]
        """
        log = self.log(point=point_b, base_point=point_a)
        sq_dist = self.squared_norm(vector=log, base_point=point_a)

        return sq_dist

    def dist(self, point_a, point_b):
        """Geodesic distance between two points.

        Note: It only works for positive definite
        Riemannian metrics.

        Parameters
        ----------
        point_a : array-like, shape=[n_samples, dimension]
        point_b : array-like, shape=[n_samples, dimension]

        Returns
        -------
        dist : array-like, shape=[n_samples,]
        """
        sq_dist = self.squared_dist(point_a, point_b)
        dist = gs.sqrt(sq_dist)
        return dist

    def variance(self,
                 points,
                 weights=None,
                 base_point=None,
                 point_type='vector'):
        """Variance of (weighted) points wrt a base point.

        Parameters
        ----------
        points: array-like, shape=[n_samples, dimension]

        weights: array-like, shape=[n_samples, 1], optional
        """
        if point_type == 'vector':
            points = gs.to_ndarray(points, to_ndim=2)
        if point_type == 'matrix':
            points = gs.to_ndarray(points, to_ndim=3)
        n_points = gs.shape(points)[0]

        if weights is None:
            weights = gs.ones((n_points, 1))

        weights = gs.array(weights)
        weights = gs.to_ndarray(weights, to_ndim=2, axis=1)

        sum_weights = gs.sum(weights)

        if base_point is None:
            base_point = self.mean(points, weights)

        variance = 0.

        sq_dists = self.squared_dist(base_point, points)
        variance += gs.einsum('nk,nj->j', weights, sq_dists)

        variance = gs.array(variance)
        variance /= sum_weights

        variance = gs.to_ndarray(variance, to_ndim=1)
        variance = gs.to_ndarray(variance, to_ndim=2, axis=1)
        return variance

    def mean(self, points,
             weights=None,
             n_max_iterations=32,
             epsilon=EPSILON,
             point_type='vector',
             mean_method='default',
             verbose=False):
        """Frechet mean of (weighted) points.

        Parameters
        ----------
        points : array-like, shape=[n_samples, dimension]
        weights : array-like, shape=[n_samples, 1], optional
        verbose : bool, optional

        Returns
        -------
        mean : array-like
            the Frechet mean of points, a point on the manifold
        """
        if mean_method == 'default':

            # TODO(nina): Profile this code to study performance,
            # i.e. what to do with sq_dists_between_iterates.
            def while_loop_cond(iteration, mean, variance, sq_dist):
                result = ~gs.logical_or(
                    gs.isclose(variance, 0.),
                    gs.less_equal(sq_dist, epsilon * variance))
                return result[0, 0] or iteration == 0

            def while_loop_body(iteration, mean, variance, sq_dist):

                logs = self.log(point=points, base_point=mean)

                tangent_mean = gs.einsum('nk,nj->j', weights, logs)

                tangent_mean /= sum_weights

                mean_next = self.exp(
                    tangent_vec=tangent_mean,
                    base_point=mean)

                sq_dist = self.squared_dist(mean_next, mean)
                sq_dists_between_iterates.append(sq_dist)

                variance = self.variance(points=points,
                                         weights=weights,
                                         base_point=mean_next)

                mean = mean_next
                iteration += 1
                return [iteration, mean, variance, sq_dist]

            if point_type == 'vector':
                points = gs.to_ndarray(points, to_ndim=2)
            if point_type == 'matrix':
                points = gs.to_ndarray(points, to_ndim=3)
            n_points = gs.shape(points)[0]

            if weights is None:
                weights = gs.ones((n_points, 1))

            weights = gs.array(weights)
            weights = gs.to_ndarray(weights, to_ndim=2, axis=1)

            sum_weights = gs.sum(weights)

            mean = points[0]
            if point_type == 'vector':
                mean = gs.to_ndarray(mean, to_ndim=2)
            if point_type == 'matrix':
                mean = gs.to_ndarray(mean, to_ndim=3)

            if n_points == 1:
                return mean

            sq_dists_between_iterates = []
            iteration = 0
            sq_dist = gs.array([[0.]])
            variance = gs.array([[0.]])

            last_iteration, mean, variance, sq_dist = gs.while_loop(
                lambda i, m, v, sq: while_loop_cond(i, m, v, sq),
                lambda i, m, v, sq: while_loop_body(i, m, v, sq),
                loop_vars=[iteration, mean, variance, sq_dist],
                maximum_iterations=n_max_iterations)

            if last_iteration == n_max_iterations:
                print('Maximum number of iterations {} reached.'
                      'The mean may be inaccurate'.format(n_max_iterations))

            if verbose:
                print('n_iter: {}, final variance: {}, final dist: {}'.format(
                    last_iteration, variance, sq_dist))

            mean = gs.to_ndarray(mean, to_ndim=2)
            return mean

        if mean_method == 'frechet-poincare-ball':

            lr = 1e-3
            tau = 5e-3

            if len(points) == 1:
                return points

            iteration = 0
            convergence = math.inf
            barycenter = points.mean(0, keepdims=True) * 0

            while convergence > tau and n_max_iterations > iteration:

                iteration += 1

                expand_barycenter = gs.repeat(barycenter, points.shape[0], 0)

                grad_tangent = 2 * self.log(points, expand_barycenter)

                cc_barycenter = self.exp(lr * grad_tangent.sum(0,
                                                               keepdims=True),
                                         barycenter)

                convergence = self.dist(cc_barycenter, barycenter).max().item()

                barycenter = cc_barycenter

            if iteration == n_max_iterations:
                warnings.warn(
                    'Maximum number of iterations {} reached. The '
                    'mean may be inaccurate'.format(n_max_iterations))

            return barycenter

    def adaptive_gradientdescent_mean(self, points,
                                      weights=None,
                                      n_max_iterations=40,
                                      epsilon=1e-12,
                                      init_points=[],
                                      verbose=False):
        """Compute Frechet mean of (weighted) points using adaptive time-steps.

        Frechet mean of (weighted) points using adaptive time-steps
        The loss function optimized is ||M_1(x)||_x (where M_1(x) is
        the tangent mean at x) rather than the mean-square-distance (MSD)
        because this saves computation time.

        Parameters
        ----------
        points: array-like, shape=[n_samples, dimension]

        weights: array-like, shape=[n_samples, 1], optional

        init_points: array-like, shape=[n_init, dimension]
        epsilon: tolerance for stopping the gradient descent
        verbose: verbose mode printing the surrogate value

        epsilon: tolerance for stopping the gradient descent
        """
        # TODO(Xavier): This function assumes that all points are lists
        #  of vectors and not of matrices
        n_points = gs.shape(points)[0]

        if n_points == 1:
            return gs.to_ndarray(points[0], to_ndim=2)

        if weights is None:
            weights = gs.ones((n_points, 1))

        weights = gs.array(weights)
        weights = gs.to_ndarray(weights, to_ndim=2, axis=1)

        sum_weights = gs.sum(weights)

        n_init = len(init_points)

        if n_init == 0:
            current_mean = points[0]
        else:
            current_mean = init_points[0]

        tau = 1.0
        iteration = 0

        logs = self.log(point=points, base_point=current_mean)
        current_tangent_mean = gs.einsum('nk,nj->j', weights, logs)
        current_tangent_mean /= sum_weights
        norm_current_tangent_mean = gs.linalg.norm(current_tangent_mean)

        while (norm_current_tangent_mean > epsilon
                and iteration < n_max_iterations):
            iteration = iteration + 1
            shooting_vector = gs.to_ndarray(
                tau * current_tangent_mean,
                to_ndim=2)
            next_mean = self.exp(
                tangent_vec=shooting_vector,
                base_point=current_mean)
            logs = self.log(point=points, base_point=next_mean)
            next_tangent_mean = gs.einsum('nk,nj->j', weights, logs)
            next_tangent_mean /= sum_weights
            norm_next_tangent_mean = gs.linalg.norm(next_tangent_mean)
            if verbose:
                print(
                    "Iter {0}: tau= {1}, "
                    "norm_current_tangent_mean = {2}".format(
                        iter, tau, norm_current_tangent_mean))
            if norm_next_tangent_mean < norm_current_tangent_mean:
                current_mean = next_mean
                current_tangent_mean = next_tangent_mean
                norm_current_tangent_mean = norm_next_tangent_mean
                tau = max(1.0, 1.0511111 * tau)
            else:
                tau = tau * 0.8

        if iteration == n_max_iterations:
            warnings.warn(
                'Maximum number of iterations {} reached.'
                'The mean may be inaccurate'.format(n_max_iterations))

        return gs.to_ndarray(current_mean, to_ndim=2)

    def diameter(self, points):
        """Give the distance between two farthest points.

        Distance between the two points that are farthest away from each other
        in points.

        Parameters
        ----------
        points

        Returns
        -------
        diameter

        """
        diameter = 0.0
        n_points = points.shape[0]

        for i in range(n_points - 1):
            dist_to_neighbors = self.dist(points[i, :], points[i + 1:, :])
            dist_to_farthest_neighbor = gs.amax(dist_to_neighbors)
            diameter = gs.maximum(diameter, dist_to_farthest_neighbor)

        return diameter

    def closest_neighbor_index(self, point, neighbors):
        """Closest neighbor of point among neighbors.

        Parameters
        ----------
        point
        neighbors
        Returns
        -------
        closest_neighbor_index

        """
        dist = self.dist(point, neighbors)
        closest_neighbor_index = gs.argmin(dist)

        return closest_neighbor_index

[((1052, 1096), 'geomstats.backend.transpose', 'gs.transpose', (['inner_prod_mat'], {'axes': '(0, 2, 1)'}), '(inner_prod_mat, axes=(0, 2, 1))\n', (1064, 1096), True, 'import geomstats.backend as gs\n'), ((2163, 2191), 'geomstats.backend.linalg.inv', 'gs.linalg.inv', (['metric_matrix'], {}), '(metric_matrix)\n', (2176, 2191), True, 'import geomstats.backend as gs\n'), ((2509, 2553), 'autograd.jacobian', 'autograd.jacobian', (['self.inner_product_matrix'], {}), '(self.inner_product_matrix)\n', (2526, 2553), False, 'import autograd\n'), ((3164, 3242), 'geomstats.backend.einsum', 'gs.einsum', (['"""nim,nmkl->nikl"""', 'cometric_mat_at_point', 'metric_derivative_at_point'], {}), "('nim,nmkl->nikl', cometric_mat_at_point, metric_derivative_at_point)\n", (3173, 3242), True, 'import geomstats.backend as gs\n'), ((3314, 3392), 'geomstats.backend.einsum', 'gs.einsum', (['"""nim,nmlk->nilk"""', 'cometric_mat_at_point', 'metric_derivative_at_point'], {}), "('nim,nmlk->nilk', cometric_mat_at_point, metric_derivative_at_point)\n", (3323, 3392), True, 'import geomstats.backend as gs\n'), ((4363, 4402), 'geomstats.backend.to_ndarray', 'gs.to_ndarray', (['tangent_vec_a'], {'to_ndim': '(2)'}), '(tangent_vec_a, to_ndim=2)\n', (4376, 4402), True, 'import geomstats.backend as gs\n'), ((4427, 4466), 'geomstats.backend.to_ndarray', 'gs.to_ndarray', (['tangent_vec_b'], {'to_ndim': '(2)'}), '(tangent_vec_b, to_ndim=2)\n', (4440, 4466), True, 'import geomstats.backend as gs\n'), ((4662, 4702), 'geomstats.backend.to_ndarray', 'gs.to_ndarray', (['inner_prod_mat'], {'to_ndim': '(3)'}), '(inner_prod_mat, to_ndim=3)\n', (4675, 4702), True, 'import geomstats.backend as gs\n'), ((5224, 5278), 'geomstats.backend.einsum', 'gs.einsum', (['einsum_str_a', 'tangent_vec_a', 'inner_prod_mat'], {}), '(einsum_str_a, tangent_vec_a, inner_prod_mat)\n', (5233, 5278), True, 'import geomstats.backend as gs\n'), ((5299, 5312), 'geomstats.backend.shape', 'gs.shape', (['aux'], {}), '(aux)\n', (5307, 5312), True, 'import geomstats.backend as gs\n'), ((5768, 5811), 'geomstats.backend.einsum', 'gs.einsum', (['einsum_str_b', 'aux', 'tangent_vec_b'], {}), '(einsum_str_b, aux, tangent_vec_b)\n', (5777, 5811), True, 'import geomstats.backend as gs\n'), ((5833, 5877), 'geomstats.backend.to_ndarray', 'gs.to_ndarray', (['inner_prod'], {'to_ndim': '(2)', 'axis': '(1)'}), '(inner_prod, to_ndim=2, axis=1)\n', (5846, 5877), True, 'import geomstats.backend as gs\n'), ((7133, 7149), 'geomstats.backend.sqrt', 'gs.sqrt', (['sq_norm'], {}), '(sq_norm)\n', (7140, 7149), True, 'import geomstats.backend as gs\n'), ((8055, 8107), 'geomstats.backend.to_ndarray', 'gs.to_ndarray', (['initial_point'], {'to_ndim': '(point_ndim + 1)'}), '(initial_point, to_ndim=point_ndim + 1)\n', (8068, 8107), True, 'import geomstats.backend as gs\n'), ((8834, 8863), 'geomstats.backend.array', 'gs.array', (['initial_tangent_vec'], {}), '(initial_tangent_vec)\n', (8842, 8863), True, 'import geomstats.backend as gs\n'), ((8894, 8952), 'geomstats.backend.to_ndarray', 'gs.to_ndarray', (['initial_tangent_vec'], {'to_ndim': '(point_ndim + 1)'}), '(initial_tangent_vec, to_ndim=point_ndim + 1)\n', (8907, 8952), True, 'import geomstats.backend as gs\n'), ((11249, 11265), 'geomstats.backend.sqrt', 'gs.sqrt', (['sq_dist'], {}), '(sq_dist)\n', (11256, 11265), True, 'import geomstats.backend as gs\n'), ((11977, 11994), 'geomstats.backend.array', 'gs.array', (['weights'], {}), '(weights)\n', (11985, 11994), True, 'import geomstats.backend as gs\n'), ((12013, 12054), 'geomstats.backend.to_ndarray', 'gs.to_ndarray', (['weights'], {'to_ndim': '(2)', 'axis': '(1)'}), '(weights, to_ndim=2, axis=1)\n', (12026, 12054), True, 'import geomstats.backend as gs\n'), ((12078, 12093), 'geomstats.backend.sum', 'gs.sum', (['weights'], {}), '(weights)\n', (12084, 12093), True, 'import geomstats.backend as gs\n'), ((12279, 12319), 'geomstats.backend.einsum', 'gs.einsum', (['"""nk,nj->j"""', 'weights', 'sq_dists'], {}), "('nk,nj->j', weights, sq_dists)\n", (12288, 12319), True, 'import geomstats.backend as gs\n'), ((12340, 12358), 'geomstats.backend.array', 'gs.array', (['variance'], {}), '(variance)\n', (12348, 12358), True, 'import geomstats.backend as gs\n'), ((12411, 12445), 'geomstats.backend.to_ndarray', 'gs.to_ndarray', (['variance'], {'to_ndim': '(1)'}), '(variance, to_ndim=1)\n', (12424, 12445), True, 'import geomstats.backend as gs\n'), ((12465, 12507), 'geomstats.backend.to_ndarray', 'gs.to_ndarray', (['variance'], {'to_ndim': '(2)', 'axis': '(1)'}), '(variance, to_ndim=2, axis=1)\n', (12478, 12507), True, 'import geomstats.backend as gs\n'), ((18505, 18522), 'geomstats.backend.array', 'gs.array', (['weights'], {}), '(weights)\n', (18513, 18522), True, 'import geomstats.backend as gs\n'), ((18541, 18582), 'geomstats.backend.to_ndarray', 'gs.to_ndarray', (['weights'], {'to_ndim': '(2)', 'axis': '(1)'}), '(weights, to_ndim=2, axis=1)\n', (18554, 18582), True, 'import geomstats.backend as gs\n'), ((18606, 18621), 'geomstats.backend.sum', 'gs.sum', (['weights'], {}), '(weights)\n', (18612, 18621), True, 'import geomstats.backend as gs\n'), ((18911, 18947), 'geomstats.backend.einsum', 'gs.einsum', (['"""nk,nj->j"""', 'weights', 'logs'], {}), "('nk,nj->j', weights, logs)\n", (18920, 18947), True, 'import geomstats.backend as gs\n'), ((19028, 19064), 'geomstats.backend.linalg.norm', 'gs.linalg.norm', (['current_tangent_mean'], {}), '(current_tangent_mean)\n', (19042, 19064), True, 'import geomstats.backend as gs\n'), ((20452, 20490), 'geomstats.backend.to_ndarray', 'gs.to_ndarray', (['current_mean'], {'to_ndim': '(2)'}), '(current_mean, to_ndim=2)\n', (20465, 20490), True, 'import geomstats.backend as gs\n'), ((21467, 21482), 'geomstats.backend.argmin', 'gs.argmin', (['dist'], {}), '(dist)\n', (21476, 21482), True, 'import geomstats.backend as gs\n'), ((3466, 3544), 'geomstats.backend.einsum', 'gs.einsum', (['"""nim,nklm->nikl"""', 'cometric_mat_at_point', 'metric_derivative_at_point'], {}), "('nim,nklm->nikl', cometric_mat_at_point, metric_derivative_at_point)\n", (3475, 3544), True, 'import geomstats.backend as gs\n'), ((4493, 4516), 'geomstats.backend.shape', 'gs.shape', (['tangent_vec_a'], {}), '(tangent_vec_a)\n', (4501, 4516), True, 'import geomstats.backend as gs\n'), ((4546, 4569), 'geomstats.backend.shape', 'gs.shape', (['tangent_vec_b'], {}), '(tangent_vec_b)\n', (4554, 4569), True, 'import geomstats.backend as gs\n'), ((4720, 4744), 'geomstats.backend.shape', 'gs.shape', (['inner_prod_mat'], {}), '(inner_prod_mat)\n', (4728, 4744), True, 'import geomstats.backend as gs\n'), ((5894, 5913), 'geomstats.backend.ndim', 'gs.ndim', (['inner_prod'], {}), '(inner_prod)\n', (5901, 5913), True, 'import geomstats.backend as gs\n'), ((8405, 8453), 'geomstats.backend.to_ndarray', 'gs.to_ndarray', (['end_point'], {'to_ndim': '(point_ndim + 1)'}), '(end_point, to_ndim=point_ndim + 1)\n', (8418, 8453), True, 'import geomstats.backend as gs\n'), ((9291, 9313), 'geomstats.backend.cast', 'gs.cast', (['t', 'gs.float32'], {}), '(t, gs.float32)\n', (9298, 9313), True, 'import geomstats.backend as gs\n'), ((9330, 9357), 'geomstats.backend.to_ndarray', 'gs.to_ndarray', (['t'], {'to_ndim': '(1)'}), '(t, to_ndim=1)\n', (9343, 9357), True, 'import geomstats.backend as gs\n'), ((9374, 9409), 'geomstats.backend.to_ndarray', 'gs.to_ndarray', (['t'], {'to_ndim': '(2)', 'axis': '(1)'}), '(t, to_ndim=2, axis=1)\n', (9387, 9409), True, 'import geomstats.backend as gs\n'), ((9442, 9494), 'geomstats.backend.to_ndarray', 'gs.to_ndarray', (['initial_point'], {'to_ndim': '(point_ndim + 1)'}), '(initial_point, to_ndim=point_ndim + 1)\n', (9455, 9494), True, 'import geomstats.backend as gs\n'), ((9566, 9624), 'geomstats.backend.to_ndarray', 'gs.to_ndarray', (['initial_tangent_vec'], {'to_ndim': '(point_ndim + 1)'}), '(initial_tangent_vec, to_ndim=point_ndim + 1)\n', (9579, 9624), True, 'import geomstats.backend as gs\n'), ((11723, 11755), 'geomstats.backend.to_ndarray', 'gs.to_ndarray', (['points'], {'to_ndim': '(2)'}), '(points, to_ndim=2)\n', (11736, 11755), True, 'import geomstats.backend as gs\n'), ((11812, 11844), 'geomstats.backend.to_ndarray', 'gs.to_ndarray', (['points'], {'to_ndim': '(3)'}), '(points, to_ndim=3)\n', (11825, 11844), True, 'import geomstats.backend as gs\n'), ((11864, 11880), 'geomstats.backend.shape', 'gs.shape', (['points'], {}), '(points)\n', (11872, 11880), True, 'import geomstats.backend as gs\n'), ((11935, 11957), 'geomstats.backend.ones', 'gs.ones', (['(n_points, 1)'], {}), '((n_points, 1))\n', (11942, 11957), True, 'import geomstats.backend as gs\n'), ((14698, 14715), 'geomstats.backend.array', 'gs.array', (['weights'], {}), '(weights)\n', (14706, 14715), True, 'import geomstats.backend as gs\n'), ((14738, 14779), 'geomstats.backend.to_ndarray', 'gs.to_ndarray', (['weights'], {'to_ndim': '(2)', 'axis': '(1)'}), '(weights, to_ndim=2, axis=1)\n', (14751, 14779), True, 'import geomstats.backend as gs\n'), ((14807, 14822), 'geomstats.backend.sum', 'gs.sum', (['weights'], {}), '(weights)\n', (14813, 14822), True, 'import geomstats.backend as gs\n'), ((15190, 15207), 'geomstats.backend.array', 'gs.array', (['[[0.0]]'], {}), '([[0.0]])\n', (15198, 15207), True, 'import geomstats.backend as gs\n'), ((15230, 15247), 'geomstats.backend.array', 'gs.array', (['[[0.0]]'], {}), '([[0.0]])\n', (15238, 15247), True, 'import geomstats.backend as gs\n'), ((15940, 15970), 'geomstats.backend.to_ndarray', 'gs.to_ndarray', (['mean'], {'to_ndim': '(2)'}), '(mean, to_ndim=2)\n', (15953, 15970), True, 'import geomstats.backend as gs\n'), ((18310, 18326), 'geomstats.backend.shape', 'gs.shape', (['points'], {}), '(points)\n', (18318, 18326), True, 'import geomstats.backend as gs\n'), ((18376, 18411), 'geomstats.backend.to_ndarray', 'gs.to_ndarray', (['points[0]'], {'to_ndim': '(2)'}), '(points[0], to_ndim=2)\n', (18389, 18411), True, 'import geomstats.backend as gs\n'), ((18463, 18485), 'geomstats.backend.ones', 'gs.ones', (['(n_points, 1)'], {}), '((n_points, 1))\n', (18470, 18485), True, 'import geomstats.backend as gs\n'), ((19236, 19288), 'geomstats.backend.to_ndarray', 'gs.to_ndarray', (['(tau * current_tangent_mean)'], {'to_ndim': '(2)'}), '(tau * current_tangent_mean, to_ndim=2)\n', (19249, 19288), True, 'import geomstats.backend as gs\n'), ((19538, 19574), 'geomstats.backend.einsum', 'gs.einsum', (['"""nk,nj->j"""', 'weights', 'logs'], {}), "('nk,nj->j', weights, logs)\n", (19547, 19574), True, 'import geomstats.backend as gs\n'), ((19657, 19690), 'geomstats.backend.linalg.norm', 'gs.linalg.norm', (['next_tangent_mean'], {}), '(next_tangent_mean)\n', (19671, 19690), True, 'import geomstats.backend as gs\n'), ((21010, 21036), 'geomstats.backend.amax', 'gs.amax', (['dist_to_neighbors'], {}), '(dist_to_neighbors)\n', (21017, 21036), True, 'import geomstats.backend as gs\n'), ((21060, 21107), 'geomstats.backend.maximum', 'gs.maximum', (['diameter', 'dist_to_farthest_neighbor'], {}), '(diameter, dist_to_farthest_neighbor)\n', (21070, 21107), True, 'import geomstats.backend as gs\n'), ((4856, 4889), 'geomstats.backend.squeeze', 'gs.squeeze', (['tangent_vec_a'], {'axis': '(0)'}), '(tangent_vec_a, axis=0)\n', (4866, 4889), True, 'import geomstats.backend as gs\n'), ((5402, 5425), 'geomstats.backend.squeeze', 'gs.squeeze', (['aux'], {'axis': '(0)'}), '(aux, axis=0)\n', (5412, 5425), True, 'import geomstats.backend as gs\n'), ((8694, 8748), 'geomstats.backend.allclose', 'gs.allclose', (['shooting_tangent_vec', 'initial_tangent_vec'], {}), '(shooting_tangent_vec, initial_tangent_vec)\n', (8705, 8748), True, 'import geomstats.backend as gs\n'), ((9729, 9779), 'geomstats.backend.einsum', 'gs.einsum', (['"""il,nk->ik"""', 't', 'new_initial_tangent_vec'], {}), "('il,nk->ik', t, new_initial_tangent_vec)\n", (9738, 9779), True, 'import geomstats.backend as gs\n'), ((13728, 13764), 'geomstats.backend.einsum', 'gs.einsum', (['"""nk,nj->j"""', 'weights', 'logs'], {}), "('nk,nj->j', weights, logs)\n", (13737, 13764), True, 'import geomstats.backend as gs\n'), ((14420, 14452), 'geomstats.backend.to_ndarray', 'gs.to_ndarray', (['points'], {'to_ndim': '(2)'}), '(points, to_ndim=2)\n', (14433, 14452), True, 'import geomstats.backend as gs\n'), ((14517, 14549), 'geomstats.backend.to_ndarray', 'gs.to_ndarray', (['points'], {'to_ndim': '(3)'}), '(points, to_ndim=3)\n', (14530, 14549), True, 'import geomstats.backend as gs\n'), ((14573, 14589), 'geomstats.backend.shape', 'gs.shape', (['points'], {}), '(points)\n', (14581, 14589), True, 'import geomstats.backend as gs\n'), ((14652, 14674), 'geomstats.backend.ones', 'gs.ones', (['(n_points, 1)'], {}), '((n_points, 1))\n', (14659, 14674), True, 'import geomstats.backend as gs\n'), ((14915, 14945), 'geomstats.backend.to_ndarray', 'gs.to_ndarray', (['mean'], {'to_ndim': '(2)'}), '(mean, to_ndim=2)\n', (14928, 14945), True, 'import geomstats.backend as gs\n'), ((15008, 15038), 'geomstats.backend.to_ndarray', 'gs.to_ndarray', (['mean'], {'to_ndim': '(3)'}), '(mean, to_ndim=3)\n', (15021, 15038), True, 'import geomstats.backend as gs\n'), ((16418, 16459), 'geomstats.backend.repeat', 'gs.repeat', (['barycenter', 'points.shape[0]', '(0)'], {}), '(barycenter, points.shape[0], 0)\n', (16427, 16459), True, 'import geomstats.backend as gs\n'), ((4996, 5030), 'geomstats.backend.squeeze', 'gs.squeeze', (['inner_prod_mat'], {'axis': '(0)'}), '(inner_prod_mat, axis=0)\n', (5006, 5030), True, 'import geomstats.backend as gs\n'), ((5538, 5571), 'geomstats.backend.squeeze', 'gs.squeeze', (['tangent_vec_b'], {'axis': '(0)'}), '(tangent_vec_b, axis=0)\n', (5548, 5571), True, 'import geomstats.backend as gs\n'), ((9934, 9986), 'geomstats.backend.einsum', 'gs.einsum', (['"""il,nkm->ikm"""', 't', 'new_initial_tangent_vec'], {}), "('il,nkm->ikm', t, new_initial_tangent_vec)\n", (9943, 9986), True, 'import geomstats.backend as gs\n'), ((13418, 13443), 'geomstats.backend.isclose', 'gs.isclose', (['variance', '(0.0)'], {}), '(variance, 0.0)\n', (13428, 13443), True, 'import geomstats.backend as gs\n'), ((13464, 13506), 'geomstats.backend.less_equal', 'gs.less_equal', (['sq_dist', '(epsilon * variance)'], {}), '(sq_dist, epsilon * variance)\n', (13477, 13506), True, 'import geomstats.backend as gs\n')]