"""
	Code for use with SAGE.
	
	This file contains code calculate the query set required to eliminate a given polynomial
		
	It is much slower than it could be, primarily because SAGE doesn't allow deletion of elements from
	a Poset -- which means we are creating a new one at every step of the algorithm :-(
	
	This code is released under GPL
	
	Author:
		thomas.dullien@rub.de
	
	Dagstuhl Symmetric Crypto Seminar 2009
"""

import string, sets

#==========================================================================
# Garbage utility functions. Scroll down to the end of the file for the "meat"
#==========================================================================


def __create_huge_polynomial( ring, inputstring ):
	"""
		This function creates a boolean polynomial from a particularly large string. It's
		useful because SAGE dies on very large boolean polynomials (the parse it uses
		seems to recurse, and it runs out of stack space)
	"""
	monomials = inputstring.split("+")
	collect_monomial = ring( monomials[0] )
	for i in range( 1, len( monomials ) ):
		collect_monomial = collect_monomial + ring( monomials[ i ] )
		if i % 10000 == 0:
			print "parsed %d monomials so far..." % i
	return collect_monomial

def __poset_to_gml( poset ):
	"""	Constructs a GML output string from a Poset that can be used to load the graph
		into yEd for drawing, printing and editing. Fantastic if you like to have posters from
		large Hasse Diagrams.
	
	"""
	headerstring = 'Creator	"Sage"\r\nVersion	"0.0"\r\ngraph\r\n[\r\n\thierarchic	1\r\n\tlabel	""\r\n\tdirected	1\r\n\t'
	vertices = poset.hasse_diagram().vertices()
	vertice_to_name_map = dict( zip( range( len( vertices )), poset.linear_extension() ))
	verticestrings = []
	for vertice in vertices:
		verticestrings.append( '\r\n\tnode\r\n\t[\r\n\t\tid	%d\r\n\t\tlabel	"%s"\r\n\t]' % ( vertice, vertice_to_name_map[ vertice ] ) )
	edgestrings = []
	for edge in poset.hasse_diagram().edges():
		edgestrings.append( '\r\n\tedge\r\n\t[\r\n\t\tsource %d\r\n\t\ttarget %d\r\n\t]' % ( edge[1], edge[0] ))
	return "".join( [ headerstring ] + verticestrings + edgestrings + ["\r\n]"] )
	
	
def __boolean_polynomial_to_poset( bp ):
	"""	Builds the divisibility lattice, attempting to save some time on this essentially quadratic operation
		by not testing divisibility if the degrees are lower.
		
		This function should be fed the boolean polynomial formed by just the "tweakable" components of
		the polynomial to eliminate. 
	
		The function returns a tuple
			( PO	-- the poset that is constructed
			cover_dictionary -- a dictionary that maps each monomial to the monomials that "cover" it )
	"""
	cover_dictionary = {}
	degree_dict = {}
	max_deg = 0
	monomials = bp.monomials()
	# Initialize the cover dictionary
	for i in monomials:
		cover_dictionary[ i ] = []
	# sort the monomials by degree
	for i in monomials:
		degree = i.deg()
		if degree_dict.has_key( degree ):
			degree_dict[ degree ].append( i )
		else:
			degree_dict[ degree ] = [ i ]
		if degree > max_deg:
			max_deg = degree
	# Ok, we have the monomials sorted by degree
	for i in range( 1, max_deg + 1 ):
		#print "Processing %d monomials of degree %d" % ( len( degree_dict[i] ), i )
		for monomialA in degree_dict[ i ]:
			for j in range( i+1, max_deg + 1 ):
				#print "   Checking against %d monomials of degree %d" % ( len( degree_dict[j] ), j )
				for monomialB in degree_dict[ j ]:
					if monomialB.reducibleBy( monomialA ):
						cover_dictionary[ monomialA ].append( monomialB )
	PO = Poset( cover_dictionary, cover_relations=False )
	return (PO, cover_dictionary)


