"""Module providing an implementation of MatrixLieAlgebras.

There are two main forms of representation for elements of a MatrixLieAlgebra
implemented here. The first one is as a matrix, as elements of R^(n x n).
The second is by choosing a base and remembering the coefficients of an element
in that base. This base will be provided in child classes
(e.g. SkewSymmetricMatrices).

Lead author: Stefan Heyder.
"""

import abc

import geomstats.backend as gs
import geomstats.errors
from geomstats.geometry.base import VectorSpace
from geomstats.geometry.matrices import Matrices

from ._bch_coefficients import BCH_COEFFICIENTS


class MatrixLieAlgebra(VectorSpace, abc.ABC):
    """Class implementing matrix Lie algebra related functions.

    Parameters
    ----------
    dim : int
        Dimension of the Lie algebra as a real vector space.
    n : int
        Amount of rows and columns in the matrix representation of the
        Lie algebra.
    """

    def __init__(self, dim, n, **kwargs):
        super(MatrixLieAlgebra, self).__init__(shape=(n, n), **kwargs)
        geomstats.errors.check_integer(dim, "dim")
        geomstats.errors.check_integer(n, "n")
        self.dim = dim
        self.basis = None
        self.n = n

    bracket = Matrices.bracket

    def baker_campbell_hausdorff(self, matrix_a, matrix_b, order=2):
        """Calculate the Baker-Campbell-Hausdorff approximation of given order.

        The implementation is based on [CM2009a]_ with the pre-computed
        constants taken from [CM2009b]_. Our coefficients are truncated to
        enable us to calculate BCH up to order 15.

        This represents Z = log(exp(X)exp(Y)) as an infinite linear combination
        of the form Z = sum z_i e_i where z_i are rational numbers and e_i are
        iterated Lie brackets starting with e_1 = X, e_2 = Y, each e_i is given
        by some i',i'': e_i = [e_i', e_i''].

        Parameters
        ----------
        matrix_a, matrix_b : array-like, shape=[..., n, n]
            Matrices.
        order : int
            The order to which the approximation is calculated. Note that this
            is NOT the same as using only e_i with i < order.
            Optional, default 2.

        References
        ----------
        .. [CM2009a] F. Casas and A. Murua. An efficient algorithm for
           computing the Baker–Campbell–Hausdorff series and some of its
           applications. Journal of Mathematical Physics 50, 2009
        .. [CM2009b] http://www.ehu.eus/ccwmuura/research/bchHall20.dat
        """
        if order > 15:
            raise NotImplementedError("BCH is not implemented for order > 15.")

        number_of_hom_degree = gs.array(
            [2, 1, 2, 3, 6, 9, 18, 30, 56, 99, 186, 335, 630, 1161, 2182]
        )
        n_terms = gs.sum(number_of_hom_degree[:order])

        el = [matrix_a, matrix_b]
        result = matrix_a + matrix_b

        for i in gs.arange(2, n_terms):
            i_p = BCH_COEFFICIENTS[i, 1] - 1
            i_pp = BCH_COEFFICIENTS[i, 2] - 1

            el.append(self.bracket(el[i_p], el[i_pp]))
            result += (
                float(BCH_COEFFICIENTS[i, 3]) / float(BCH_COEFFICIENTS[i, 4]) * el[i]
            )
        return result

    @abc.abstractmethod
    def basis_representation(self, matrix_representation):
        """Compute the coefficients of matrices in the given basis.

        Parameters
        ----------
        matrix_representation : array-like, shape=[..., n, n]
            Matrix.

        Returns
        -------
        basis_representation : array-like, shape=[..., dim]
            Coefficients in the basis.
        """
        raise NotImplementedError("basis_representation not implemented.")

    def matrix_representation(self, basis_representation):
        """Compute the matrix representation for the given basis coefficients.

        Sums the basis elements according to the coefficents given in
        basis_representation.

        Parameters
        ----------
        basis_representation : array-like, shape=[..., dim]
            Coefficients in the basis.

        Returns
        -------
        matrix_representation : array-like, shape=[..., n, n]
            Matrix.
        """
        if self.basis is None:
            raise NotImplementedError("basis not implemented")

        return gs.einsum("...i,ijk ->...jk", basis_representation, self.basis)
