from sage.algebras.quatalg.quaternion_algebra import QuaternionAlgebra
from sage.combinat.free_module import CombinatorialFreeModule
from sage.modules.with_basis.indexed_element import IndexedFreeModuleElement
from sage.categories.magmatic_algebras import MagmaticAlgebras
from sage.rings.all import AA, ZZ
from sage.matrix.matrix_space import MatrixSpace
from sage.misc.table import table

class Octonion(IndexedFreeModuleElement):
    def conjugate(self):
        r"""
        SETUP::

            sage: from mjo.octonions import Octonions

        EXAMPLES::

            sage: O = Octonions()
            sage: x = sum(O.gens())
            sage: x.conjugate()
            e0 - e1 - e2 - e3 - e4 - e5 - e6 - e7

        TESTS::

        Conjugating twice gets you the original element::

            sage: set_random_seed()
            sage: O = Octonions()
            sage: x = O.random_element()
            sage: x.conjugate().conjugate() == x
            True

        """
        C = MatrixSpace(ZZ,8).diagonal_matrix((1,-1,-1,-1,-1,-1,-1,-1))
        return self.parent().from_vector(C*self.to_vector())

    def real(self):
        r"""
        Return the real part of this octonion.

        The real part of an octonion is its projection onto the span
        of the first generator. In other words, the "first dimension"
        is real and the others are imaginary.

        SETUP::

            sage: from mjo.octonions import Octonions

        EXAMPLES::

            sage: O = Octonions()
            sage: x = sum(O.gens())
            sage: x.real()
            e0

        TESTS:

        This method is idempotent::

            sage: set_random_seed()
            sage: O = Octonions()
            sage: x = O.random_element()
            sage: x.real().real() == x.real()
            True

        """
        return (self + self.conjugate())/2

    def imag(self):
        r"""
        Return the imaginary part of this octonion.

        The imaginary part of an octonion is its projection onto the
        orthogonal complement of the span of the first generator. In
        other words, the "first dimension" is real and the others are
        imaginary.

        SETUP::

            sage: from mjo.octonions import Octonions

        EXAMPLES::

            sage: O = Octonions()
            sage: x = sum(O.gens())
            sage: x.imag()
            e1 + e2 + e3 + e4 + e5 + e6 + e7

        TESTS:

        This method is idempotent::

            sage: set_random_seed()
            sage: O = Octonions()
            sage: x = O.random_element()
            sage: x.imag().imag() == x.imag()
            True

        """
        return (self - self.conjugate())/2

    def _norm_squared(self):
        return (self*self.conjugate()).coefficient(0)

    def norm(self):
        r"""
        Return the norm of this octonion.

        SETUP::

            sage: from mjo.octonions import Octonions

        EXAMPLES::

            sage: O = Octonions()
            sage: O.one().norm()
            1

        TESTS:

        The norm is nonnegative and belongs to the base field::

            sage: set_random_seed()
            sage: O = Octonions()
            sage: n = O.random_element().norm()
            sage: n >= 0 and n in O.base_ring()
            True

        The norm is homogeneous::

            sage: set_random_seed()
            sage: O = Octonions()
            sage: x = O.random_element()
            sage: alpha = O.base_ring().random_element()
            sage: (alpha*x).norm() == alpha.abs()*x.norm()
            True

        """
        return self._norm_squared().sqrt()

    def inverse(self):
        r"""
        Return the inverse of this element if it exists.

        SETUP::

            sage: from mjo.octonions import Octonions

        EXAMPLES::

            sage: O = Octonions()
            sage: x = sum(O.gens())
            sage: x*x.inverse() == O.one()
            True

        ::

            sage: O = Octonions()
            sage: O.one().inverse() == O.one()
            True

        TESTS::

            sage: set_random_seed()
            sage: O = Octonions()
            sage: x = O.random_element()
            sage: x.is_zero() or ( x*x.inverse() == O.one() )
            True

        """
        if self.is_zero():
            raise ValueError("zero is not invertible")
        return self.conjugate()/self._norm_squared()


    def cayley_dickson(self, Q=None):
        r"""
        Return the Cayley-Dickson representation of this element in terms
        of the quaternion algebra ``Q``.

        The Cayley-Dickson representation is an identification of
        octionions `x` and `y` with pairs of quaternions `(a,b)` and
        `(c,d)` respectively such that:

        * `x + y = (a+b, c+d)`
        * `xy` = (ac - \bar{d}*b, da + b\bar{c})`
        * `\bar{x} = (a,-b)`

        where `\bar{x}` denotes the conjugate of `x`.

        SETUP::

            sage: from mjo.octonions import Octonions

        EXAMPLES::

            sage: O = Octonions()
            sage: x = sum(O.gens())
            sage: x.cayley_dickson()
            (1 + i + j + k, 1 + i + j + k)

        """
        if Q is None:
            Q = QuaternionAlgebra(self.base_ring(), -1, -1)

        i,j,k = Q.gens()
        a = (self.coefficient(0)*Q.one() +
             self.coefficient(1)*i +
             self.coefficient(2)*j +
             self.coefficient(3)*k )
        b = (self.coefficient(4)*Q.one() +
             self.coefficient(5)*i +
             self.coefficient(6)*j +
             self.coefficient(7)*k )

        from sage.categories.sets_cat import cartesian_product
        P = cartesian_product([Q,Q])
        return P((a,b))


