""" enigma.py
"""

import operator
from string import ascii_uppercase as UPPER
from copy import deepcopy
from permutations import Permutation,cycleform,permform

def char2int(key):
    try:
        c = key.upper()
        return UPPER.index(c)
    except:
        return key

class Rotor(Permutation):

    def __init__(self,perm,ringsetting='A',key='A',notches=()):

        Permutation.__init__(self,perm)
        self.notches = notches
        self.setRing(ringsetting)
        self.setKey(key)

    def setRing(self,A):
        r = char2int(A)
        cycles = [[(x+r) % 26 for x in cycle] for cycle in self.C]
        perm = permform(cycles)
        self.P = perm
        self.C = cycleform(perm)

    def setKey(self,A):
        k = char2int(A)
        cycles = [[(x-k) % 26 for x in cycle] for cycle in self.C]
        perm = permform(cycles)
        self.P = perm
        self.C = cycleform(perm)
        self.key = UPPER[k]

    def step(self):
        cycles = [[(x-1) % 26 for x in cycle] for cycle in self.C]
        perm = permform(cycles)
        self.P = perm
        self.C = cycleform(perm)

        k = char2int(self.key)
        k = (k+1) % 26
        self.key = UPPER[k]

    def trigger(self):
        return self.key in self.notches

def rotor(strng):
    return Permutation(tuple(strng.lower()))

# The rotors
I     = Rotor("EKMFLGDQVZNTOWYHXUSPAIBRCJ",notches=("Q"))
II    = Rotor("AJDKSIRUXBLHWTMCQGZNPYFVOE",notches=("E"))
III   = Rotor("BDFHJLCPRTXVZNYEIWGAKMUSQO",notches=("V"))
IV    = Rotor("ESOVPZJAYQUIRHXLNFTGKDCMWB",notches=("J"))
V     = Rotor("VZBRGITYUPSDNHLXAWMJQOFECK",notches=("Z"))
VI    = Rotor("JPGVOUMFYQBENHZRDKASXLICTW",notches=("Z","M"))
VII   = Rotor("NZJHGRCXMYSWBOUFAIVLPEKQDT",notches=("Z","M"))
VIII  = Rotor("FKQHTLXOCBJSPDZRAMEWNIUYGV",notches=("Z","M"))

# Used only in Model M-4 with thinB and thinC reflectors
beta  = Rotor("LEYJVCNIXWPBQMDRTAKZGFUHOS")
gamma = Rotor("FSOKANUERHMBTIYCWLQPZXVGJD")

# The reflectors
B     = Rotor("YRUHQSLDPXNGOKMIEBFZCWVJAT")
C     = Rotor("FVPJIAOYEDRZXWGCTKUQSBNMHL")
# The reflectors used in Model M-4 with beta and gamma rotors.
thinB = Rotor("ENKQAUYWJICOPBLMDXZVFTHRGS")
thinC = Rotor("RDOBJNTKVEHMLFCWZAXGYIPSUQ")

# The cyclic permutation and its inverse
rho   = Permutation("BCDEFGHIJKLMNOPQRSTUVWXYZA")

identity =  Permutation("ABCDEFGHIJKLMNOPQRSTUVWXYZ")

# The 'QWERTZU' in military and civilian types. We
# have the 'QWERTZU' act on the left so we use the
# inverse of the input permutation.
military = Permutation("ABCDEFGHIJKLMNOPQRSTUVWXYZ")
civilian = Permutation("QWERTZUIOASDFGHJKPYXCVBNML")

# This class still needs to be checked to see if it
# performs correctly.
class Enigma(object):

    def __init__(self,
            rotorList = (I,II,III),
            reflector = B,
            ringSettings = (),
            stecker = identity,
            qwertzu = military):
        """ The rotor (wheel) list should be given from slowest to fastest:
            U, W_L, W_M, W_R. The ring settings default to 'AA...', the
            stecker and 'qwertzu' both default to the identity. The key
            settings are made after the machine is initialized.
        """
        self.qwertzu = qwertzu
        self.rotors = [deepcopy(R) for R in rotorList]
        self.rotors.insert(0,deepcopy(reflector))
        self.nR = len(self.rotors)

        self.__setRings(ringSettings)
        self.setStecker(stecker)

    def __repr__(self):
        return repr(self.permutation)

    def __str__(self):
        return str(self.permutation)

    def __call__(self,text):
        return "".join(map(self.encipher,text))

    def setStecker(self,stecker):
        try:
            self.SQ = stecker * self.qwertzu
        except:
            self.SQ = Permutation(stecker) * self.qwertzu

        self.setPermutation()

    def __setRings(self,ringSettings):
        # Convert ring setting letters to numbers.
        ringSettings = map(char2int,ringSettings)

        # All ring settings default to 'A'.
        nS = len(ringSettings)
        ringSettings = (0,)*(self.nR - nS) + tuple(ringSettings)
        for k,R in enumerate(self.rotors):
            R.setRing(ringSettings[k])

    def setKeys(self,keySettings):
        # Convert key setting letters to numbers.
        keySettings = map(char2int,keySettings)

        # All key settings default to 'A'.
        nK = len(keySettings)
        keySettings = (0,)*(self.nR - nK) + tuple(keySettings)
        for k,R in enumerate(self.rotors):
            R.setKey(keySettings[k])

        self.setPermutation()

    def setPermutation(self):
        self.permutation = (self.SQ >> reduce(operator.lshift,self.rotors))

    def step(self):
        if self.rotors[-2].trigger():
            motion = (1,1,1)
        elif self.rotors[-1].trigger():
            motion = (0,1,1)
        else:
            motion = (0,0,1)

        for i in range(3):
            if motion[i]:
                self.rotors[i-3].step()

        self.setPermutation()

    def encipher(self,c):
        self.step()
        return self.permutation.encipher(c)