# sage_setup: distribution = sagemath-modules
"""
Root lattices and root spaces
"""
# ****************************************************************************
#       Copyright (C) 2008-2009 Nicolas M. Thiery <nthiery at users.sf.net>
#
#  Distributed under the terms of the GNU General Public License (GPL)
#                  https://www.gnu.org/licenses/
# ****************************************************************************
from sage.misc.cachefunc import cached_method, cached_in_parent_method
from sage.rings.integer_ring import ZZ
from sage.combinat.free_module import CombinatorialFreeModule
from .root_lattice_realizations import RootLatticeRealizations
import functools


class RootSpace(CombinatorialFreeModule):
    r"""
    The root space of a root system over a given base ring.

    INPUT:

    - ``root_system`` -- a root system
    - ``base_ring`` -- a ring `R`

    The *root space* (or lattice if ``base_ring`` is `\ZZ`) of a root
    system is the formal free module `\bigoplus_i R \alpha_i`
    generated by the simple roots `(\alpha_i)_{i\in I}` of the root system.

    This class is also used for coroot spaces (or lattices).

    .. SEEALSO::

        - :meth:`RootSystem`
        - :meth:`RootSystem.root_lattice` and :meth:`RootSystem.root_space`
        - :meth:`~sage.combinat.root_system.root_lattice_realizations.RootLatticeRealizations`

    Todo: standardize the variable used for the root space in the examples (P?)

    TESTS::

        sage: for ct in CartanType.samples(crystallographic=True)+[CartanType(["A",2],["C",5,1])]:
        ....:     TestSuite(ct.root_system().root_lattice()).run()
        ....:     TestSuite(ct.root_system().root_space()).run()
        sage: r = RootSystem(['A',4]).root_lattice()
        sage: r.simple_root(1)
        alpha[1]
        sage: latex(r.simple_root(1))
        \alpha_{1}
    """

    def __init__(self, root_system, base_ring):
        """
        EXAMPLES::

            sage: P = RootSystem(['A',4]).root_space()
            sage: s = P.simple_reflections()
        """
        from sage.categories.morphism import SetMorphism
        from sage.categories.homset import Hom
        from sage.categories.sets_with_partial_maps import SetsWithPartialMaps
        self.root_system = root_system
        CombinatorialFreeModule.__init__(self, base_ring,
                                         root_system.index_set(),
                                         prefix="alphacheck" if root_system.dual_side else "alpha",
                                         latex_prefix="\\alpha^\\vee" if root_system.dual_side else "\\alpha",
                                         category=RootLatticeRealizations(base_ring))
        if base_ring is not ZZ:
            # Register the partial conversion back from ``self`` to the root lattice
            # See :meth:`_to_root_lattice` for tests
            root_lattice = self.root_system.root_lattice()
            SetMorphism(Hom(self, root_lattice, SetsWithPartialMaps()),
                        self._to_root_lattice
                        ).register_as_conversion()

    def _repr_(self):
        """
        EXAMPLES::

            sage: RootSystem(['A',4]).root_lattice()    # indirect doctest
            Root lattice of the Root system of type ['A', 4]
            sage: RootSystem(['B',4]).root_space()
            Root space over the Rational Field of the Root system of type ['B', 4]
            sage: RootSystem(['A',4]).coroot_lattice()
            Coroot lattice of the Root system of type ['A', 4]
            sage: RootSystem(['B',4]).coroot_space()
            Coroot space over the Rational Field of the Root system of type ['B', 4]
        """
        return self._name_string()

    def _name_string(self, capitalize=True, base_ring=True, type=True):
        """
        EXAMPLES::

            sage: RootSystem(['A',4]).root_space()._name_string()
            "Root space over the Rational Field of the Root system of type ['A', 4]"
        """
        return self._name_string_helper("root", capitalize=capitalize, base_ring=base_ring, type=type)

    simple_root = CombinatorialFreeModule.monomial

    @cached_method
    def to_coroot_space_morphism(self):
        """
        Return the ``nu`` map to the coroot space over the same base ring,
        using the symmetrizer of the Cartan matrix.

        It does not map the root lattice to the coroot lattice, but
        has the property that any root is mapped to some scalar
        multiple of its associated coroot.

        EXAMPLES::

            sage: R = RootSystem(['A',3]).root_space()
            sage: alpha = R.simple_roots()
            sage: f = R.to_coroot_space_morphism()                                      # needs sage.graphs
            sage: f(alpha[1])                                                           # needs sage.graphs
            alphacheck[1]
            sage: f(alpha[1] + alpha[2])                                                # needs sage.graphs
            alphacheck[1] + alphacheck[2]

            sage: R = RootSystem(['A',3]).root_lattice()
            sage: alpha = R.simple_roots()
            sage: f = R.to_coroot_space_morphism()                                      # needs sage.graphs
            sage: f(alpha[1])                                                           # needs sage.graphs
            alphacheck[1]
            sage: f(alpha[1] + alpha[2])                                                # needs sage.graphs
            alphacheck[1] + alphacheck[2]

            sage: S = RootSystem(['G',2]).root_space()
            sage: alpha = S.simple_roots()
            sage: f = S.to_coroot_space_morphism()                                      # needs sage.graphs
            sage: f(alpha[1])                                                           # needs sage.graphs
            alphacheck[1]
            sage: f(alpha[1] + alpha[2])                                                # needs sage.graphs
            alphacheck[1] + 3*alphacheck[2]
        """
        R = self.base_ring()
        C = self.cartan_type().symmetrizer().map(R)
        return self.module_morphism(diagonal=C.__getitem__,
                                    codomain=self.coroot_space(R))

    def _to_root_lattice(self, x):
        """
        Try to convert ``x`` to the root lattice.

        INPUT:

        - ``x`` -- an element of ``self``

        EXAMPLES::

            sage: R = RootSystem(['A',3])
            sage: root_space = R.root_space()
            sage: x = root_space.an_element(); x
            2*alpha[1] + 2*alpha[2] + 3*alpha[3]
            sage: root_space._to_root_lattice(x)
            2*alpha[1] + 2*alpha[2] + 3*alpha[3]
            sage: root_space._to_root_lattice(x).parent()
            Root lattice of the Root system of type ['A', 3]

        This will fail if ``x`` does not have integral coefficients::

            sage: root_space._to_root_lattice(x/2)
            Traceback (most recent call last):
            ...
            ValueError: alpha[1] + alpha[2] + 3/2*alpha[3] does not have integral coefficients

        .. NOTE::

            For internal use only; instead use a conversion::

                sage: R.root_lattice()(x)
                2*alpha[1] + 2*alpha[2] + 3*alpha[3]
                sage: R.root_lattice()(x/2)
                Traceback (most recent call last):
                ...
                ValueError: alpha[1] + alpha[2] + 3/2*alpha[3] does not have integral coefficients

        .. TODO:: generalize diagonal module morphisms to implement this
        """
        try:
            return self.root_system.root_lattice().sum_of_terms((i, ZZ(c)) for i, c in x)
        except TypeError:
            raise ValueError("%s does not have integral coefficients" % x)

    @cached_method
    def _to_classical_on_basis(self, i):
        r"""
        Implement the projection onto the corresponding classical root space or lattice, on the basis.

        EXAMPLES::

            sage: L = RootSystem(["A",3,1]).root_space()
            sage: L._to_classical_on_basis(0)                                           # needs sage.graphs
            -alpha[1] - alpha[2] - alpha[3]
            sage: L._to_classical_on_basis(1)
            alpha[1]
            sage: L._to_classical_on_basis(2)
            alpha[2]
        """
        if i == self.cartan_type().special_node():
            return self._classical_alpha_0()
        else:
            return self.classical().simple_root(i)

    @cached_method
    def to_ambient_space_morphism(self):
        r"""
        The morphism from ``self`` to its associated ambient space.

        EXAMPLES::

            sage: CartanType(['A',2]).root_system().root_lattice().to_ambient_space_morphism()
            Generic morphism:
            From: Root lattice of the Root system of type ['A', 2]
            To:   Ambient space of the Root system of type ['A', 2]
        """
        if self.root_system.dual_side:
            L = self.cartan_type().dual().root_system().ambient_space()
            basis = L.simple_coroots()
        else:
            L = self.cartan_type().root_system().ambient_space()
            basis = L.simple_roots()

        def basis_value(basis, i):
            return basis[i]
        return self.module_morphism(on_basis=functools.partial(basis_value, basis) , codomain=L)


