import random
from ngram_score import ngram_score
import re

fitness = ngram_score('quadgrams.txt') # load our quadgram model

# helper function, converts an integer 0-25 into a character
def i2a(i): return 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'[i%26]

# decipher a piece of text using the substitution cipher and a certain key
def sub_decipher(text,key):
    invkey = [i2a(key.index(i)) for i in 'ABCDEFGHIJKLMNOPQRSTUVWXYZ']


    ret = ''
    for c in text:
        if c.isalpha(): ret += invkey[ord(c.upper())-ord('A')]
        else: ret += c
    return ret

def break_simplesub(ctext,startkey=None):
    ''' perform hill-climbing with a single start. This function may have to be called many times
        to break a substitution cipher. '''
    # make sure ciphertext has all spacing/punc removed and is uppercase
    ctext = re.sub('[^A-Z]','',ctext.upper())
    parentkey,parentscore = startkey or list('ABCDEFGHIJKLMNOPQRSTUVWXYZ'),-99e99
    if not startkey: random.shuffle(parentkey)
    parentscore = fitness.score(sub_decipher(ctext,parentkey))
    count = 0
    while count < 1000:
        a = random.randint(0,25)
        b = random.randint(0,25)
        child = parentkey[:]
        # swap two characters in the child
        child[a],child[b] = child[b],child[a]
        score = fitness.score(sub_decipher(ctext,child))
        # if the child was better, replace the parent with it
        if score > parentscore:
            parentscore, parentkey = score, child[:]
            count = 0 # reset the counter
        count += 1
    return parentscore, parentkey

ctext = 'ag nc nud uvfznavw bsvgadcvzauy zs xuf, nc erszc az av bahncr, znuz ax, pf xs bnuvwavw znc srdcr sg znc yczzcrx sg znc uyhnupcz, znuz vsz uesrd bsmyd pc qudc smz. Ig uvfsvc eaxncx zs dcbahncr zncxc, uvd wcz uz zncar qcuvavw, nc qmxz xmpxzazmzc znc gsmrzn yczzcr sg znc uyhnupcz, vuqcyf D, gsr A, uvd xs eazn znc szncrx'

print "Substitution Cipher solver, you may have to wait several iterations"
print "for the correct result. Press ctrl+c to exit program."

# keep going until we are killed by the user
i = 0
maxscore = -99e99
while 1:
    i = i+1 # keep track of how many iterations we have done
    score, key = break_simplesub(ctext,list('ABCDEFGHIJKLMNOPQRSTUVWXYZ'))
    if score > maxscore:
        maxscore,maxkey = score,key[:]
        print '\nbest score so far:',maxscore,'on iteration',i
        print '    best key: '+''.join(maxkey)
        print '    plaintext: '+ sub_decipher(ctext,maxkey)