# 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)