class RootSpaceElement(CombinatorialFreeModule.Element):
    def scalar(self, lambdacheck):
        """
        The scalar product between the root lattice and
        the coroot lattice.

        EXAMPLES::

            sage: L = RootSystem(['B',4]).root_lattice()
            sage: alpha      = L.simple_roots()
            sage: alphacheck = L.simple_coroots()
            sage: alpha[1].scalar(alphacheck[1])                                        # needs sage.graphs
            2
            sage: alpha[1].scalar(alphacheck[2])                                        # needs sage.graphs
            -1

        The scalar products between the roots and coroots are given by
        the Cartan matrix::

            sage: matrix([ [ alpha[i].scalar(alphacheck[j])                             # needs sage.graphs
            ....:            for i in L.index_set() ]
            ....:          for j in L.index_set() ])
            [ 2 -1  0  0]
            [-1  2 -1  0]
            [ 0 -1  2 -1]
            [ 0  0 -2  2]

            sage: L.cartan_type().cartan_matrix()                                       # needs sage.graphs
            [ 2 -1  0  0]
            [-1  2 -1  0]
            [ 0 -1  2 -1]
            [ 0  0 -2  2]

        TESTS:

        Verify that :issue:`15325` (A) is fixed::

            sage: rt = RootSystem(['E', 8])
            sage: lat = rt.root_lattice()
            sage: spc = rt.ambient_space()
            sage: lat.simple_root(1).scalar(spc.simple_coroot(2))
            0

        Verify that directionality is correct for roots of different lengths::

            sage: lat = RootSystem(['B', 3]).root_lattice()
            sage: lat.simple_root(2).scalar(lat.simple_coroot(3))                       # needs sage.graphs
            -2
        """
        if lambdacheck in self.parent().coroot_lattice() or lambdacheck in self.parent().coroot_space():
            # This is the mathematically canonical case, where we use the Cartan matrix to find the scalar product
            zero = self.parent().base_ring().zero()
            cartan_matrix = self.parent().dynkin_diagram()
            return sum( (sum( (lambdacheck[i]*s for i,s in cartan_matrix.column(j)), zero) * c for j,c in self), zero)

        if lambdacheck in self.parent().root_system.ambient_space():
            # lambdacheck lives in the ambient space of the root space, so we take the usual dot product in the ambient space
            return self.to_ambient().dot_product(lambdacheck)

        raise TypeError(f"{lambdacheck} is not in a coroot lattice/space")

    def is_positive_root(self):
        """
        Check whether an element in the root space lies in the
        nonnegative cone spanned by the simple roots.

        EXAMPLES::

            sage: R = RootSystem(['A',3,1]).root_space()
            sage: B = R.basis()
            sage: w = B[0] + B[3]
            sage: w.is_positive_root()
            True
            sage: w = B[1] - B[2]
            sage: w.is_positive_root()
            False
        """
        return all(c >= 0 for c in self.coefficients())

    @cached_in_parent_method
    def associated_coroot(self):
        r"""
        Return the coroot associated to this root.

        OUTPUT:

        An element of the coroot space over the same base ring; in
        particular the result is in the coroot lattice whenever
        ``self`` is in the root lattice.

        EXAMPLES::

            sage: L = RootSystem(["B", 3]).root_space()
            sage: alpha = L.simple_roots()
            sage: alpha[1].associated_coroot()                                          # needs sage.graphs
            alphacheck[1]
            sage: alpha[1].associated_coroot().parent()                                 # needs sage.graphs
            Coroot space over the Rational Field of the Root system of type ['B', 3]

            sage: L.highest_root()                                                      # needs sage.graphs
            alpha[1] + 2*alpha[2] + 2*alpha[3]
            sage: L.highest_root().associated_coroot()                                  # needs sage.graphs
            alphacheck[1] + 2*alphacheck[2] + alphacheck[3]

            sage: alpha = RootSystem(["B", 3]).root_lattice().simple_roots()
            sage: alpha[1].associated_coroot()                                          # needs sage.graphs
            alphacheck[1]
            sage: alpha[1].associated_coroot().parent()                                 # needs sage.graphs
            Coroot lattice of the Root system of type ['B', 3]
        """
        # assert self in self.parent().roots()
        scaled_coroot = self.parent().to_coroot_space_morphism()(self)
        s = self.scalar(scaled_coroot)
        return scaled_coroot.map_coefficients(lambda c: (2 * c) // s)

    def quantum_root(self):
        r"""
        Check whether ``self`` is a quantum root.

        INPUT:

        - ``self`` -- an element of the nonnegative integer span of simple roots

        A root `\alpha` is a quantum root if `\ell(s_\alpha) = \langle 2 \rho, \alpha^\vee \rangle - 1`
        where `\ell` is the length function, `s_\alpha` is the reflection across the hyperplane
        orthogonal to `\alpha`, and `2\rho` is the sum of positive roots.

        .. warning::

            This implementation only handles finite Cartan types and assumes that ``self`` is a root.

        .. TODO:: Rename to is_quantum_root

        EXAMPLES::

            sage: Q = RootSystem(['C',2]).root_lattice()
            sage: positive_roots = Q.positive_roots()
            sage: for x in sorted(positive_roots):                                      # needs sage.graphs
            ....:     print("{} {}".format(x, x.quantum_root()))
            alpha[1] True
            alpha[1] + alpha[2] False
            2*alpha[1] + alpha[2] True
            alpha[2] True
        """
        return len(self.associated_reflection()) == -1 + (self.parent().nonparabolic_positive_root_sum(())).scalar(self.associated_coroot())

    def max_coroot_le(self):
        r"""
        Return the highest positive coroot whose associated root is less than or equal to ``self``.

        INPUT:

        - ``self`` -- an element of the nonnegative integer span of simple roots

        Returns None for the zero element.

        Really ``self`` is an element of a coroot lattice and this method returns the highest root whose
        associated coroot is <= ``self``.

        .. warning::

            This implementation only handles finite Cartan types

        EXAMPLES::

            sage: # needs sage.graphs
            sage: root_lattice = RootSystem(['C',2]).root_lattice()
            sage: root_lattice.from_vector(vector([1,1])).max_coroot_le()
            alphacheck[1] + 2*alphacheck[2]
            sage: root_lattice.from_vector(vector([2,1])).max_coroot_le()
            alphacheck[1] + 2*alphacheck[2]
            sage: root_lattice = RootSystem(['B',2]).root_lattice()
            sage: root_lattice.from_vector(vector([1,1])).max_coroot_le()
            2*alphacheck[1] + alphacheck[2]
            sage: root_lattice.from_vector(vector([1,2])).max_coroot_le()
            2*alphacheck[1] + alphacheck[2]

            sage: root_lattice.zero().max_coroot_le() is None                           # needs sage.graphs
            True
            sage: root_lattice.from_vector(vector([-1,0])).max_coroot_le()              # needs sage.graphs
            Traceback (most recent call last):
            ...
            ValueError: -alpha[1] is not in the positive cone of roots
            sage: root_lattice = RootSystem(['A',2,1]).root_lattice()
            sage: root_lattice.simple_root(1).max_coroot_le()                           # needs sage.graphs
            Traceback (most recent call last):
            ...
            NotImplementedError: Only implemented for finite Cartan type
        """
        if not self.parent().cartan_type().is_finite():
            raise NotImplementedError("Only implemented for finite Cartan type")
        if not self.is_positive_root():
            raise ValueError(f"{self} is not in the positive cone of roots")
        coroots = self.parent().coroot_lattice().positive_roots_by_height(increasing=False)
        for beta in coroots:
            if beta.quantum_root():
                gamma = self - beta.associated_coroot()
                if gamma.is_positive_root():
                    return beta
        return None

    def max_quantum_element(self):
        r"""
        Return a reduced word for the longest element of the Weyl group
        whose shortest path in the quantum Bruhat graph to the identity
        has sum of quantum coroots at most ``self``.

        INPUT:

        - ``self`` -- an element of the nonnegative integer span of simple roots

        Really ``self`` is an element of a coroot lattice.

        .. warning::

            This implementation only handles finite Cartan types

        EXAMPLES::

            sage: # needs sage.graphs sage.libs.gap
            sage: Qvee = RootSystem(['C',2]).coroot_lattice()
            sage: Qvee.from_vector(vector([1,2])).max_quantum_element()
            [2, 1, 2, 1]
            sage: Qvee.from_vector(vector([1,1])).max_quantum_element()
            [1, 2, 1]
            sage: Qvee.from_vector(vector([0,2])).max_quantum_element()
            [2]
        """
        Qvee = self.parent()
        word = []
        while self != Qvee.zero():
            beta = self.max_coroot_le()
            word += list(beta.associated_reflection())
            self = self - beta.associated_coroot()
        W = self.parent().weyl_group()
        return (W.demazure_product(word)).reduced_word()

    def to_ambient(self):
        r"""
        Map ``self`` to the ambient space.

        EXAMPLES::

            sage: alpha = CartanType(['B',2]).root_system().root_lattice().an_element(); alpha
            2*alpha[1] + 2*alpha[2]
            sage: alpha.to_ambient()
            (2, 0)
            sage: alphavee = CartanType(['B',2]).root_system().coroot_lattice().an_element(); alphavee
            2*alphacheck[1] + 2*alphacheck[2]
            sage: alphavee.to_ambient()
            (2, 2)
        """
        return self.parent().to_ambient_space_morphism()(self)


RootSpace.Element = RootSpaceElement