def __construct_tweakable_sub_poly( booleanRing, bp, tweak_variables ):
	"""	This constructs a sub-polynomial consisting of all the tweakable monomials in bp. 
		
		The arguments are:
			booleanRing	: The ring over which to work
			bp			: The input polynomial
			tweakvars	: A list of strings describing the tweakable variables
		
		We first construct a polynomial consisting of only the tweakable components
	"""
	temppoly =  booleanRing( "+".join( tweak_variables )) #[ "x%d" % i for i in range( 80, 144 ) ] ) )
	tweakable_monomials = temppoly.monomials()
	#
	#	First retrieve all monomials that contain tweakable variables
	#
	monomials_to_keep = []
	for monomial in bp.monomials():
		for tweakable in tweakable_monomials:
			if monomial.reducibleBy( tweakable ):
				monomials_to_keep.append( monomial )
	temppoly = booleanRing("0")
	#
	#	Now retrieve the tweakable part, create a monomial from it, and add it to monset
	#
	varset = sets.Set()
	monset = sets.Set()
	for mon in monomials_to_keep:	
		str = mon.__repr__()	# Get the string rep for the monomial
		vars = str.split("*")		# Cut it at the *
		tweakvars = []		
		for var in vars:		
			if var in tweak_variables:		# Is the current variable tweakable ?
				varset.add( var )
				tweakvars.append( var )	# Yes ? Then collect it in tweakvars
		monset.add( "*".join( tweakvars ) )
	print "[!] --> Total number of %d tweakable vars in this poly" % len( varset )
	#
	#	We should have all tweakable-only monomials in monset now
	#
	return booleanRing( "+".join( list( monset ) ) )
		
	
def increase_counters( L, t, counter_dictionary ):
	for element in L.lower_covers( t ):
		#print "%s is covered by %s, counter is %d" % (element, t, counter_dictionary[element])
		counter_dictionary[ element.__str__() ] = counter_dictionary[ element.__str__() ] + 1
		#print "Counter is now %d" % counter_dictionary[ element ]
	counter_dictionary[ t.__str__() ] = counter_dictionary[ t.__str__() ] + 1 
	# Done
	
def remove_even_elements_top_down( L, counter_dictionary, cover_dictionary ):
	"""
		In order to remove the elements in question
	"""
	hasse_diagram = L.hasse_diagram()
	# We need a dictionary that maps lattice elements to hasse diagram vertices
	vertices = hasse_diagram.vertices()
	vertice_to_monomial = dict( zip( range( len( vertices )), L.linear_extension() ))
	monomial_to_vertice = dict( zip( L.linear_extension(), range( len( vertices ))))
	
	string_dict = dict( [ (x.__str__(), x) for x in cover_dictionary.keys()])
	templist = [ top for top in L.maximal_elements() if ( counter_dictionary[ top.__str__() ] % 2 == 0 ) ]	
	while len( templist ) > 0:
		# Remove the top elements with even count. We have to do this in a particularly braindead way
		# due to the fact that SAGE doesn't allow me to just remove the elements -- I have to construct 
		# a new poset
		relations = L.cover_relations()
		temporary_cover_dict = dict( [ (element.element, [] )  for element in L if element not in templist ] )
		for relation in relations:
			if relation[ 1 ] not in templist:
				temporary_cover_dict[ relation[0].element ].append( relation[1].element )
		L = Poset( temporary_cover_dict )
		templist = [ top for top in L.maximal_elements() if ( counter_dictionary[ top.__str__() ] % 2 == 0 ) ]
	return L

#==========================================================================
# Interesting stuff starts here
#==========================================================================

