"""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

[(51, 'geomstats.backend.transpose', 'gs.transpose', True, 'import geomstats.backend as gs\n'), (84, 'geomstats.backend.linalg.inv', 'gs.linalg.inv', True, 'import geomstats.backend as gs\n'), (94, 'autograd.jacobian', 'autograd.jacobian', False, 'import autograd\n'), (112, 'geomstats.backend.einsum', 'gs.einsum', True, 'import geomstats.backend as gs\n'), (115, 'geomstats.backend.einsum', 'gs.einsum', True, 'import geomstats.backend as gs\n'), (143, 'geomstats.backend.to_ndarray', 'gs.to_ndarray', True, 'import geomstats.backend as gs\n'), (144, 'geomstats.backend.to_ndarray', 'gs.to_ndarray', True, 'import geomstats.backend as gs\n'), (149, 'geomstats.backend.to_ndarray', 'gs.to_ndarray', True, 'import geomstats.backend as gs\n'), (164, 'geomstats.backend.einsum', 'gs.einsum', True, 'import geomstats.backend as gs\n'), (165, 'geomstats.backend.shape', 'gs.shape', True, 'import geomstats.backend as gs\n'), (179, 'geomstats.backend.einsum', 'gs.einsum', True, 'import geomstats.backend as gs\n'), (180, 'geomstats.backend.to_ndarray', 'gs.to_ndarray', True, 'import geomstats.backend as gs\n'), (222, 'geomstats.backend.sqrt', 'gs.sqrt', True, 'import geomstats.backend as gs\n'), (252, 'geomstats.backend.to_ndarray', 'gs.to_ndarray', True, 'import geomstats.backend as gs\n'), (266, 'geomstats.backend.array', 'gs.array', True, 'import geomstats.backend as gs\n'), (267, 'geomstats.backend.to_ndarray', 'gs.to_ndarray', True, 'import geomstats.backend as gs\n'), (339, 'geomstats.backend.sqrt', 'gs.sqrt', True, 'import geomstats.backend as gs\n'), (364, 'geomstats.backend.array', 'gs.array', True, 'import geomstats.backend as gs\n'), (365, 'geomstats.backend.to_ndarray', 'gs.to_ndarray', True, 'import geomstats.backend as gs\n'), (367, 'geomstats.backend.sum', 'gs.sum', True, 'import geomstats.backend as gs\n'), (375, 'geomstats.backend.einsum', 'gs.einsum', True, 'import geomstats.backend as gs\n'), (377, 'geomstats.backend.array', 'gs.array', True, 'import geomstats.backend as gs\n'), (380, 'geomstats.backend.to_ndarray', 'gs.to_ndarray', True, 'import geomstats.backend as gs\n'), (381, 'geomstats.backend.to_ndarray', 'gs.to_ndarray', True, 'import geomstats.backend as gs\n'), (552, 'geomstats.backend.array', 'gs.array', True, 'import geomstats.backend as gs\n'), (553, 'geomstats.backend.to_ndarray', 'gs.to_ndarray', True, 'import geomstats.backend as gs\n'), (555, 'geomstats.backend.sum', 'gs.sum', True, 'import geomstats.backend as gs\n'), (568, 'geomstats.backend.einsum', 'gs.einsum', True, 'import geomstats.backend as gs\n'), (570, 'geomstats.backend.linalg.norm', 'gs.linalg.norm', True, 'import geomstats.backend as gs\n'), (603, 'geomstats.backend.to_ndarray', 'gs.to_ndarray', True, 'import geomstats.backend as gs\n'), (643, 'geomstats.backend.argmin', 'gs.argmin', True, 'import geomstats.backend as gs\n'), (118, 'geomstats.backend.einsum', 'gs.einsum', True, 'import geomstats.backend as gs\n'), (145, 'geomstats.backend.shape', 'gs.shape', True, 'import geomstats.backend as gs\n'), (146, 'geomstats.backend.shape', 'gs.shape', True, 'import geomstats.backend as gs\n'), (150, 'geomstats.backend.shape', 'gs.shape', True, 'import geomstats.backend as gs\n'), (182, 'geomstats.backend.ndim', 'gs.ndim', True, 'import geomstats.backend as gs\n'), (259, 'geomstats.backend.to_ndarray', 'gs.to_ndarray', True, 'import geomstats.backend as gs\n'), (281, 'geomstats.backend.cast', 'gs.cast', True, 'import geomstats.backend as gs\n'), (282, 'geomstats.backend.to_ndarray', 'gs.to_ndarray', True, 'import geomstats.backend as gs\n'), (283, 'geomstats.backend.to_ndarray', 'gs.to_ndarray', True, 'import geomstats.backend as gs\n'), (284, 'geomstats.backend.to_ndarray', 'gs.to_ndarray', True, 'import geomstats.backend as gs\n'), (287, 'geomstats.backend.to_ndarray', 'gs.to_ndarray', True, 'import geomstats.backend as gs\n'), (356, 'geomstats.backend.to_ndarray', 'gs.to_ndarray', True, 'import geomstats.backend as gs\n'), (358, 'geomstats.backend.to_ndarray', 'gs.to_ndarray', True, 'import geomstats.backend as gs\n'), (359, 'geomstats.backend.shape', 'gs.shape', True, 'import geomstats.backend as gs\n'), (362, 'geomstats.backend.ones', 'gs.ones', True, 'import geomstats.backend as gs\n'), (446, 'geomstats.backend.array', 'gs.array', True, 'import geomstats.backend as gs\n'), (447, 'geomstats.backend.to_ndarray', 'gs.to_ndarray', True, 'import geomstats.backend as gs\n'), (449, 'geomstats.backend.sum', 'gs.sum', True, 'import geomstats.backend as gs\n'), (462, 'geomstats.backend.array', 'gs.array', True, 'import geomstats.backend as gs\n'), (463, 'geomstats.backend.array', 'gs.array', True, 'import geomstats.backend as gs\n'), (479, 'geomstats.backend.to_ndarray', 'gs.to_ndarray', True, 'import geomstats.backend as gs\n'), (544, 'geomstats.backend.shape', 'gs.shape', True, 'import geomstats.backend as gs\n'), (547, 'geomstats.backend.to_ndarray', 'gs.to_ndarray', True, 'import geomstats.backend as gs\n'), (550, 'geomstats.backend.ones', 'gs.ones', True, 'import geomstats.backend as gs\n'), (575, 'geomstats.backend.to_ndarray', 'gs.to_ndarray', True, 'import geomstats.backend as gs\n'), (582, 'geomstats.backend.einsum', 'gs.einsum', True, 'import geomstats.backend as gs\n'), (584, 'geomstats.backend.linalg.norm', 'gs.linalg.norm', True, 'import geomstats.backend as gs\n'), (625, 'geomstats.backend.amax', 'gs.amax', True, 'import geomstats.backend as gs\n'), (626, 'geomstats.backend.maximum', 'gs.maximum', True, 'import geomstats.backend as gs\n'), (154, 'geomstats.backend.squeeze', 'gs.squeeze', True, 'import geomstats.backend as gs\n'), (169, 'geomstats.backend.squeeze', 'gs.squeeze', True, 'import geomstats.backend as gs\n'), (264, 'geomstats.backend.allclose', 'gs.allclose', True, 'import geomstats.backend as gs\n'), (292, 'geomstats.backend.einsum', 'gs.einsum', True, 'import geomstats.backend as gs\n'), (418, 'geomstats.backend.einsum', 'gs.einsum', True, 'import geomstats.backend as gs\n'), (438, 'geomstats.backend.to_ndarray', 'gs.to_ndarray', True, 'import geomstats.backend as gs\n'), (440, 'geomstats.backend.to_ndarray', 'gs.to_ndarray', True, 'import geomstats.backend as gs\n'), (441, 'geomstats.backend.shape', 'gs.shape', True, 'import geomstats.backend as gs\n'), (444, 'geomstats.backend.ones', 'gs.ones', True, 'import geomstats.backend as gs\n'), (453, 'geomstats.backend.to_ndarray', 'gs.to_ndarray', True, 'import geomstats.backend as gs\n'), (455, 'geomstats.backend.to_ndarray', 'gs.to_ndarray', True, 'import geomstats.backend as gs\n'), (498, 'geomstats.backend.repeat', 'gs.repeat', True, 'import geomstats.backend as gs\n'), (157, 'geomstats.backend.squeeze', 'gs.squeeze', True, 'import geomstats.backend as gs\n'), (172, 'geomstats.backend.squeeze', 'gs.squeeze', True, 'import geomstats.backend as gs\n'), (296, 'geomstats.backend.einsum', 'gs.einsum', True, 'import geomstats.backend as gs\n'), (410, 'geomstats.backend.isclose', 'gs.isclose', True, 'import geomstats.backend as gs\n'), (411, 'geomstats.backend.less_equal', 'gs.less_equal', True, 'import geomstats.backend as gs\n')]