from pyrobot.brain.ga import *
from pyrobot.brain.conx import *

class NNGA(GA):
    def __init__(self, cnt):
        n = Network()
        n.add( Layer('input', 2) )
        n.add( Layer('hidden', 3) )
        n.add( Layer('output', 1) )
        n.connect('input', 'hidden')
        n.connect('hidden','output')
        n.setInputs([[0.0, 0.0],
                     [0.0, 1.0],
                     [1.0, 0.0],
                     [1.0, 1.0]])
        n.setOutputs([[0.0],
                      [1.0],
                      [1.0],
                      [0.0]])
        n.setVerbosity(0)
        n.setTolerance(.4)
        n.setLearning(0)
        g = n.arrayify()
        self.network = n
        GA.__init__(self, 
                    Population(cnt, Gene, size=len(g), verbose=1,
                               min=-10, max=10, elitePercent = .1),
                    mutationRate = 0.5, crossoverRate = 0.25,
                    maxGeneration = 400, verbose = 1)
    def fitnessFunction(self, genePos):
        self.network.unArrayify(self.pop.individuals[genePos].genotype)
        error, correct, count, pcorrect = self.network.sweep()
        return 4 - error
    def isDone(self):
        self.network.unArrayify(self.pop.bestMember.genotype)
        error, correct, count, pcorrect = self.network.sweep()
        print "Correct:", correct
        return correct == 4

ga = NNGA(20)
ga.evolve()
ga.network.unArrayify(ga.pop.bestMember.genotype)
ga.network.setInteractive(1)
ga.network.sweep()