#FORMAT python

# Evolve a robot that chases after a particular color.  The only object
# of that color is another robot which is running a simple avoid obstacles
# brain.  The evolving robot is called adapter, the other robot is called
# opponent.

from pyrobot.engine import Engine
from pyrobot.brain.ga import *
from pyrobot.brain.conx import *
from pyrobot.system.config import *
import time
import os

class chaseGA(GA):
    def __init__(self, cnt, filename):
        self.filename = filename
        self.popsize = cnt
        self.output = open(self.filename + ".fit", "w")
        config = Configuration()
        config.put("pyrobot", "gui", "tk")  # or tty
        self.adapter = Engine("Player6665",
                              "NNFindBlob",
                              "StageSimulator",
                              config = config,
                              worldfile = "chase.cfg")
        self.opponent =Engine("Player6666",
                              "BBWander")
        # Wait for them to get initialized
        time.sleep(1)

        # Create a dummy network to determine the genome length
        n = Network()
        n.addLayers(8,6,2)
        g = n.arrayify()

        # Initial genes will have values between -10 and 10.
        GA.__init__(self,
                    Population( cnt, Gene, min = -10, max = 10, size = len(g),
                                elitePercent=0.1, verbose = 1),
                    mutationRate = 0.1,
                    crossoverRate = 0.2,
                    maxGeneration = 10,
                    verbose = 1)

    # Fitness at each step is the product of the color blob's range,
    # the blob's location, and the robot's speed.  If no blob is
    # visible or if the robot is blocked from moving, then its fitness
    # is 0 for those steps.
    def fitnessFunction(self, genePos):
        total = 0
        # Reset the robots back to their starting positions.
        self.adapter.robot.simulation[0].setPose('GreenRobot',-0.575,-1.247,45)
        self.opponent.robot.simulation[0].setPose('RedRobot',0.194,1.236, 315)
        self.adapter.brain.net.unArrayify(self.pop.individuals[genePos].genotype)
        time.sleep(1)
        self.adapter.pleaseRun()
        self.opponent.pleaseRun()
        # inputs 0:Bsonar 1:Lsonar 2:Fsonar 3:Rsonar
        # 4:blobRange 5:blobToLeft 6:blobCentered 7:blobToRight
        location = 0
        for i in range(250):
            time.sleep(0.1) # allow the robot to move
            inputs = self.adapter.brain.currentInputs
            translate, rotate = self.adapter.brain.getOutputs()
            if self.adapter.brain.robot.stall:
                translate = 0
            print "moving robot", translate, rotate
            if inputs[0]>0.92 or inputs[1]>0.92 or inputs[2]>0.92 or inputs[3]>0.92:
                notblocked = 0
            else:
                notblocked = 1
            if inputs[6]:
                location = 1.0
            elif inputs[5] or inputs[7]:
                location = 0.5
            else:
                location = 0
            current = notblocked * abs(translate) * location * inputs[4]
            #print current
            total += current
        self.adapter.pleaseStop()
        self.opponent.pleaseStop()
        self.output.write("Gene " + str(genePos) + " fitness " + str(total) + "\n")
        #print "**** TOTAL FITNESS ****", total
        return total

    def isDone(self):
        self.output.write("Generation " + str(self.generation) + "\n")
        self.output.write("Average fitness " + str(self.pop.avgFitness) +"\n")
        self.output.write("Best Fitness " + str(self.pop.bestMember.fitness) +"\n")
        self.output.write("-------------------------------------------------\n")
        self.adapter.brain.net.unArrayify(self.pop.bestMember)
        name = self.filename + "-gen" + str(self.generation) + ".wts"
        self.adapter.brain.net.saveWeightsToFile(name)
        if self.generation >= self.maxGeneration:
            self.adapter.shutdown()
            self.opponent.shutdown()
            self.output.close()
            return 1
        else:
            return 0

ga = chaseGA(10, "try")
ga.evolve()
ga.adapter.shutdown()
ga.opponent.shutdown()