# vim: set ts=3:

# COPYRIGHT (c) 2016, Chris Huebsch <chris@huebsch-gemacht.de>
# Inspired by the amazing game of Holoken
#

import random, argparse

from reportlab.pdfgen.canvas import Canvas
from reportlab.lib.pagesizes import A4
from reportlab.lib.units import mm

class GridCage():
	# Koordinates of other cells in a cage (starting cage "0" *not* included)
	__CAGES__=[
		# #0)
		# 0 x
		[(0,1)],
		# #1)
		# 0
		# x
		[(1,0)],
		# #2)
		# 0 x x
		[(0,1), (0,2)],
		# #3)
		# 0
		# x
		# x
		[(1,0), (2,0)],
		# #4)
		# 0 x
		#   x
		[(0,1), (1,1)],
		# #5)
		#   x
		# 0 x
		[(0,1), (-1,1)],
		# #6)
		# 0
		# x x
		[(1,0), (1,1)],
		# #7)
		# 0 x
		# x
		[(1,0), (0,1)],
		# #8)
		# 0 x
		# x x
		[(1,0), (0,1), (1,1)],
		# #9)
		# 0 x x
		# x
		[(1,0), (0,1), (0,2)],
		# #10)
		#     x
		# 0 x x
		[(0,1), (0,2), (-1,2)],
		# #11)
		# 0 x
		# x
		# x
		[(1,0), (2,0), (0,1)],
		# #12)
		# 0
		# x
		# x x
		[(1,0), (2,0), (2,1)],
		]

	# for each entry of __CAGES__ one classification entry
	# same order as in __CAGES__
	__CAGE_CLASSIFICATION__ = [
		# simple 1x2
		"2", "2",
		# simple 1x3
		"3", "3",
		# around corner 3 cells
		"l", "l", "l", "l",
		# 2x2
		"Q",
		# around corner long leg, 4 cells 
		"L", "L", "L", "L"]

	def __init__(self):
		# No cells in cage
		self.cells = []

	# Add a cell to a cage
	def addCell(self, row, col):
		self.cells.append((row, col))

	# Set value and operation for this cage
	def setValOp(self, value, operation):
		self.value = value
		self.operation = operation

	# Return value + operatin string
	def getValOp(self):
		return "%d%s"%(self.value, self.operation)

	# Checks if cell is in this cage
	def containsCell(self, row, col):
		return (row, col) in self.cells

	# Returns all cells for this cage
	def getCells(self):
		return self.cells

	def __repr__(self):
		return "%d%s: %s"%(self.value, self.operation, str(self.cells))
	

