"""
 vigenere.py

 Playing around with the Vigenere cipher;
 see http://en.wikipedia.org/wiki/Vigenere_cipher for its details.

 To keep things simple, the only characters allowed are 
 space, A-Z, and a-z converted to uppercase.
 All other characters are ignored.

 Encoding and decoding requires a secret key, also 
 composed of (or converted to) uppercase letters.

 The cipher works like this :

      HELLO       7  4 11 14     = plain text; space=0, A-Z = 1-26
    + ABABA    +  1  2  1  2     = secret key, repeated as many times as needed
      -----      -----------       -- addition is mod 27 --
      IGMNP       8  6 12 16     = encrypted message

 The Letters() class defined below implements the cipher with this API :

   >>> Letters('hello')
   Letters('HELLO')

   >>> Letters('hello').encode('ab')
   Letters('IGMNP')

   >>> Letters('hello').encode('ab').decode('ab')
   Letters('HELLO')

   >>> Letters("So hows it going? Yeah, I thought so.").encode('strange')
   Letters('KHRIBCXSBKAUVNF RZSHMSBRUVVZZAKAFV')

   >>> Letters('KHRIBCXSBKAUVNF RZSHMSBRUVVZZAKAFV').decode('strange')
   Letters('SO HOWS IT GOING YEAH I THOUGHT SO')

 So there you are.

 Jim Mahoney | March 2012 | MIT License
"""

class Letter:
    """ A-Z with values 0-25 and addition mod 26. """
    def __init__(self, char):
        """ char is a character or Letter """
        self.char = char.upper()
    def upper(self):
        # This is a trick to make Letter(Letter('a')) same as Letter('a'),
        # since both 'a'.upper() and Letter('a').upper() will return 'A'.
        return self.char
    def __repr__(self):
        return "Letter('" + self.char + "')"
    def value(self):
        """ Return 0 for space, 1 to 26 for 'A' to 'Z'.
            >>> Letter('b').value()
            2
        """
        return 0 if self.char == ' ' else ord(self.char) - ord('A') + 1
    def __add__(self, other, sign=1):
        result = (self.value() + sign * other.value()) % 27
        return Letter(' ') if result==0 else Letter(chr(result + ord('A') - 1))
    def __sub__(self, other):
        return self.__add__(other, -1)

class Letters:
    """ A Letter sequence, with array subscripts, encode(key), decode(key) """
    def __init__(self, chars):
        """ chars is a list of char or Letter """
        chars = map(lambda x: x.upper(), chars)
        self.letters = [Letter(c) for c in \
                        filter(lambda x: x==' ' or 'A'<=x<='Z', chars)]
    def __len__(self):
        return len(self.letters)
    def __repr__(self):
        return "Letters('" + ''.join([l.char for l in self.letters]) + "')"
    def __getitem__(self, n):
        """ Array-like integer or slice subscripts.
            >>> Letters('abcdefg')[0]
            Letter('A')
            >>> Letters('abcdefg')[2:5]
            Letters('CDE')
        """
        if type(n)==slice:
            (start, stop, step) = (n.start or 0, n.stop, n.step or 1)
            return Letters([self.letters[i] for i in range(start, stop, step)])
        else:        
            return self.letters[n]
    def __iter__(self):
        self.index = -1
        return self
    def next(self):
        self.index += 1
        if self.index < len(self):
            return self[self.index]
        else:
            raise StopIteration
    def zip(self, key):
        """ Return list of letters and key chars (repeated as needed)
            >>> Letters('TWO').zip('ab')
            [(Letter('T'), 'a'), (Letter('W'), 'b'), (Letter('O'), 'a')]
        """
        return zip(self.letters, key * (len(self)/len(key) + 1))
    def encode(self, key):
        return Letters([l + Letter(c) for (l, c) in self.zip(key)])
    def decode(self, key):
        return Letters([l - Letter(c) for (l, c) in self.zip(key)])

if __name__ == "__main__":
    import doctest
    doctest.testmod()