#FORMAT python

from pyrobot.brain.ga import *

def safe_queen(new_row, new_col, sol):
    for row in range(new_row):
        if (sol[row] == new_col or
            sol[row] + row == new_col + new_row or
            sol[row] - row == new_col - new_row):
            return 0
    return 1

def fitness(sol):
    set = []
    row = 0
    sum = 0
    for col in sol:
        set.append( col )
        if col < 0 or col >= len(sol):
            return 0
        sum += safe_queen( row, col, set)
        row += 1
    return sum

class GAQueens(GA):
    def fitnessFunction(self, genePos):
        return fitness(self.pop.individuals[genePos].genotype)

    def isDone(self):
        best = fitness(self.pop.bestMember.genotype)
        print "Safe queens", best
        return  best == len(self.pop.bestMember.genotype)

    def mutate(self, **args):
        pos = int(random.random() * len(self.genotype))
        self.genotype[pos] =  math.floor(random.random() * \
                                   (self.max - self.min + 1)) + self.min

class MyGene(Gene):
    def display(self):
        for row in range(len(self.genotype)):
            for col in self.genotype:
               if col == row:
                    print "X",
               else:
                    print ".",
            print ""
        print ""

ga = GAQueens(Population(300, MyGene, size = 16, mode = 'integer', max = 16, elitePercent = .1),
              maxGenerations = 200, verbose = 1)
ga.evolve()
ga.pop.bestMember.display()