class Octonions(CombinatorialFreeModule):
    r"""
    SETUP::

        sage: from mjo.octonions import Octonions

    EXAMPLES::

        sage: Octonions()
        Octonion algebra with base ring Algebraic Real Field
        sage: Octonions(field=QQ)
        Octonion algebra with base ring Rational Field

    """
    def __init__(self,
                 field=AA,
                 prefix="e"):

        # Not associative, not commutative
        category = MagmaticAlgebras(field).FiniteDimensional()
        category = category.WithBasis().Unital()

        super().__init__(field,
                         range(8),
                         element_class=Octonion,
                         category=category,
                         prefix=prefix,
                         bracket=False)

        # The product of each basis element is plus/minus another
        # basis element that can simply be looked up on
        # https://en.wikipedia.org/wiki/Octonion
        e0, e1, e2, e3, e4, e5, e6, e7 = self.gens()
        self._multiplication_table = (
            (e0, e1, e2, e3, e4, e5, e6, e7),
            (e1,-e0, e3,-e2, e5,-e4,-e7, e6),
            (e2,-e3,-e0, e1, e6, e7,-e4,-e5),
            (e3, e2,-e1,-e0, e7,-e6, e5,-e4),
            (e4,-e5,-e6,-e7,-e0, e1, e2, e3),
            (e5, e4,-e7, e6,-e1,-e0,-e3, e2),
            (e6, e7, e4,-e5,-e2, e3,-e0,-e1),
            (e7,-e6, e5, e4,-e3,-e2, e1,-e0),
        )

    def product_on_basis(self, i, j):
        return self._multiplication_table[i][j]

    def one_basis(self):
        r"""
        Return the monomial index (basis element) corresponding to the
        octonion unit element.

        SETUP::

            sage: from mjo.octonions import Octonions

        TESTS:

        This gives the correct unit element::

            sage: set_random_seed()
            sage: O = Octonions()
            sage: x = O.random_element()
            sage: x*O.one() == x and O.one()*x == x
            True

        """
        return 0

    def _repr_(self):
        return ("Octonion algebra with base ring %s" % self.base_ring())

    def multiplication_table(self):
        """
        Return a visual representation of this algebra's multiplication
        table (on basis elements).

        SETUP::

            sage: from mjo.octonions import Octonions

        EXAMPLES:

        The multiplication table is what Wikipedia says it is::

            sage: Octonions().multiplication_table()
            +----++----+-----+-----+-----+-----+-----+-----+-----+
            | *  || e0 | e1  | e2  | e3  | e4  | e5  | e6  | e7  |
            +====++====+=====+=====+=====+=====+=====+=====+=====+
            | e0 || e0 | e1  | e2  | e3  | e4  | e5  | e6  | e7  |
            +----++----+-----+-----+-----+-----+-----+-----+-----+
            | e1 || e1 | -e0 | e3  | -e2 | e5  | -e4 | -e7 | e6  |
            +----++----+-----+-----+-----+-----+-----+-----+-----+
            | e2 || e2 | -e3 | -e0 | e1  | e6  | e7  | -e4 | -e5 |
            +----++----+-----+-----+-----+-----+-----+-----+-----+
            | e3 || e3 | e2  | -e1 | -e0 | e7  | -e6 | e5  | -e4 |
            +----++----+-----+-----+-----+-----+-----+-----+-----+
            | e4 || e4 | -e5 | -e6 | -e7 | -e0 | e1  | e2  | e3  |
            +----++----+-----+-----+-----+-----+-----+-----+-----+
            | e5 || e5 | e4  | -e7 | e6  | -e1 | -e0 | -e3 | e2  |
            +----++----+-----+-----+-----+-----+-----+-----+-----+
            | e6 || e6 | e7  | e4  | -e5 | -e2 | e3  | -e0 | -e1 |
            +----++----+-----+-----+-----+-----+-----+-----+-----+
            | e7 || e7 | -e6 | e5  | e4  | -e3 | -e2 | e1  | -e0 |
            +----++----+-----+-----+-----+-----+-----+-----+-----+

        """
        n = self.dimension()
        # Prepend the header row.
        M = [["*"] + list(self.gens())]

        # And to each subsequent row, prepend an entry that belongs to
        # the left-side "header column."
        M += [ [self.monomial(i)] + [ self.monomial(i)*self.monomial(j)
                                    for j in range(n) ]
               for i in range(n) ]

        return table(M, header_row=True, header_column=True, frame=True)


