""" aristocrats.py
    c odenthal
    2005/03/01

    This module contains a class designed to be
    an aid in solving "ARISTOCRATS" cryptograms.
"""

from __future__ import division
import plottools, stringtools

cipher2plain = dict(zip(stringtools.ascii_uppercase,"_"*26))
plain2cipher = dict(zip(stringtools.ascii_lowercase,"_"*26))
startTable = stringtools.maketrans(stringtools.ascii_uppercase,"_"*26)

class Aristocrat(str):

    def __init__(self,S=""):

        str.__init__(self,S)
        self.cipher2plain = cipher2plain.copy()
        self.plain2cipher = plain2cipher.copy()
        self.cipherText = self.upper()
        self.decode()

    def __repr__(self):
        return self.cipherText

    def __str__(self):
        return "\n"+stringtools.displayStrings(
                    self.cipherText,
                    (self.plainText, self.cipherText),
                    linelength=80,
                    spaces=1)

    def decode(self):
        self.makeTable()
        self.plainText = self.cipherText.translate(self.transTable)

    def makeTable(self):
        """ This makes a decipherment table that will
            translate the cipher text to plain text.
        """
        cp = self.cipher2plain.items()
        ct = "".join([c for c,p in cp])
        pt = "".join([p for c,p in cp])
        self.transTable = stringtools.maketrans(ct,pt)

    def _map(self,CLett,pLett):
        """ Decodes the cipher text 'Ctext' as plain text 'ptext'.
        """
        C = CLett.upper()
        p = pLett.lower()
        oldC, self.plain2cipher[p] = self.plain2cipher[p], C
        self.cipher2plain[oldC] = "_"
        self.cipher2plain[C] = p
        self.decode()

    def _unmap(self,C):
        C = C.upper()
        p = self.cipher2plain[C]
        self.plain2cipher[p] = "_"
        self.cipher2plain[C] = "_"
        self.decode()

    def update(self,CT,PT=""):
        """ This calls '_map' or '_unmap' possibly multiple
            times to update the current deciphering guess.
            If no 'PT' is used, then the ciphertext 'CT' is
            unmapped. Otherwise, the plain text 'PT' is
            assigned to the cipher text 'CT'.
        """
        CT = list(CT)
        PT = list(PT)
        while PT:
            self._map(CT.pop(0),PT.pop(0))
        while CT:
            self._unmap(CT.pop())
        print self

    def load(self,filename):
        self.cipherText = open(filename).read().upper().strip()
        self.decode()

    def reset(self):
        self.cipher2plain = cipher2plain.copy()
        self.plain2cipher = plain2cipher.copy()
        self.decode()

    def getLetterFreq(self):
        """ Gets the letter fequencies of the cipher text.
        """
        D = {}
        for char in self.cipherText:
            D[char] = D.get(char,0)+1
        return D

    def plotLetterFreq(self):
        """ Plots the letter frequencies of the cipher text.
        """
        plottools.freqPlot(self)

    def showLetterFreq(self,linelength=80):
        """ Displays the frequencies of the cipher letters.
        """
        freqDict = self.getLetterFreq()
        freqList = []
        for char in self.cipherText:
            freq = freqDict[char]
            if char not in stringtools.ascii_uppercase:
                nextChar = " "
            elif freq < 10:
                nextChar = str(freq)
            else:
                nextChar = "+"
            freqList.append(nextChar)
        freqString = "".join(freqList)

        print stringtools.displayStrings(
                    self.cipherText,
                    (self.plainText, self.cipherText, freqString),
                    linelength=linelength,
                    spaces=1)

    def showLowFreqContacts(self,lowFreq=0.04,spaces=1):
        """ Shows the contact numbers that the characters
            of the cipher text have with the low frequency
            characters.
        """
        freqDict = self.getLetterFreq()
        freqList = [freqDict.get(char,0) for char in stringtools.letters]
        totaLett = sum(freqList)
        cutOff = lowFreq*totaLett

        print cutOff
        lowFreqList = [count<cutOff for count in freqList]
        lowFreqDict = dict(zip(stringtools.letters,lowFreqList))

        CT = self.cipherText
        contacts = {}
        for k,c in enumerate(CT[1:]):
            if not lowFreqDict.get(c,True):
                if lowFreqDict.get(CT[k],False):
                    contacts[c] = contacts.get(c,0) + 1
        for k,c in enumerate(CT[:-1]):
            if not lowFreqDict.get(c,True):
                if lowFreqDict.get(CT[k+1],False):
                    contacts[c] = contacts.get(c,0) + 1

        contactList = [(contacts[c],c) for c in contacts]
        contactList.sort()
        contactList.reverse()

        chars = [" "+c for (f,c) in contactList]
        freqs = [f for (f,c) in contactList]
        for k,f in enumerate(freqs):
            if f < 10:
                freqs[k] = " "+str(f)
            else:
                freqs[k] = str(f)

        s = " "*spaces
        chars = s.join(chars)
        freqs = s.join(freqs)

        contactFreq = "\n"+ \
            "CT   " + chars +"\n"+  \
            "freq " + freqs

        print contactFreq

    def showAlphabets(self,spaces=1):
        """ Displays the cipher text and plain text alphabets.
        """

        s = " "*spaces
        pt1 = stringtools.ascii_lowercase
        ct1 = s.join([self.plain2cipher[k] for k in pt1])

        ct2 = stringtools.ascii_uppercase
        pt2 = s.join([self.cipher2plain[k] for k in ct2])

        alphabets =  "\n"+ \
            "pt  " + s.join(list(pt1)) +"\n"+  \
            "CT  " + ct1 +"\n"+  \
            "\n"+ \
            "CT  " + s.join(list(ct2)) +"\n"+  \
            "pt  " + pt2

        print alphabets

    def showCipherFreq(self,spaces=1):
        """ Displays the cipher text frequencies in order.
        """
        freqDict = self.getLetterFreq()
        freqList = [(freqDict.get(char,0), char)
                        for char in stringtools.letters]
        freqList.sort()
        freqList.reverse()

        chars = [c for (f,c) in freqList]
        freqs = [f for (f,c) in freqList]
        for k,f in enumerate(freqs):
            if f < 10:
                freqs[k] = str(f)
            else:
                freqs[k] = '+'

        s = " "*spaces
        chars = s.join(chars)
        freqs = s.join(freqs)

        cipherFreq = "\n"+ \
            "CT   " + chars +"\n"+  \
            "freq " + freqs

        print cipherFreq

    def showPositionChart(self):
        """ Shows the position frequencies of cipher text
            letters in the cipher text words.
        """
        wordList = self.cipherText.split()
        wordList = [word.strip(stringtools.punctuation)
                        for word in wordList]

        posFreqs = {}
        for c in stringtools.letters:
            posFreqs[c] = {1:0, 2:0, 3:0, -1:0, -2:0, -3:0, 'mid':0}

        for word in wordList:
            k = 1
            while len(word)>1:
                if k > 3:
                    for c in word:
                        posFreqs[c]['mid'] += 1
                    word = ""
                else:
                    init,word,final = word[0],word[1:-1],word[-1]
                    posFreqs[init][k] += 1
                    posFreqs[final][-k] += 1
                    k += 1
            if len(word) == 1:
                if k > 3:
                    posFreqs[word]['mid'] += 1
                else:
                    posFreqs[word][k] += 1
                    posFreqs[word][-k] += 1

        positions = [1,2,3,'mid',-3,-2,-1]
        posFreqList = [[posFreqs[c][p]
                            for p in positions]
                            for c in stringtools.letters]

        for posF in posFreqList:
            posF.insert(0,sum(posF))

        for k,posF in enumerate(posFreqList):
            for j,f in enumerate(posF):
                if f < 10:
                    posFreqList[k][j] = " "+str(f)
                else:
                    posFreqList[k][j] = str(f)
            posFreqList[k].insert(0,stringtools.letters[k]+" ")

        posFreqList = [" "*3 + " ".join(posFreq)
                            for posFreq in posFreqList]

        positionString = "\n".join(posFreqList)
        headerString = "\n  CT  FQ  1  2  3  m -3 -2 -1\n"

        print headerString + positionString

def loadAristocrat(filename):
    """ Loads the 'aristocrat' stored in file 'filename'.
    """
    SList = open(filename).readlines()
    SList = [strng.strip() for strng in SList]
    Strng = " ".join(SList)
    return Aristocrat(Strng)