# Some imports in order to get ready to transfer this to a file for inclusion into Sage once stabilized
from sage.combinat.crystals.letters import CrystalOfLetters
from sage.combinat.rigged_configurations.tensor_product_kr_tableaux import TensorProductOfKirillovReshetikhinTableaux
from sage.combinat.rigged_configurations.kr_tableaux import KirillovReshetikhinTableaux

class SolitonCellularAutomata(object):
    """
    A Soliton Cellular Automata system.

    INPUT:

    - ``initial_state`` -- the list of elements, can also be a string
      when ``vacuum`` is 1 and ``n`` is `sl_n`.
    - ``n`` -- (default: 2) the value `sl_n` or a Cartan type
    - ``vacuum`` -- (default: 1) the height of the vacuum element

    EXAMPLES::

        sage: B = SolitonCellularAutomata('341111112241111222311111112222344', 4)
        sage: B
        341111112241111222311111112222344

        sage: B = SolitonCellularAutomata([[4,1],[4,1],[2,1],[2,1],[2,1],[2,1],[3,1],[3,1],[3,2]], 4, 2)
        sage: B
        442222333
        111111112
    """
    def __init__(self, initial_state, n=2, vacuum=1):
        if n in ZZ:
            cartan_type = CartanType(['A',n-1,1])
        else:
            cartan_type = CartanType(n)
        if isinstance(initial_state, str):
            # We consider things 1-9
            initial_state = [[ZZ(x) if x != '.' else ZZ.one()] for x in initial_state]

        self._cartan_type = cartan_type
        self._vacuum = vacuum
        try:
            KRT = TensorProductOfKirillovReshetikhinTableaux(self._cartan_type,
                                                             [[vacuum,len(st)//vacuum]
                                                              for st in initial_state])
            self._states = [KRT(pathlist=initial_state)]
        except TypeError:
            KRT = TensorProductOfKirillovReshetikhinTableaux(self._cartan_type,
                                                             [[vacuum,1]
                                                              for st in initial_state])
            self._states = [KRT(*initial_state)]
        self._evolutions = []
        self._nballs = len(self._states[0])

    def evolve(self, carrier_capacity=None, carrier_index=None):
        """
        INPUT:

        - ``carrier_capacity`` -- (default: the number of balls in
          the system) the size of carrier
        """
        if carrier_capacity is None:
            carrier_capacity = self._nballs
        if carrier_index is None:
            carrier_index = self._vacuum
        passed = False
        empty_carrier = KirillovReshetikhinTableaux(self._cartan_type, carrier_index, carrier_capacity).module_generator()
        vacuum = KirillovReshetikhinTableaux(self._cartan_type, self._vacuum, 1).module_generator()
        carrier_factor = (carrier_index, carrier_capacity)
        last_final_carrier = empty_carrier
        state = self._states[-1]
        dims = state.parent().dims
        while not passed:
            KRT = TensorProductOfKirillovReshetikhinTableaux(self._cartan_type,
                                                             dims + (carrier_factor,))
            elt = KRT(*(list(state) + [empty_carrier]))
            RC = RiggedConfigurations(self._cartan_type, (carrier_factor,) + dims)
            elt2 = RC(*elt.to_rigged_configuration()).to_tensor_product_of_kirillov_reshetikhin_tableaux()
            # Back to an empty carrier or we are not getting any better
            if elt2[0] == empty_carrier or elt2[0] == last_final_carrier:
                passed = True
                KRT = TensorProductOfKirillovReshetikhinTableaux(self._cartan_type, dims)
                self._states.append(KRT(*elt2[1:]))
                self._evolutions.append(carrier_factor)
                if elt2[0] != empty_carrier:
                    print("Last carrier:")
                    print(ascii_art(last_final_carrier))
            else:
                # We need to add more vacuum states
                last_final_carrier = elt2[0]
                dims = tuple([(self._vacuum, 1)]*carrier_capacity) + dims

    # Output functions
    # ----------------

    def __repr__(self):
        return repr(sum((ascii_art('\n'.join(str(letter) for letter in b))
                        for b in self._states[-1]), ascii_art('')))

    def print_states(self, num=None, square=True):
        if num is None:
            num = len(self._states)
        if num > len(self._states):
            for _ in range(num - len(self._states)):
                self.evolve()

        vacuum = KirillovReshetikhinTableaux(self._cartan_type, self._vacuum, 1).module_generator()
        if square:
            num_factors = len(self._states[-1])
            for i,state in enumerate(self._states[:num]):
                state = [vacuum]*(num_factors - len(state)) + list(state)
                print("t: %s"%i)
                print( sum((ascii_art('\n'.join(str(letter) for letter in b)) if b != vacuum else ascii_art('.')
                            for b in state), ascii_art('')) )
        else:
            for i in range(num):
                print("t: %s"%i)
                print(sum((ascii_art('\n'.join(str(letter) for letter in b)) if b != vacuum else ascii_art('.')
                           for b in self._states[n]), ascii_art('')))

    def latex_states(self, num=None, square=True, box_width='5pt'):
        """
        Return a latex verion of the states.

        EXAMPLES::

            sage: B = SolitonCellularAutomata([[2],[-1],[1],[1],[1],[1],[1],[1],[1],[1],[1],[2],[2],[3],[-2],[1],[1],[1],[1],[1],[2],[-1],[1],[1],[1],[1],[1],[2],[3],[3],[-3],[-2]], ['C',3,1])
            sage: view(B.latex_states(15))

            sage: B = SolitonCellularAutomata([[2],[-2],[1],[1],[1],[1],[2],[0],[-3],[1],[1],[1],[1],[1],[2],[2],[3],[-3],], ['D',4,2])
            sage: view(B.latex_states(16))

        Example 4.13 from Yamada:

            sage: B = SolitonCellularAutomata([[3],[3],[1],[1],[1],[1],[2],[2],[2]], ['D',4,3])
            sage: view(B.latex_states(15))

        Example 4.14 from Yamada:

            sage: B = SolitonCellularAutomata([[3],[1],[1],[1],[2],[3],[1],[1],[1],[2],[3],[3]], ['D',4,3])
            sage: view(B.latex_states(15))
        """
        if num is None:
            num = len(self._states)
        if num > len(self._states):
            for _ in range(num - len(self._states)):
                self.evolve()
        from sage.misc.latex import latex
        from sage.misc.latex import LatexExpr

        vacuum = KirillovReshetikhinTableaux(self._cartan_type, self._vacuum, 1).module_generator()
        def compact_repr(b):
            if b == vacuum:
                return "\\makebox[%s]{.}"%box_width
            if self._vacuum == 1:
                return latex(list(b)[0])
            temp = "\\begin{array}{c}"
            temp += r"\\".join(latex(letter) for letter in reversed(b))
            temp += "\\end{array}"
            return "\\makebox[%s]{$%s$}"%(box_width, temp)

        ret = "{\\arraycolsep=0.5pt \\begin{array}"
        if square:
            num_factors = len(self._states[-1])
            ret += "{c|c%s}\n"%('c'*num_factors)
            for i,state in enumerate(self._states[:num]):
                state = [vacuum]*(num_factors - len(state)) + list(state)
                ret += "t = %s & \\cdots & %s \\\\\n"%(i, r" & ".join(compact_repr(b) for b in state))
        else:
            ret += "{cc}\n"
            for i in range(num):
                ret += "t = %s & \\cdots %s \\\\\n"%(i, r" \; ".join(compact_repr(b) for b in state))
        ret += "\\end{array}}\n"
        return LatexExpr(ret)

    def latex_states_word(self, num=None, square=True, box_width='6pt'):
        """
        Return a latex verion of the states.
        """
        if num is None:
            num = len(self._states)
        if num > len(self._states):
            for _ in range(num - len(self._states)):
                self.evolve()
        from sage.misc.latex import latex
        from sage.misc.latex import LatexExpr

        vacuum = KirillovReshetikhinTableaux(self._cartan_type, self._vacuum, 1).module_generator()
        def compact_repr(b):
            if b == vacuum:
                return "{\\color{gray} 1}"
                #return "\\makebox[%s]{.}"%box_width
            if self._vacuum == 1:
                return latex(list(b)[0])
            temp = "\\begin{array}{c}"
            temp += r"\\".join(latex(letter) for letter in reversed(b))
            temp += "\\end{array}"
            #return temp
            return "\\makebox[%s]{$%s$}"%(box_width, temp)

        ret = "{\\begin{array}"
        if square:
            num_factors = len(self._states[-1])
            ret += "{c|c}\n"
            for i,state in enumerate(self._states[:num]):
                state = [vacuum]*(num_factors - len(state)) + list(state)
                ret += "t = %s & \\cdots %s \\\\\n"%(i, r" ".join(compact_repr(b) for b in state))
        else:
            ret += "{cc}\n"
            for i in range(num):
                ret += "t = %s & \\cdots %s \\\\\n"%(i, r" ".join(compact_repr(b) for b in state))
        ret += "\\end{array}}\n"
        return LatexExpr(ret)

    def latex_tikz(self, num=None, vacuum_repr='.'):
        """
        Return a tikz latex version of the states.
        """
        if num is None:
            num = len(self._states)
        if num > len(self._states):
            for _ in range(num - len(self._states)):
                self.evolve()

        from sage.misc.latex import latex
        from sage.misc.latex import LatexExpr
        latex.add_package_to_preamble_if_available("tikz")
        latex.add_to_preamble("\usetikzlibrary{matrix}")
        res = "\\begin{tikzpicture}\n"
        vacuum = KirillovReshetikhinTableaux(self._cartan_type, self._vacuum, 1).module_generator()
        def compact_repr(b):
            if b == vacuum:
                return vacuum_repr
            if self._vacuum == 1:
                return latex(list(b)[0])
            return latex(b)

        num_factors = len(self._states[-1])
        # 
        ret = "\\begin{tikzpicture}\n\\matrix [matrix of math nodes, column sep=-0.6em, nodes={minimum width=12pt}]{\n"
        for i,state in enumerate(self._states[:num]):
            ret += "t = {}\quad & \\cdots".format(i)
            state = [vacuum]*(num_factors - len(state)) + list(state)
            for b in state:
                ret += ' & {}'.format(compact_repr(b))
            ret += "\\\\\n"
        ret += "};\n\\end{tikzpicture}\n"
        return LatexExpr(ret)

def evolve_state(state, carrier, print_carrier=False):
    final = []
    try:
        u = [carrier.maximal_vector()]
    except (ValueError, TypeError, AttributeError):
        u = [carrier.module_generators[0]]
    # Assume every element has the same parent
    R = state[0].parent().R_matrix(carrier)
    for elt in reversed(state):
        up, eltp = R(R.domain()(elt, u[0]))
        u.insert(0, up)
        final.insert(0, eltp)
    def simple_repr(x):
        return ''.join(repr(x).strip('[]').split(', '))
    art = sum((ascii_art(
"""
{!s:^7}
   |
 --+-- {}
   |
{!s:^7}
""".format(simple_repr(state[i]), simple_repr(u[i+1]), simple_repr(final[i])))
             for i in range(len(state))), ascii_art(''))
    art._baseline = 2
    if print_carrier:
        print(ascii_art(simple_repr(u[0])) + art)
    return final

def evolve_state_latex(state, carrier, scale=1):
    from sage.graphs.graph_latex import setup_latex_preamble
    setup_latex_preamble()
    final = []
    try:
        u = [carrier.maximal_vector()]
    except (ValueError, TypeError, AttributeError):
        u = [carrier.module_generators[0]]
    # Assume every element has the same parent
    R = state[0].parent().R_matrix(carrier)
    for elt in reversed(state):
        up, eltp = R(R.domain()(elt, u[0]))
        u.insert(0, up)
        final.insert(0, eltp)
    def simple_repr(x):
        return ''.join(repr(x).strip('[]').split(', '))
    ret = '\\begin{{tikzpicture}}[scale={}]\n'.format(scale)
    for i,val in enumerate(reversed(state)):
        ret += '\\node (i{}) at ({},0.9) {{${}$}};\n'.format(i, -2*i, simple_repr(val))
    for i,val in enumerate(reversed(final)):
        ret += '\\node (t{}) at ({},-1) {{${}$}};\n'.format(i, -2*i, simple_repr(val))
    for i,val in enumerate(reversed(u)):
        ret += '\\node (u{}) at ({},0) {{${}$}};\n'.format(i, -2*i+1, simple_repr(val))
    for i in range(len(state)):
        ret += '\\draw[->] (i{}) -- (t{});\n'.format(i, i)
        ret += '\\draw[->] (u{}) -- (u{});\n'.format(i, i+1)
    ret += '\\end{tikzpicture}'
    return LatexExpr(ret)