from enigma import *

class Zygalski(Enigma):

    def setkey(self,key):
        """ This sets the key and computes the composition of
            the permutation applied to the first and fourth letters.
        """
        key = key2exp(key)
        self.key = key

        keyedrotors = [self.conjugate(k,P) for (k,P) in zip(key,self.rotors)]
        permutation = self.reflector ^ ~reduce(operator.mul,keyedrotors)

        keyplus3 = self.advancekey(3,key)
        keyedrotorsplus3 = [self.conjugate(k,P) for (k,P) in zip(keyplus3,self.rotors)]
        permutationplus3 = self.reflector ^ ~reduce(operator.mul,keyedrotorsplus3)

        self.permutation = permutation * permutationplus3

    def advancekey(self,steps,key='current'):
        """ This advances the 'key' by the indicated number of 'steps' of
            the fast rotor. The key is a tuple from the fastest to the
            slowest rotor. Since we don't know where the The correct enigma stepping isn't implemented,
            only a (reversed) odometer-like stepping. 8^(
        """
        if key == 'current':
            key = self.key[:]

        if key:
            key = list(key)
            key[0] += steps
            key[0] %= 26
            if key[0] in range(steps):
                key0 = key[0]
                key = self.advancekey(1,key[1:])
                key.insert(0,key0)

        return key

    def printsheet(self,slow,start=1):
        """ 'slow' is the setting of the slow rotor """
        slow = key2exp(slow)[0]
        shift = start-1

        print "   " + " ".join(list(LOWER))

        for i in range(26):
            results = [" "+LOWER[i]]
            for j in range(26):
                self.setkey((i+shift,j,slow))
                if self.permutation.fixedpoints():
                    results.append("o")
                else:
                    results.append(" ")
            print " ".join(results)