from World import *
from random import *
from math import *

class RacistWorld(TurtleWorld):
    def step(self):
        """during each time step, each turtle moves once, then
        we print the average level of segregation"""
        for animal in self.animals:
            animal.step()

        segs = [t.seg for t in self.animals]
        avg_seg = sum(segs) / len(segs)
        print avg_seg
        
class RacistTurtle(Turtle):
    """a RacistTurtle isn't really racist; he just gets uncomfortable
    if there aren't enough neighbors like himself.
    """
    def __init__(self, world, color='red'):
        Animal.__init__(self, world)
        self.r = 5
        self.heading = randint(1, 360)
        self.pen = False
        self.color = color
        self.delay = 0
        self.x = randint(-200, 200)
        self.y = randint(-200, 200)
        self.draw()
        world.register(self)

    def fd(self, dist=1):
        Turtle.fd(self, dist)
        self.keep_in_bounds()

    def keep_in_bounds(self):
        """don't let the turtle go off the canvas"""
        width = self.world.canvas.width
        height = self.world.canvas.width

        if self.x > width/2: self.x = width/2
        if self.x < -width/2: self.x = -width/2
        if self.y > height/2: self.y = height/2
        if self.y < -height/2: self.y = -height/2
        self.redraw()

    def step(self):
        self.rt(randint(-60,60))
        if self.is_happy():
            self.fd(1)
        else:
            self.fd(20)

    def is_happy(self, k=5, thresh=0.4):
        """find the closest (k) neighbors and count the fraction that are
        the same color (seg).  If seg>thresh, the turtle is happy.
        """
        others = self.k_closest(k)
        like_me = 0.0
        for other in others:
            if self.color == other.color:
                like_me += 1

        self.seg = like_me / k

        if self.seg >= thresh:
            return True
        else:
            return False
        
    def k_closest(self, k):
        """find the k other turtles closest to this one"""
        others = world.animals
        t = [(self.distance_to(t), t) for t in others if t is not self]
        t.sort()
        u = [other for (d, other) in t[0:k]]
        return u

    def distance_to(self, other):
        """find the distance between this turtle and another"""
        dx = self.x - other.x
        dy = self.y - other.y
        return sqrt(dx**2 + dy**2)


world = RacistWorld()
world.setup_run()
world.delay = 0

# create 100 turtles of each kind
for i in range(100):
    RacistTurtle(world, 'red')
    RacistTurtle(world, 'blue')

world.mainloop()