#!/usr/bin/env python -i

from math import sqrt
from math import ceil

def shank(k, n, a, b):
	# Solves a discrete log problem using Shank's algorithm.
	# (also known as baby-step, giant-step)

	# a and b are assumed to be members of the field k.
	# n is the order of a in k.

	m = int(ceil(sqrt(n)))
	
	L1 = []
	for j in range(0, m):
		L1.append((j, a**(m*j)))
	L1 = sorted(L1, key=lambda x: abs(x[1]))
	
	L2 = []
	for i in range(0, m):
		L2.append((i, b * a**(-i)))
	L2 = sorted(L2, key=lambda x: abs(x[1]))
		
	found = False
	i = 0
	j = 0
	
	while (not found) and (i < m) and (j < m):
		if L1[j][1] == L2[i][1]:
			found = True
		elif abs(L1[j][1]) > abs(L2[i][1]):
			i = i + 1
		else:
			j = j + 1
			
	if found:
		return m*L1[j][0]+L2[i][0] % n
	
def pi1(k, x, n):
	# The projection to the other subgroup than the one with order n
	
	# pow is ahelluvalot quicker than my own exponentiation, and that
	# is really making a huge difference here.
	return k(pow(abs(x), n, k.n))

def ph1(k, n, g, x, Q, C):
	# For each factor, push the element into a subgroup by multiplying
	# by all the other factors.
	D = []
	for q, c in zip(Q, C):
		print "Pohlig-Hellman in the subgroup", q, c
		m = 1 # n/(q**c)
		for a, b in zip(Q, C):
			if q != a:
				m = m * a**b
		h = pi1(k, g, m)
		y = pi1(k, x, m)
		d = ph2(k, q, c, h, y)
		D.append(d)
		
	n = [q**c for q, c in zip(Q, C)]
	return crt(D, n)

def pi2(k, l, r, i, y):
	a = l**(r-i-1)
	# We just assume that y is an element in the correct group
	return k(pow(abs(y), a, k.n)) #  y**a

def ph2(k, l, r, g, x):
	# Here, we push the elements further down into smaller subgroups
	# of prime order p instead of p to some power.
	# Finally, we use Shank's algorithm to solve the problem there.
	y = x
	a = []
	for i in range(0, r):
		d = shank(k, l, pi2(k, l, r, 0, g), pi2(k, l, r, i, y))
		a.append(d)
		e = -(l**i) * d
		y = y*g**e
	
	sum = 0
	for i in range(0, r):
		sum = sum + a[i]*(l**i) 
		
	return sum
		
def crt(a, n):
	# Chinese remainder theorem
	# Both a and n are lists.
	tb = ToolBox()
	l = len(a)
	
	m = 1
	for p in n:
		m = m*p
	
	
	sum = 0
	for ai, p in zip(a, n):
		si = tb.modinv(m/p, p)
		sum = (sum + ai*si*m/p) % m
		
	return sum
		
class ToolBox(object):
	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 for '+str(a)+' modulo '+str(m)+
	        	' does not exist')
	    else:
	        return x % m

def FiniteRing(n):
	class RingElement(object):
		def __init__(self, a):
			self.a = a % n
			
		def __add__(self, other):
			if isinstance(other, int):
				return RingElement(self.a + other)
			elif isinstance(other, RingElement):
				return RingElement(self.a + other.a)
			else:
				raise TypeError('Incompatible types')
			
		def __radd__(self, other):
			if isinstance(other, int):
				return RingElement(self.a + other)
			elif isinstance(other, RingElement):
				return RingElement(self.a + other.a)
			else:
				raise TypeError('Incompatible types')
			
		def __sub__(self, other): 
			if isinstance(other, int):
				return RingElement(self.a - other)
			elif isinstance(other, RingElement):
				return RingElement(self.a - other.a)
			else:
				raise TypeError('Incompatible types')
			
		def __rsub__(self, other): 
			if isinstance(other, int):
				return RingElement(self.a - other)
			elif isinstance(other, RingElement):
				return RingElement(self.a - other.a)
			else:
				raise TypeError('Incompatible types')
			
		def __mul__(self, other): 
			if isinstance(other, int):
				return RingElement(self.a * other)
			elif isinstance(other, RingElement):
				return RingElement(self.a * other.a)
			else:
				raise TypeError('Incompatible types')
			
		def __rmul__(self, other): 
			if isinstance(other, int):
				return RingElement(self.a * other)
			elif isinstance(other, RingElement):
				return RingElement(self.a * other.a)
			else:
				raise TypeError('Incompatible types')
		
		def __truediv__(self, other):
			if isinstance(other, RingElement):
				tb = ToolBox()
				try:
					f = self * RingElement(tb.modinv(other.a, n))
					return f
				except:
					raise Exception('Division not well-defined, ' + str(other) + ' is not invertible in ' + self.__name__)
			else:
				raise TypeError('Incompatible types')
		
		def __div__(self, other): 
			if isinstance(other, RingElement):
				tb = ToolBox()
				try:
					f = self * RingElement(tb.modinv(other.a, n))
					return f
				except:
					raise Exception('Division not well-defined, ' + str(other) + ' is not invertible in ' + self.__name__)
			else:
				raise TypeError('Incompatible types')
		
		def __neg__(self): 
			return RingElement(-self.a)
		
		def __eq__(self, other): 
			return isinstance(other, RingElement) and self.a == other.a
		
		def __abs__(self): 
			return abs(self.a)
		
		def __str__(self): 
			return str(self.a)
		
		def __repr__(self): 
			return '%d (mod %d)' % (self.a, self.n)
			
		def __pow__(self, e):
			# g is the base, e is the exponent, m is the modulus.
		
			# We want to handle all integers e
			if e < 0:
				tb = ToolBox()
				try:
					g = RingElement(1) / self
				except:
					raise Exception('Cannot raise to negative power, ' + repr(self) + ' is not invertible.')
				e = -e
			elif e == 0:
				return RingElement(1)
			else:
				g = self
			
			# Now assume that e > 0.

			# Get binary representation of e, and note that the first two
			# symbols are 0b when e is positive, so we cut them away.
			be = bin(e)[2:]
		
			a = RingElement(1)
			for b in be:
				a = (a * a)
				if b == '1':
					a = (a * g)
			return a
	
	RingElement.n = n
	RingElement.__name__ = 'Z/%d' % (n)
	return RingElement
	
print """
Usage: python -i shank.py
   or: ./affine.py (if you have set the executable flag)
	
The file implements Shank's algorithm and Pohlig-Hellman algorithm
to solve the discrete log problem.

	shank(k, n, a, n): 
		k is a field, as generated by the FiniteRing class
		n is the cardinality of the field
		a is the base
		b is the element with the unknown logarithm
		
	ph1(k, n, g, x, Q, C):
		k is a field, as generated by the FiniteRing class
		n is the cardinality of the field
		g is the base
		x is the element with the unknown logarithm
		Q is a list of prime factors of n-1
		C is a list of prime multiplicities for Q
	
Example
	
	n = 101
	k = AffineCipher(101)
	g = 3
	x = 43
	Q = [2, 5]
	C = [2, 2]
	
	# n-1 = 100 = 10 * 10 = 2^2 * 5*2
"""