import random


# Value that each row and column must add up to.
target         = 24
# Length and width of the grid used in the puzzle.
gridSize       = 3

# Number bits in the genome (4 for each value 0-15 in each box of the grid).
genomeLength   = (gridSize*gridSize)*4
# Number of genomes per population.
populationSize = 50
# The program will give up after this many generations.
maxGenerations = 1000

crossoverRate  = 0.8
mutationRate   = 0.02

solution = 0
solutionFound = 0

class genome:
    def __init__(self,genes="",fitness=0.0):
        self.genes = genes
        self.fitness = fitness

def nybbleToInt(nybble):
    if   nybble == "0000": return 0
    elif nybble == "0001": return 1
    elif nybble == "0010": return 2
    elif nybble == "0011": return 3
    elif nybble == "0100": return 4
    elif nybble == "0101": return 5
    elif nybble == "0110": return 6
    elif nybble == "0111": return 7
    elif nybble == "1000": return 8
    elif nybble == "1001": return 9
    elif nybble == "1010": return 10
    elif nybble == "1011": return 11
    elif nybble == "1100": return 12
    elif nybble == "1101": return 13
    elif nybble == "1110": return 14
    elif nybble == "1111": return 15
    else: return -1

def createRandomPopulation():
    population = []

    while len(population) < populationSize:
        genes = ""
        
        for i in range(genomeLength):
            genes += str(random.randint(0,1))

        population.append(genome(genes))
    return population

def parseGenes(genome):
    values = []

    for pos in range(0,genomeLength,4):
        nybble = genome.genes[pos:pos+4]
        values.append(nybbleToInt(nybble))
    return values

# Calculates and stores the fitness of each genome in the population.
def calculateFitness(population):
    for genome in population:
        
        #Keeps track of total differences.
        total = 0

        values = parseGenes(genome)

        # Calculate the result of adding each box in each row.
        for row in range(gridSize):
            result = 0

            for box in range(gridSize):
                result += values[row*gridSize+box]

            # Add the difference between the target and result for this row.
            total += abs(target-result)

        # Calculate the result of adding each box in each column.
        for col in range(gridSize):
            result = 0

            for box in range(gridSize):
                result += values[col+box*gridSize]

            # Add the difference between the target and result for this column.
            total += abs(target-result)

        # If all results equal the target, a solution is found, so the globals
        # are set accordingly and the function returns immediately.
        if total == 0:
            solution = genome
            solutionFound = 1
            return 

        genome.fitness = 1.0/total

# Selects two random genomes based on fitness.     
def selection(population):
    totalFitness = 0.0

    # Calculate the total fitness in the population.
    for genome in population:
        totalFitness += genome.fitness

    # Place ball randomly between 0 and the total fitness.
    ball = random.random()*totalFitness
    # Keeps track of how far through the pie we are.
    pie = 0.0

    for genome in population:
        # Add this genome's slice to the pie.
        pie += genome.fitness
        # Select this genome if the ball was in its slice of the pie.
        if ball<pie: return genome

# High chance of crossing the genes of the two genomes.
def crossover(mother,father):
    # Crossover rate determines chance of undergoing crossover.
    if random.random() < crossoverRate:
        # Position at which genes are swapped.
        pos = random.randint(0,genomeLength)
        daughter = genome()
        son = genome()

        # Daughter's genes are first part of mother's, second part of father's.
        daughter.genes = mother.genes[0:pos]+father.genes[pos:]

        # Son's genes are first part of father's, second part of mother's.
        son.genes = father.genes[0:pos]+mother.genes[pos:]

        mother = daughter
        father = son

# Low chance of changing some of the genes in the genome.
def mutate(genome):
    for bit in genome.genes:
        if random.random() < mutationRate:
            bit = str((int(bit)+1)%2)

def printSolution(selection):
    values = parseGenes(solution)

    for val in range(values):
        print values[val]
        if val%gridSize == gridSize-1:
            print "\n\n"
        else:
            print "\t"

def go():
    print "Target:%i"%target
    print "Size:%ix%i"%(gridSize,gridSize)
    print "Thinking..."
    
    # First, we create an initial population of genomes with random genes.
    population = createRandomPopulation()

    print len(population)
    # This loop executes the entire genetic algorithm, and only exits when a
    # solution has been found or the maximum number of generations has been
    # exceeded.
    for generation in range(maxGenerations):
        # The fitness of each genome is used in selection, and must be
        # calculated before all else.  If a solution is found, CalculateFitness
        # will set the globals Solution to the solving genome and SolutionFound
        # to true.
        calculateFitness(population)

        if solutionFound: break

        newPopulation = []

        # Genomes from the current population are now selected, crossed,
        # mutated, and then added to the new population until it is full.
        while len(newPopulation) < populationSize:
            genome1 = selection(population)
            genome2 = selection(population)

            crossover(genome1, genome2)

            mutate(genome1)
            mutate(genome2)

            newPopulation.append(genome1)
            newPopulation.append(genome2)

        # Copy new population into current one.
        population = newPopulation[:]

    if solutionFound:
        print "Solution found in %i generations."%generation
        printSolution(solution)
    else:
        print "No solution found in %i generations."%generation

    raw_input("press Enter")
go()