def calculate_query_set( bp ):
	"""
		This algorithm attempts to construct a small set of vectors with which to eliminate
		ALL of the input polynomial.
	
		The input polynomial is assumed to have variables x0 to x79 unknown and x80 to x143
		tweakable.
		
			bp is the input polynomial 
	"""
	S = sets.Set()
	(L, cover_dictionary) = __boolean_polynomial_to_poset( bp )
	# Construct the hasse diagram. We need this
	
	counter_dictionary = dict( zip( [ t.__str__() for t in L ], [ 1 for _ in xrange( len(L)) ]))
	while len( L ) > 0:
		T = L.maximal_elements( )
		for t in T:
			S.add( t )
			increase_counters( L, t, counter_dictionary )
		L = remove_even_elements_top_down( L, counter_dictionary, cover_dictionary )		
	return S
		
def run_old():
	variables = [ "x%d" % i for i in range( 200 ) ]
	booleanRing = BooleanPolynomialRing( len(variables), variables )
	lines = file("/home/thomas/Desktop/present.4.rounds.5.plain.25.key.variables.txt", "rt").readlines()

	polys = []
	i = 0
	for polystring in lines:
		print "poly %d" % i
		i = i + 1
		polynomial = __create_huge_polynomial( booleanRing, polystring )
		print "polynomial with %d terms and degree %d" % (len(polynomial), polynomial.degree())
		if polynomial.degree() > 8:
			polys.append( polynomial )
		if len( polys ) > 3:
			return polys
	"""
	# Good polys are for example 17 (small) or 21+ (huge)
	polystring = lines[ 17 ]
	polynomial = __create_huge_polynomial( booleanRing, polystring )
	print "polynomial with %d terms and degree %d" % (len(polynomial), polynomial.degree())
	
	#Write the poly as graph for yed
	print "Constructing lattice for large poly" 
	po, coverdict = __boolean_polynomial_to_poset( polynomial )
	outfile = file("/home/thomas/Desktop/test.gml" , "wt" )
	outfile.write( __poset_to_gml( po ) )
	outfile.close()

	# Now construct the tweakable subpoly
	smallpoly = __construct_tweakable_sub_poly( booleanRing, polynomial, [ "x%d" % i for i in range( 80, 144 ) ]  )
	print "small polynomial with %d terms" % len(smallpoly)
	print smallpoly
	print "Constructing lattice for small poly "
	po2, coverdict2 = __boolean_polynomial_to_poset( smallpoly )
	print po2
	outfile = file("/home/thomas/Desktop/testsmall.gml", "wt" )
	outfile.write( __poset_to_gml( po2 ) )
	outfile.close()
	return po2
	"""
	return polys

def run():
	querysets = []
	# Take the first 8 present equations
	count = 36
	inputs = file("/home/thomas/Desktop/Present.Full.2.Rounds.txt").readlines()[count:64]
	inputs = [ inputs[ i ] for i in range( 0, len(inputs), 4 )]
	for line in inputs:
		tweakable = [ "p%d" % i for i in range(64) ]
		ring = BooleanPolynomialRing( 80+64,  tweakable + [ "k%d" % i for i in range( 80 ) ] )
		bp = __create_huge_polynomial( ring,  line )
		print "[!] Processing PRESENT polynomial %d with %d variables and %d monomials" % (count, len(bp.variables()), len(bp)),
		bp2 = __construct_tweakable_sub_poly( ring, bp, tweakable )
		print "%d tweakable vars, %d tweakable components" % (len([ x for x in bp.variables() if x.__str__() in tweakable]), len(bp2) )
		qs = calculate_query_set( bp2 )
		print "Queryset has %d elements" % len(qs)
		count = count + 1 
		
	
	for i in range( 8, 16 ):
		inputstring = file( "/home/thomas/Desktop/%d.vars.random.txt" % i , "rt" ).read()
		print "="*80
		print "[!] Processing random polynomial with %d variables" % i,
		ring = BooleanPolynomialRing( i, [ "x%d" % j for j in range(i)])
		bp = ring( inputstring )
		print ",%d monomials" % len(bp),
		qs = calculate_query_set( bp )
		print ",queryset is of size %d/%d"  % (len(qs), 2**(i-2))
		querysets.append( qs )
	
	for qs in querysets:
		print qs
		print len(qs)
	
	
	return qs