import itertools


class SimpleFST(object):
    """
    An implementation of a simple finite state transducer.

    Inputs are words in {0,1}*
    Outputs are words in {+,-}*
    Subcubes are words in {0,1,*}*, where * means that dimension is spanned by the subcube.

    transition table is a dict that maps (state, input) to (state, output)

    Interpreting the input as a vertex of a hypercube, the output gives the orientation of the incident
    edges, where + means incoming. This induces a USO on any hypercube, as long as each state outputs + for one
    input and - for the other.
    """

    def __init__(self, table=None, start=0, edges=None):
        # two ways to initialize the automaton, keep both for compatibility
        if table is None:
            assert edges is not None
            table = self.table_from_edges(edges)
        else:
            assert edges is None
        self.table = table
        self.states = set(state for (state, v) in self.table)
        self.start_state = start
        self.k = len(self.states)
        self.validate_table()
        self.build_inverse_table()


    def validate_table(self):
        """ Check if table describes a simple FST."""
        assert self.start_state in self.states

        for state in self.states:
            # make sure transitions exist and are valid
            nextstate, out0 = self.table[state, "0"]
            assert nextstate in self.states
            assert out0 in "+-"

            nextstate, out1 = self.table[state, "1"]
            assert nextstate in self.states
            assert out1 in "+-"

            # make sure state is either even or odd
            assert out0 != out1

    
    def build_inverse_table(self):
        """ Precalculate inverse state transition table to find + and - transitions easily."""
        self.inv_table = {}

        for state in self.states:
            nextstate, out = self.table[state, "0"]
            self.inv_table[state, out] = nextstate, "0"
            nextstate, out = self.table[state, "1"]
            self.inv_table[state, out] = nextstate, "1"

   
    def transduce(self, vertex):
        # implements the outmap of the USO generated by the automaton
        state = self.start_state
        for v in vertex:
            state, out = self.table[state, v]
            yield out


    def outmap(self, vertex, subcube):
        # NOTE: subcubes are described by words in {0,1,*}*
        # this method makes no attempt to validate that the vertex actually belongs to the subcube
        return [o for (o, s) in zip(self.transduce(vertex), subcube) if s == "*"]
    

    def find_sink(self, subcube):
        # find sink through structure of automaton, i.e. cheating
        res = []
        state = self.start_state
        for i, dim in enumerate(subcube):
            if dim == "*":
                state, v = self.inv_table[state, "+"]
            else:
                # jump
                state, _ = self.table[state, dim]
                v = dim
            res.append(v)
        return "".join(res)


    def fingerprint_subcube(self, subcube):
        """
        produce a fingerprint of the transducer of a subcube.
        doesn't guarantee absolutely identical automatons, only identical up to the dimension of the subcube
        they're being evaluated on.
        """
        # Forward pass to accumulate data for each active dimension, namely a set of reachable states and
        # an updated state transition table
        states_per_dim = []
        table_per_dim = []
        FIRST = True
        # each chunk is a number of skipped dimensions followed by a spanned dimension
        for chunk in subcube.split("*"):
            transitions = dict((q, reduce(lambda q, dim: self.table[q, dim][0], chunk, q))
                               for q in self.states)
            table = dict(((q1, a), (transitions[q2], b)) for ((q1, a), (q2, b)) in self.table.iteritems())
            table_per_dim.append(table)

            if FIRST:
                states = set([transitions[self.start_state]])
                FIRST = False
            else:
                states = set()
                previous_states = states_per_dim[-1]
                for state in previous_states:
                    states.add(table[state, "0"][0])
                    states.add(table[state, "1"][0])
            states_per_dim.append(states)

        # Backward pass to create fingerprint. In the last dimension, all states of the same parity are 
        # identical. For previous dimensions, states are defined by their parity and the parity of both
        # their children.
        state_index = dict((q, "") for q in self.states)
        res = []
        for (reachable_states, table) in reversed(zip(states_per_dim[:-1], table_per_dim[1:])):
            # assign a class triple to each state, namely the parity of the state, and the class of its
            # two neighbors
            state_lst = [(q, table[q, "0"][1] + state_index[table[q, "0"][0]] + state_index[table[q, "1"][0]])
                         for q in reachable_states]

            # find out what classes we have, put them in some canonical order,
            # then assign an index to each class (basically a reverse lookup of the list of classes)
            classes = sorted(set(c for (state, c) in state_lst))
            class_index = dict((c, str(i)) for (i,c) in enumerate(classes))
            
            # assign to each state the index of its class
            state_index = dict((q, class_index[c]) for (q, c) in state_lst)
    
            # add the result of this iteration to the fingerprint
            #res.extend(sorted(set(c for (state, c) in state_lst if state in states)))
            res.extend(classes)
            res.append("|")

        return "".join(res)


    def get_edges(self):
        edges = [(q_old, a, b, q_new) for ((q_old, a), (q_new, b)) in self.table.iteritems()]
        return sorted(edges)


    def table_from_edges(self, edges):
        return dict(((q_old, a), (q_new,b)) for (q_old, a, b, q_new) in edges)