class HoloGen:
	def __init__(self, opts):
		self.opts = opts

		# n is dimension of grid
		self.n = opts.size

		# default number of cages of size 1x1 is n/2
		if opts.onecages == -1:
			self.singleCageCount = self.n/2
		else:
			self.singleCageCount = opts.onecages

		# define a grid of n*n filled with 0s
		self.grid = [0 for i in xrange(self.n**2)]

		# no cages defined yet
		self.cages = []

		# Track wich cell is in which cage
		self.cageForCell={}

	# Generation process
	def generate(self):
		# first position all numbers according to base rule
		self.generateGrid()
		# generate cage layout
		self.generateCages()
		# assign each cage a value and an operation
		self.generateCageValues()
		#print (self.cages)

	# fill grid of n*n elements with n numbers
	# each number exactly once per row and colum
	# random positioning
	def generateGrid(self):
		value = 0
		# position each of the n differen values
		while value < self.n:
			value += 1
			# in every row
			for row in xrange (0, self.n):
				attempts = 0
				column = -1
				# try differnt positions
				while True:
					# at a random position in the row
					column = random.randint(0, self.n-1)
					# max 20 attempts
					attempts += 1
					if attempts == 20:
						break
					# check if cell has already a value
					if self.getCell(row, column) != 0:
						continue
					# check if value already in this column
					if self.valueInColumn(column, value):
						continue
					break
				# check if all attempts have been w/o success
				if attempts == 20:
					# remove all occurences of this value and start over
					self.clearValue(value)
					value -= 1
					break
				# otherwise fix that value into the cell
				self.setCell(row, column, value)

	# generate cages on grid
	# without any holes
	# without assinging a cell to multiple cages
	# randomized
	def generateCages(self):
		# restart if no solution has been found
		restart = True
		attempts = 0
		while restart:
			# optimistic, it did work
			restart = False
			attempts += 1

			# maximum attempts n*n
			if (attempts > self.n**2):
				raise Exception("Cannot create Cages")

			# position single cages (only one cell)
			self.createSingleCages()

			# for each cell on grid
			for row in xrange(self.n):
				for col in xrange(self.n):

					# check if cell is in a cage already
					if self.inCage(row, col):
						# skip
						continue

					# get possible cages for this cell
					possibleCages = self.getPossibleCages(row, col)

					# no cage have been found
					if len(possibleCages) == 0:
						# delete all cages
						self.cages = []
						# start over
						restart = True
						break

					# get one random cage layout
					cageOffsets = random.choice(possibleCages)

					# create new Cage
					cage = GridCage()

					# add base cell
					cage.addCell(row, col)
					#print ("New Cage %d, %d"%(row,col))

					# add all other cells
					for offset in cageOffsets:
						newRow = row+offset[0]
						newCol = col+offset[1]
						#print ("Adding %d, %d"%(newRow,newCol))
						cage.addCell(newRow, newCol)

					# store cage
					self.cages.append(cage)

				# break outer loop for restart
				if restart:
					break

	# check if cell is in any cage
	def inCage(self, row, col):
		# check all cages 
		for cage in self.cages:
			# check if cell is in that cage
			if cage.containsCell(row, col):
				return True
		return False			

	# position single cages in grid
	# create as many single cages
	# select random position, check position and value
	def createSingleCages(self):
		# rows and columns with single cages
		singleRows = []
		singleCols = []
		# values already used in a single cage
		singleVals = []
		for i in xrange(self.singleCageCount):
			attempts = 0
			while True:
				attempts += 1
				if (attempts > 10*self.singleCageCount):
					raise Exception("Cannot create Single Cages")

				# get a random position for that single cage
				row = random.randint(0, self.n-1)
				col = random.randint(0, self.n-1)
				
				# get the value (remember numbers already positioned!)
				val = self.getCell(row, col)

				# create that single cage only if there is no
				# other single cell in that row and column
				# and value is not in a single cage yet
				if not row in singleRows \
					and not col in singleCols \
					and not val in singleVals:
					break

			# otherwise remember position and value
			singleRows.append(row)
			singleCols.append(col)
			singleVals.append(val)

			# create the single cage
			cage = GridCage()
			cage.addCell(row, col)
			self.cages.append(cage)

	# get all possible cages for the given cell
	# if no cages are found, return []
	def getPossibleCages(self, row, col):
		pc = []

		# get all cageoffsets and its list position
		for idx, offsets in enumerate(GridCage.__CAGES__):
			# find out cage type classification
			ct = GridCage.__CAGE_CLASSIFICATION__[idx]
			# if not allowed for this grid - use next
			if not ct in self.opts.cagetypes:
				continue

			# optimistic approach
			valid = True
			#print ("Offsets %s"%str(offsets))

			# check all cells of this cage candidate
			for offset in offsets:
				newRow = row+offset[0]
				newCol = col+offset[1]
				# not outside of grid
				if newRow < 0 or newCol < 0 or newRow >= self.n or newCol >= self.n:
					valid = False
				# not in any other cage yet
				if self.inCage(newRow, newCol):
					valid = False
			# add valid cage to cage candidate list
			if valid:
				#print ("Appending")
				# multiple times according to program option
				for _ in range(self.opts.cagetypes.count(ct)):
					pc.append(offsets)

		return pc

	# generate values and operations for each cage
	# and while iterating all cages, record which cell is in wich cage (needed for plotting)
	# random again
	def generateCageValues(self):
		cageNumber = 0
		for cage in self.cages:
			cageNumber += 1

			# get all cells for this cage
			cellList = cage.getCells()

			# assign each cell of this cage the same number
			for cell in cellList:
				self.cageForCell[cell] = cageNumber

			# if this is a single cage, store its only value without an operation
			if len(cellList) == 1:
				cell  = cellList[0]
				cage.setValOp(self.getCell(cell[0], cell[1]), "")

			# if the cage has two cells, all mathematical operations are allowed
			elif len(cellList) == 2:
				operations = [x for x in self.opts.mathops]

				# shuffle in place
				random.shuffle(operations)

				# find operation for this cell
				# (operation, val1, val2, larger and smaller remain defined after while!)
				for operation in operations:
					# get values in both cells, find the larger and the smaller one
					val1 = self.getCell(cellList[0][0], cellList[0][1])
					val2 = self.getCell(cellList[1][0], cellList[1][1])
					larger = max(val1, val2)
					smaller = min(val1, val2)

					# selected operation was divide
					if operation is "d":
						# check if division is possible (natural number)
						if larger % smaller == 0:
							break
					else:
						# all other operations are ok
						break
				else:
					# for loop came to an end without finding an operation
					raise Exception("No operation can be assigned to a cage.")

				# assign value and operation to the cage
				if operation is "a":
					cage.setValOp(val1+val2, "+")
				elif operation is "s":
					cage.setValOp(larger-smaller, "-")
				elif operation is "m":
					cage.setValOp(val1*val2, "*")
				elif operation is "d":
					cage.setValOp(larger/smaller, "/")
				else:
					raise Exception("Invalid Operation %s."%operation)

			# cages with more than 2 cells can only be plus or multiply
			else:
				# remove all s(ubstracts) and d(ivides) from the operation candidates
				operations = [x for x in self.opts.mathops.replace("s","").replace("d","")]

				# get a random operation
				operation = random.choice(operations)

				# add or multipy all values and store it
				if operation is "a":
					s = 0
					for cell in cellList:
						s += self.getCell(cell[0], cell[1])
					cage.setValOp(s, "+")
				elif operation is "m":
					p = 1
					for cell in cellList:
						p *= self.getCell(cell[0], cell[1])
					cage.setValOp(p, "*")
				else:
					raise Exception("Invalid Operation %s."%operation)
	
	# print an ascii version of the field to stdout (including cage number)
	def dump(self):
		for row in reversed(xrange(self.n)):
			for col in xrange (self.n):
				print ("%2d|%d" % (self.getCell(row, col), self.cageForCell[(row, col)])),
			print

	# creates a PDF with or without solution
	# all changes to the pdf-context are made in this function
	def drawPDF(self, name, solution = False):
		# New A4 Page
		pdf = Canvas(name, pagesize=A4)

		# Move origin to top left corner
		pdf.translate(0*mm, 297.0*mm)

		# Gridsize is n-th part of page width (plus two invisible cells on each side)
		gs = (210.0/(self.n+2)) * mm

		# Move origin down+right one cell
		pdf.translate(gs, -gs)

		# Print title if defined
		if not self.opts.title is None:
			pdf.saveState()
			pdf.setFontSize(gs/3)
			pdf.drawString(0, gs/10.0, self.opts.title)
			pdf.restoreState()

		# Move origin down n cells
		pdf.translate(0, -self.n*gs)

		# Plot basic grid
		pdf.saveState()
		pdf.setLineWidth(0.3*mm)
		pdf.setStrokeColorRGB(0.5,0.5,0.5)	
		self.drawBaseGrid(pdf, gs)
		pdf.restoreState()

		# Plot cages and labels
		pdf.saveState()
		pdf.setFont("Helvetica", gs/4)
		pdf.setLineCap(1)
		pdf.setLineWidth(1*mm)
		self.drawCageBorders(pdf, gs)
		self.drawCageLabels(pdf, gs)
		pdf.restoreState()

		# Plot solution if requested
		if solution:
			pdf.saveState()
			num = 0
			pdf.setFont("Helvetica", gs/2)
			for row in xrange(self.n):
				for col in xrange (self.n):
					num += 1
					pdf.drawCentredString((col+0.5)*gs, (row+0.3)*gs  , \
					"%d"%self.getCell(row, col))
					#"%d"%num)
			pdf.restoreState()

		pdf.showPage()
		pdf.save()

	# basic grid, n*n cells, each gs in size
	def drawBaseGrid(self, pdf, gs):
		for row in xrange(self.n):
			for col in xrange (self.n):
				pdf.rect(col*gs, row*gs, gs, gs, stroke = True, fill = False)

	# draw borders for each cage
	def drawCageBorders(self, pdf, gs):
		# big outside border
		pdf.rect(0, 0, self.n*gs, self.n*gs, stroke=True, fill=False)

		# check each cell
		for row in xrange(self.n):
			for col in xrange (self.n):
				# get cage number for current cell
				currentCage = self.cageForCell[(row, col)]
				# if there is a cell above this one
				if row != self.n-1:
					# get the cage number for that cell above
					northCage = self.cageForCell[(row+1, col)]
					# draw thick line if different cage
					if currentCage != northCage:
						pdf.line((col)*gs, (row+1)*gs, (col+1)*gs, (row+1)*gs)

				# is there a cell right to this one
				if col != self.n-1:
					# get the cage number			
					eastCage = self.cageForCell[(row, col+1)]
					# draw thick line if different cage
					if currentCage != eastCage:
						pdf.line((col+1)*gs, (row)*gs, (col+1)*gs, (row+1)*gs)

	# draw value and operation for each cage in first cell of cage in lower left corner
	def drawCageLabels(self, pdf, gs):
		for cage in self.cages:
			baseCell = cage.getCells()[0]
			pdf.drawString((baseCell[1]+0.05) * gs, (baseCell[0]+0.05) * gs, cage.getValOp())

	# get position of cell in linear array
	def index(self, row, col):
		return (row-1)*self.n+(col-1)

	# get value of cell at row/col
	def getCell(self, row, col):
		return self.grid[self.index(row, col)]

	# set value of cell at row/col
	def setCell(self, row, col, val):
		#print ("Set Cell %d, %d to %d"%(row, col, val))
		self.grid[self.index(row, col)] = val

	# check if value is in any cells of that column
	def valueInColumn(self, column, value):
		for row in xrange(0, self.n):
			if self.getCell(row, column) == value:
				return True
		return False

	# remove value from grid
	def clearValue(self, value):
		for row in xrange(0, self.n):
			for col in xrange (0, self.n):
				if self.getCell(row, col) == value:
					self.setCell(row, col, 0)

