"""
    Turing machine module

    Represented by:
    - An alphabet of symbols
    - A set of states
    - A sequence of transitions
    - A data tape
    - A read/write head

    An instruction is a tuple (s0, a, b, d, s1)
    such that:
        s0, s1 are states (int)
        a, b are symbols in the alphabet
        d is either the left or right direction
"""

class TM:
    """
        Turing Machine class
    """

    # Blank symbol
    blank = None
    # Directions
    L = -1
    R = +1

    def __init__(self, transitions = [], tape = [], start_state = 0, alphabet = {0,1}):
        self.states = self.build_states(alphabet, transitions)
        self.head = 0
        self.tape = tape
        self.cur_state = start_state
        self._tape_root = 0
        # Check if start state is valid
        assert(self.states[start_state] != None)

    def current_state(self):
        return self.states[self.cur_state]

    def head_index(self):
        return self._tape_root + self.head

    def move(self):
        s = self.current_state()

        # Read head and find matching transition
        value = self.tape[self.head_index()]

        if value in s:
            (b, d, next) = s[value]
            self.cur_state = next
            self.tape[self.head_index()] = b
            self.head = self.head + d

            # Resize tape if head out of range, to simulate infinite tape
            if self.head < 0:
                self.tape.insert(0, TM.blank)
                self._tape_root = self._tape_root + 1
            elif self.head >= len(self.tape):
                self.tape.append(TM.blank)

            return True
        else:
            # Can't move
            return False

    def build_states(self, alphabet, transitions):
        """
            Build the transition graph from a sequence of transitions
        """
        # State map
        states = dict()

        # Alphabet set
        t = set(alphabet)
        t.add(TM.blank)

        for (s0, a, b, d, s1) in transitions:
            # Create states if they do not already exist
            if not s0 in states:
                states[s0] = dict()
            if not s1 in states:
                states[s1] = dict()

            # Symbols must be in alphabet
            assert(a in t)
            assert(b in t)
            # Transition for this symbol must not already exist
            assert(not a in states[s0])

            # Create transition
            states[s0][a] = (b, d, s1)

        return states

    def trimmed_tape(self):
        return list(filter(lambda b: b != None, self.tape))


def from_bits(bits):
    x = 0
    for b in bits:
        x = x << 1
        x = x | b
    return x

if __name__ == "__main__":

    import pprint

    def clear():
        return TM(
            transitions = [
                (0, 0,    1,    TM.R, 0),
                (0, 1,    1,    TM.R, 0),
                (0, None, None, TM.L, 1)
            ],
            tape = [0,0,1,0,1,0],
            start_state = 0
        )

    # S -> 1S1
    # S -> 11
    def halve():
        return TM(
            alphabet = {1},
            transitions = [
                (0, 1, 1,    TM.R, 1),
                (1, 1, None, TM.R, 0)
            ],
            tape = [1,1,1,1,1,1],
            start_state = 0
        )

    def add_one():
        return TM(
            alphabet = {0,1},
            transitions = [
                (0, 0, 0, TM.R, 0),
                (0, 1, 1, TM.R, 0),
                (0, None, None, TM.L, 1),

                (1, 0, 1, TM.L, 3),
                (1, 1, 0, TM.L, 2),

                (3, 0, 0, TM.L, 3),
                (3, 1, 1, TM.L, 3),
                (3, None, None, TM.R, 4),

                (2, 0, 1, TM.L, 3),
                (2, 1, 1, TM.L, 2),
                (2, None, 1, TM.L, 5),

                (5, None, None, TM.R, 4),
            ],
            tape = [0,1,1,1,1,1],
            start_state = 0
        )

    def rshift():
        return TM(
            transitions = [
                (0, 0, None, TM.R, 1),
                (0, 1, None, TM.R, 2),

                (1, 0, 0, TM.R, 1),
                (1, 1, 0, TM.R, 2),
                (1, None, None, TM.L, 3),

                (2, 0, 1, TM.R, 1),
                (2, 1, 1, TM.R, 2),
                (2, None, None, TM.L, 3)
            ],
            tape = [1,1,1,1,1,0],
            start_state = 0
        )

    def q4():
        return TM(
            transitions = [
                (0, 0,    1, TM.R, 1),
                (1, None, 0, TM.L, 2),
                (0, 1, 1, TM.R, 3),

                (3, None, None, TM.L, 4),
                (4, 1,    0,    TM.L, 5),
                (5, None, None, TM.R, 6),
                (3, 0, 0, TM.L, 6),
                (3, 1, 1, TM.L, 6),
            ],
            #tape = [0],
            #tape = [1],
            #tape = [1, 0, 1],
            tape = [1, 0],
        )

    def q5():
        return TM(
            transitions = [
                (0, 0, 0, TM.R, 0),
                (0, 1, 1, TM.R, 0),
                (0, None, 1, TM.L, 1),

                (1, 1, 1, TM.L, 2),
                (1, 0, 1, TM.L, 2),

                (2, 1, 0, TM.L, 2),
                (2, 0, 0, TM.L, 2),
                (2, None, None, TM.R, 3)
            ],
            tape = [0],
            start_state = 0
        )

    tm = q4 ()

    print("STATES:")
    pprint.pprint(tm.states, width=1)

    print("RUN:")

    print("input:", from_bits(tm.trimmed_tape()))
    print(tm.tape)

    while(tm.move()):
        print(tm.trimmed_tape())

    print("output:", from_bits(tm.trimmed_tape()))
    print("head:", tm.head_index())