"""Chapter 6, exercise 4 in Allen Downey's Think Complexity book.

The goal of this exercise is to implement a Turing machine.
See http://en.wikipedia.org/wiki/Turing_machine.

4. Write a class named TMDrawer that generates an image that represents the
   state of the tape and the position and state of the head. For one example of
   what that might look like, see
   http://mathworld.wolfram.com/TuringMachine.html.

"""
import matplotlib.pyplot as pyplot
import numpy

from CADrawer import Drawer


class TMDrawer(Drawer):
    """Implementation of Drawer using matplotlib for Turing Machines.

    We visualize the tape and machine head/state by drawing the tape state as
    a row, then adding a row underneath that for the head & state. The tape head
    is drawn as a solid square whose color represents the state.

    """
    def __init__(self, tape_colors, state_colors, cmap=None):
        self.tape_colors = tape_colors
        self.state_colors = state_colors
        self.cmap = cmap if cmap else 'spectral'

    def draw(self, tm, start=0, end=None):
        """Draws the TM using pyplot.pcolor."""
        tape = tm.get_array(start, end)
        rows, cols = tape.shape

        # create a new array that has twice as many rows
        steps = rows
        rows *= 2
        a = numpy.zeros((rows, cols), dtype=numpy.int8)

        # copy the tape state into the array at every other row
        for i in xrange(steps):
            for j in xrange(cols):
                #a[n, j] = self.tape_colors[tape[i, j]]
                c = self.tape_colors[tape[i, j]]
                a[i * 2, j] = c

        # copy the head position & state into every other row
        n = 1
        for pos, state in tm.history:
            a[n, pos] = self.state_colors[state]
            n += 2

        # flipud puts the first row at the top;
        pyplot.pcolor(numpy.flipud(a), cmap=self.cmap)
        pyplot.axis([0, cols, 0, rows])
        pyplot.colorbar()

        # empty lists draw no ticks
        pyplot.xticks([])
        pyplot.yticks([])

    def show(self):
        """display the pseudocolor representation of the CA"""
        pyplot.show()

    def save(self, filename='tm.png'):
        """save the pseudocolor representation of the TM in (filename)."""
        pyplot.savefig(filename)