from fractions import gcd
from operator import itemgetter

class Coding:

	def encode(self, chars):
		nums = []
		for char in chars:
			char = char.upper()
			nums.append(ord(char) - ord('A'))
		return nums
		
	def decode(self, nums):
		chars = []
		for num in nums:
			chars.append(chr(num + ord('A')))
		return chars

	def diff(self,a,b):
		if ord(a) > ord(b):
			return ord(a)-ord(b)
		else:
			return ord(b)-ord(a)

class Affine:

	def __init__(self,filename = "text.txt"):
         self.filename = filename
         self.getText()
         self.ce = Coding()
         self.setAlphabet()

	def cryptAnalysis(self):
		self.makecount()
		self.printCount()
		
	def getText(self):
		result_f = open(self.filename)
		self.msg = ""	
		for line in result_f:
			self.msg = self.msg + line.strip("\n")
		self.msg.upper()

	def setAlphabet(self):
		self.alphabet = ["A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "N", "O", "P", "Q", "R", "S", "T", "U", "V", "W", "X", "Y", "Z"]

	def IsNumber(self,number):
		try:
			int(number)
			return True
		except ValueError:
			return False

	# number of instances of a letter or small phrase (th, the, etc...) in a text
	def numberOfInstances(self,msg, phrase):
	    num = 0
	    lenarray = len(msg)-len(phrase) # to not end up outside array
	    for i in range(lenarray):
	        if(msg[i:(i+len(phrase))]) == phrase:
	            num = num + 1
	    return num

	def makecount(self):
		self.count = []
		for char in self.alphabet:
			self.count.append([char,self.numberOfInstances(self.msg,char)])
		self.count.sort(key=itemgetter(1), reverse=True)

	# Print contents of the table sorted by keys with k rows.
	def printTable(self,lst, k):
		output = []
		for i in range(k):
			output.append([])
		for i in range(len(lst)):
			item = lst[i][0] + ": " + str(lst[i][1]) + "	"
			output[(i%k)].append(item)
		for i in range(k):
			print ''.join(output[i])
		print(" ")

	def printCount(self):
		self.printTable(self.count,3)

	# returns a such that a*x = y % 26
	def findA(self,x,y):
		for i in range(26):
			if(((i*x)%26) == y):
				return i
		return 0

	def egcd(self,a, b):
	    if a == 0:
	        return (b, 0, 1)
	    else:
	        g, y, x = self.egcd(b % a, a)
	        return (g, x - (b // a) * y, y)

	def modinv(self, a, m):
	    g, x, y = self.egcd(a, m)
	    if g != 1:
	        raise Exception('modular inverse does not exist')
	    else:
	        return x % m

	def suggestion(self,from1,to1,from2,to2):
		self.char1 = from1.upper()
		self.s1 = to1.upper()
		self.char2 = from2.upper()
		self.s2 = to2.upper()
		if(self.char1 != '' and self.char2 != '' and self.s1 != '' and self.s2 != ''): 
			diffs = self.ce.diff(self.s1,self.s2)
			diffc =  self.ce.diff(self.char1,self.char2)
			a = self.findA(diffc,diffs)
			if(gcd(int(a),26) == 1):
				print "Possibly correct a, proceding with decryption..."
				b = (self.ce.encode(self.s1)[0] - a*self.ce.encode(self.char1)[0])%26
				self.decrypt(a,b)
			else:
				print "Does not give an invertible a."

	def finished(self,switch):
		if switch == 1:
			print("Suggestion gives a non-invertible a, make new guess")
			self.proceed = "0"
		elif switch == 2:
			print("Invalid suggestions, try again...")
			self.proceed = "0"
		elif switch == 3:
			ans = raw_input("Was suggestions correct? (0/1): ")
			if ans == "1":
				exit()
			else:
				self.proceed = "0"


	def decrypt(self,a,b):		
		if(gcd(int(a),26) == 1):
			ainv = self.modinv(int(a),26)
			print("a = "+str(a)+", b = " + str(b) + ", inverse of a: " + str(ainv))
			nums = self.ce.encode(list(self.msg))
			for i in range(len(nums)):
				nums[i] = ((nums[i] - b)*ainv)%26
			plaintext = self.ce.decode(nums)	
			print ''.join(plaintext)
		else:
			print "Non-invertible a, decryption not performed."
			

	def encrypt(self,a,b):
		nums = self.ce.encode(list(self.msg))
		for i in range(len(nums)):
			nums[i] = (a*nums[i] + b)%26
		cipher = self.ce.decode(nums)
		print ''.join(cipher)
		


print """
Preliminary: Put a file called text.txt in the same directory
as the affine.py file.

Usage: python -i affine.py
   or: ./affine.py (if you have set the executable flag)
	
The file provides the class Affine, which you can use by typing

	cs = Affine(<filename>)
	
Where filename is name of the file containing either the message or 
ciphertext you want to work with. If no input is given, text.txt will be used.
Note that a file with the given filename input as name must sit in the same 
directory as the affine.py file.

It contains
the following methods:

	cryptAnalysis(): Makes a count of each letter and prints the result.
	encrypt(): Encrypts the message from <filename> with the given values 
			for a and b.
	decrypt(a,b): Decrypts the ciphertext from <filename> with the given
		  	values for a and b (give a, not the inverse used to 
		  	decrypt, which is calculated).
	suggestion(f1,t1,f2,t2): Assumes that when encryption was done, the 
			letter from1 went to to1 and from2 went to to2. This is used 
			for trying to calculate the key, and if possible decryption.
	               
Example
	cs = Affine()
	cs.cryptAnalysis()
	cs.suggestion("e","k","t","h")
	cs.decrypt(4,5)
	cs.encrypt(4,5)

"""