if __name__ == "__main__":
	parser = argparse.ArgumentParser(description="Generates Holokens and solutions of arbitrary sizes. More information on http://chu.in-chemnitz.de/programmieren/python/hologen.php")
	parser.add_argument('-s', '--size', default='4', type=int, help="Size of the game. Default 4.")
	parser.add_argument('-f', '--filename', default='out', help="Basename of the PDFs (without .pdf!). Creates two pdfs. <filename>.pdf and solution_<filename>.pdf. Default 'out'.")
	parser.add_argument('-o', '--onecages', default='-1', type=int, help="Number of cages of size 1. (o <= s). Default s/2.")
	parser.add_argument('-m', '--mathops', default="asmd", help="Allowd mathematical operations (a=add, s=sub, m=multipy, d=divide). Use mutiple times to increase probability. Default: asmd. Example use '-m am' for add and multiply only.")
	parser.add_argument('-c', '--cagetypes', default="23QlL", help="Use cages of this type (see webpage). Use mutiple times to increase probability. Default 23QlL")
	parser.add_argument('-t', '--title', default=None, help="Page title.")
	parser.add_argument('-v', '--verbose', default=False, const=True, action='store_const', help="Show some more information during generate.")
	opts = parser.parse_args()
	hologen = HoloGen(opts)
	hologen.generate()
	if opts.verbose:
		hologen.dump()
	hologen.drawPDF(opts.filename+".pdf")
	hologen.drawPDF("solution_"+opts.filename+".pdf", True)