from __future__ import division

import string
import plottools, stringtools

LOWER = string.ascii_lowercase
UPPER = string.ascii_uppercase
LETTER= string.letters

def stripPunc(text):
    """ This strips all punctuation from 'text'.
    """
    return filter(lambda c : c in LETTER, text)

def stripRepeats(text):
    """ This strips all reapeats from 'text'.
    """
    return "".join([ c for (k,c) in enumerate(text)
                     if c not in text[:k]])

def showVigenere():
    print " "+" ".join(list(LOWER))
    print " "+"-"*51
    for k in range(26):
        text = UPPER[k:]+UPPER[:k]
        print " "+" ".join(list(text))

def showAlberti(key,shift=0):
    key = key.lower() + LOWER
    key = stringtools.stripDups(key)
    key = key[shift:] + key[:shift]
    print " "+" ".join(list(key))
    print " "+"-"*51
    for k in range(26):
        text = UPPER[k:]+UPPER[:k]
        print " "+" ".join(list(text))

def showMixed(plainkey="",cipherkey="",plainshift=0,ciphershift=0,sorted=False):
    plain = stringtools.stripDups(plainkey.lower()+LOWER)
    plain = plain[plainshift:]+plain[:plainshift]

    cipherkey = stringtools.stripDups(cipherkey.upper())
    N = len(cipherkey)
    cipher = stringtools.stripDups(cipherkey+UPPER)
    cipherLst = stringtools.breakIntoCols(cipher,N)
    cipher = "".join(cipherLst)
    cipher = cipher[ciphershift:]+cipher[:ciphershift]

    if sorted:
        print " "+" ".join(list(LOWER))
    else:
        print " "+" ".join(list(plain))
    print " "+"-"*51
    for k in range(26):
        text = cipher[k:]+cipher[:k]
        if sorted:
            ztxt = zip(plain,text)
            ztxt.sort()
            text = "".join([c for (p,c) in ztxt])
        print " "+" ".join(list(text))




def vigenere(key, pt):
    """ This implements a Vigenere encipherment of the
        plaintext 'pt' using 'key' as the key.
    """
    # Make 'pt' lowercase and strip out all punctuation.
    pt = stripPunc(pt.lower())

    # Make 'key' lower case and strip out repetitions
    # and punctuation.
    key = stripRepeats(stripPunc(key.lower()))

    # Convert the strings 'key' and 'pt' to lists of
    # indices (in LOWER) of the characters.
    keyLst = [ string.index(LOWER, c) for c in key ]
    ptLst = [ string.index(LOWER, c) for c in pt ]

    # We'll need the length of the 'key'.
    N = len(key)

    # We'll construct a list consisting of the indices
    # (in UPPER) of the cipher text letters.
    ctLst = []
    for i,x in enumerate(ptLst):
        j = i % N
        y = (x+keyLst[j]) % 26
        ctLst.append(y)

    # Convert the list of indices into a string.
    # This is our cipher text.
    return "".join([ UPPER[y] for y in ctLst ])

def revVigenere(key, CT):
    """ This implements a Vigenere decipherment of the
        ciphertext 'CT' using 'key' as the key.
    """
    # Make 'CT' uppercase and strip out all punctuation.
    CT = stripPunc(CT.upper())

    # Make 'key' upper case and strip out repetitions
    # and punctuation.
    #key = stripRepeats(stripPunc(key.upper()))
    key = stripPunc(key.upper())

    # Convert the strings 'key' and 'pt' to lists of
    # indices (in UPPER) of the characters.
    keyLst = [ string.index(UPPER, c) for c in key ]
    CTLst = [ string.index(UPPER, c) for c in CT ]

    # We'll need the length of the 'key'.
    N = len(key)

    # We'll construct a list consisting of the indices
    # (in LOWER) of the plain text letters.
    ptLst = []
    for i,x in enumerate(CTLst):
        j = i % N
        y = (x-keyLst[j]) % 26
        ptLst.append(y)

    # Convert the list of indices into a string.
    # This is our cipher text.
    return "".join([ LOWER[y] for y in ptLst ])