import Graphics
import Myro
import math
from ai.ga import GA, Population, Gene
from ai.conx import SRN

Myro.init("sim")
robot = Myro.getRobot()
Myro.setOption("show-sensors", True)

Myro.getSimulation()

pic = robot.takePicture()
pic

Myro.joystick()
Graphics.getWindow()

pic = robot.takePicture()
pic

getImage(pic)

## Evolving a Robot Controller
## Calico Python
## D.S. Blank

def getImage(pic):
    bins = [0] * 10
    for p in pic.getPixels():
        rgb = p.getRGB()
        if rgb[0] == rgb[1] == 0:
            bins[int(p.x/256 * 10)] += 1
    bins = [v/pic.height for v in bins]
    return bins

class RobotGA(GA):
    def __init__(self, cnt):
        n = SRN()
        n.addThreeLayers(15, 5, 2)
        n.setVerbosity(0)
        n.setTolerance(.1)
        n.setLearning(0)
        n.setSequenceType("ordered-continuous")
        g = n.arrayify()
        self.network = n
        GA.__init__(self,
                    Population(cnt, Gene, size=len(g), verbose=1,
                               min=-10, max=10, maxStep = 1,
                               imin=-10, imax=10, 
                               elitePercent = .1),
                    mutationRate=0.05, crossoverRate=0.6,
                    maxGeneration=100, verbose=1)
    def fitnessFunction(self, genePos):
        robot.setPose(100, 250, 0)
        current = Graphics.Point(100, 250)
        step = 0
        fitness = .01
        step_limit = max(min(self.generation, 30), 5)
        while step < step_limit: # 2.5 steps per second
            self.network.unArrayify(self.pop.individuals[genePos].genotype)
            pic = robot.takePicture()
            Myro.show(pic)
            image = getImage(pic)
            sensors = (robot.getIR() + 
                       [v/6400 for v in robot.getObstacle()] + 
                       [0 if p == 0 else 1 for p in image])
            t, r = self.network.propagate(input=sensors)
            self.network.copyHiddenToContext()
            robot.move(t * 2.0 - 1, r * 2.0 - 1)
            fitness += sum(image)
            current = Graphics.Point(robot.frame.center)
            step += 1
        print(fitness)
        robot.stop()
        return fitness
    def isDone(self):
        return False


robotga = RobotGA(10)

robotga.evolve()

robotga.evolve(cont=1)

robotga.saveGenesToFile("genes3.pickle")

robotga.loadGenesFromFile("genes3.pickle")

robotga.evolve(cont=1)

calico.YouTubeVideo("NbZNO1Lx0hc")

calico.YouTubeVideo("hpLJ7xfLaEQ")