class OctonionMatrix:
    r"""
    A pseudo-matrix class that supports octonion entries.

    Matrices in SageMath can't have base rings that are
    non-commutative, much less non-associative. The "matrix" scaling,
    addition, and multiplication operations for this class are all
    wholly inefficient, but are hand-written to guarantee that they
    are performed in the correct order. Of course, it can't guarantee
    that you won't write something visually ambiguous like
    `A*B*C`... but you already have that problem with the
    non-associative octonions themselves.

    This class is only as sophisticated as it need to be to implement
    the Jordan and inner-products in the space of Hermitian matrices
    with octonion entries, namely ``(X*Y+Y*X)/2`` and
    ``(X*Y).trace().real()`` for two octonion matrices ``X`` and
    ``Y``.

    .. WARNING:

        These are not matrices in the usual sense! Matrix
        multiplication is associative. Multiplication of octonion
        "matrices" cannot be, since the multiplication of the
        underlying octonions is not (consider two 1-by-1 matrices each
        containing a single octonion).
    """
    def __init__(self, entries):
        r"""
        Initialize this matrix with a list of lists in (row,column) order,
        just like in SageMath.
        """
        self._nrows = len(entries)

        if self._nrows == 0:
            self._ncols = 0
        else:
            # We don't check that you haven't supplied two rows (or
            # columns) of different lengths!
            self._ncols = len(entries[0])

        self._entries = entries

    def __getitem__(self, indices):
        r"""
        SETUP::

            sage: from mjo.octonions import Octonions, OctonionMatrix

        EXAMPLES::

            sage: O = Octonions(field=QQ)
            sage: M = OctonionMatrix([ [O.one(),  O.zero()],
            ....:                      [O.zero(), O.one() ] ])
            sage: M[0,0]
            e0
            sage: M[0,1]
            0
            sage: M[1,0]
            0
            sage: M[1,1]
            e0
        """
        i,j = indices
        return self._entries[i][j]

    def nrows(self):
        r"""
        SETUP::

            sage: from mjo.octonions import Octonions, OctonionMatrix

        EXAMPLES::

            sage: O = Octonions(field=QQ)
            sage: M = OctonionMatrix([ [O.one(),  O.zero()],
            ....:                      [O.zero(), O.one() ],
            ....:                      [O.zero(), O.zero()] ])
            sage: M.nrows()
            3

        """
        return self._nrows

    def ncols(self):
        r"""

        SETUP::

            sage: from mjo.octonions import Octonions, OctonionMatrix

        EXAMPLES::

            sage: O = Octonions(field=QQ)
            sage: M = OctonionMatrix([ [O.one(),  O.zero()],
            ....:                      [O.zero(), O.one() ],
            ....:                      [O.zero(), O.zero()] ])
            sage: M.ncols()
            2

        """
        return self._ncols

    def __repr__(self):
        return table(self._entries, frame=True)._repr_()

    def __mul__(self,rhs):
        r"""

        SETUP::

            sage: from mjo.octonions import Octonions, OctonionMatrix

        EXAMPLES::

            sage: O = Octonions(QQ)
            sage: e1 = O.monomial(1)
            sage: e2 = O.monomial(2)
            sage: e1*e2
            e3
            sage: e2*e1
            -e3
            sage: E1 = OctonionMatrix([[e1]])
            sage: E2 = OctonionMatrix([[e2]])
            sage: E1*E2
            +----+
            | e3 |
            +----+
            sage: E2*E1
            +-----+
            | -e3 |
            +-----+

        """
        if not self.ncols() == rhs.nrows():
            raise ValueError("dimension mismatch")

        m = self.nrows()
        n = self.ncols()
        p = rhs.ncols()

        C = lambda i,j: sum( self[i,k]*rhs[k,j] for k in range(n) )
        return OctonionMatrix([ [C(i,j) for j in range(m)]
                                for i in range(p) ] )

    def __rmul__(self,scalar):
        r"""

        SETUP::

            sage: from mjo.octonions import Octonions, OctonionMatrix

        EXAMPLES::

            sage: O = Octonions(QQ)
            sage: M = OctonionMatrix([[O.one(), O.zero()],
            ....:                     [O.zero(),O.one() ] ])
            sage: 2*M
            +------+------+
            | 2*e0 | 0    |
            +------+------+
            | 0    | 2*e0 |
            +------+------+

        r"""
        # SCALAR GOES ON THE LEFT HERE
        return OctonionMatrix([ [scalar*self[i,j]
                                 for j in range(self.ncols())]
                                for i in range(self.nrows()) ])

    def __add__(self,rhs):
        r"""
        SETUP::

            sage: from mjo.octonions import Octonions, OctonionMatrix

        EXAMPLES::

            sage: O = Octonions(QQ)
            sage: e0,e1,e2,e3,e4,e5,e6,e7 = O.gens()
            sage: A = OctonionMatrix([ [e0,e1],
            ....:                      [e2,e3] ])
            sage: B = OctonionMatrix([ [e4,e5],
            ....:                      [e6,e7] ])
            sage: A+B
            +---------+---------+
            | e0 + e4 | e1 + e5 |
            +---------+---------+
            | e2 + e6 | e3 + e7 |
            +---------+---------+

        """
        if not self.ncols() == rhs.ncols():
            raise ValueError("column dimension mismatch")
        if not self.nrows() == rhs.nrows():
            raise ValueError("row dimension mismatch")
        return OctonionMatrix([ [self[i,j] + rhs[i,j]
                                 for j in range(self.ncols())]
                                for i in range(self.nrows()) ])