MAX_KEY_LENGTH = 20
PEAK_MIN_OFFSET = 4

NUM_LETTERS = 26
FIRST_LETTER = "A"

KEY_LENGTH_GUESS = 6

def index_of_coincidence(s, shift):
    """Find the index of coincidence of s with itself shifted by shift"""
    return len([ i for i in xrange(len(s)) if s[i] == s[(i + shift) % len(s)]])

def peak_frequency(l):
    """Guess peak frequency in l by finding the frequency for which the
    average peak value would be maximal."""
    def average_value_for_frequency(k):
        return sum([ l[i] for i in xrange(1,len(l)) if i % k == 0 ]) / \
                float(len(l)/k)

    return max(xrange(1,len(l)), key=average_value_for_frequency)

def key_length_guess(s):
    l = [index_of_coincidence(s,shift) for shift in xrange(MAX_KEY_LENGTH)]

    print "First", MAX_KEY_LENGTH, " indexes of coincidence:"
    print l

    key_length = peak_frequency(l)

    print "Guessed key length is", key_length

    if key_length == KEY_LENGTH_GUESS:
        print "That is the key length we guessed earlier by manual examination."
    else:
        print "This differs from our original guess,", KEY_LENGTH_GUESS

    return key_length

def frequency_vector(s, key_length, offset):
    v = [0] * NUM_LETTERS
    i = 0
    for c in s:
        if i % key_length == offset:
            v[ord(c) - ord(FIRST_LETTER)] += 1
        i += 1

    for i in xrange(len(v)):
        v[i] *= (1.0/float(len(s)))

    return v

def rotated_vector_match(v1, v2):
    """Return the rotation i of vector v1 such that v1 * v2 is maximal
    Note that a list of possibilities is returned. Note that we rotate
    backwards, in an attempt to guess what the forward rotation was."""
    def rotated(v, i):
        return v[i:]+v[:i]

    matchings = {}

    for i in xrange(len(v1)):
        s = sum([ rotated(v1,i)[j] * v2[j] for j in xrange(len(v1))])
        matchings[i] = s

    return [ i for i in matchings.keys()
            if matchings[i] == max(matchings.values()) ]

def guess_key(s):
    key_length = key_length_guess(s)

    key = ""

    for i in xrange(key_length):
        freq_vector = frequency_vector(s, key_length, i)
        best_rotations = rotated_vector_match(freq_vector,
                basiccrypt.LETTER_FREQUENCIES)
        print "Best rotations for key letter",i,":", best_rotations

        best_key_letters = [chr(ord(FIRST_LETTER) + k) for k in best_rotations]

        print "In letters:", best_key_letters
        print "Guessing", best_key_letters[0]
        key += best_key_letters[0]

    print "Guessed key is", key
    return key

def vigenere_decrypt(s, key, encrypt = False):
    decrypted = ""

    i = 0
    for c in s:
        offset = ord(key[i % len(key)]) - ord(FIRST_LETTER)
        if encrypt: offset = -offset
        cur_letter = ord(c) - ord(FIRST_LETTER)
        new_letter = (cur_letter - offset) % NUM_LETTERS
        decrypted += chr(ord(FIRST_LETTER) + new_letter)
        i += 1

    return decrypted


if __name__ == '__main__':
    import basiccrypt
    import sys
    filename = "cipher1.txt"
    if len(sys.argv) > 1:
        filename = sys.argv[1]
    s = basiccrypt.read_cipher_from_file(filename).upper()
    key = guess_key(s)
    print "Decryption:"
    print vigenere_decrypt(s, key)