import random

import core.model
from core.model import DIRECTION_UP, DIRECTION_DOWNRIGHT
from core.event import SimulationStartedEvent, AgentMoveRequest, AgentMovedEvent


class SAgent(core.model.Agent):
    def __init__(self, disp, rate, type):
        core.model.Agent.__init__(self, disp)
        self.disp.unregister(self)
        self.rate = rate
        self.type = type

    def update(self, env):
        count = 0
        friends = 0
        for i in range(DIRECTION_UP, DIRECTION_DOWNRIGHT+1):
            if self.sector.neighbors[i] != None:
                if self.sector.neighbors[i].occupant:
                    count += 1
                    if self.sector.neighbors[i].occupant.type == self.type:
                        friends += 1

        if count != 0:
            sat = float(friends)/count

        # We are not satisfied, let's move
        if count == 0 or sat < self.rate:
            self.teleport(env)

    def teleport(self, env):
        # Build list of free sectors and select one
        free_sectors = []
        for l in env.sectors:
            for c in l:
                if c.occupant == None:
                    free_sectors.append(c)
                    
        if len(free_sectors) == 0:
            print "Unable to move, no free space!"
            return

        
        self.sector.occupant = None
        self.sector = random.choice(free_sectors)
        self.sector.occupant = self
        self.disp.post(AgentMovedEvent(self))
        

        

class SchellingSimulation(core.model.Simulation):
    def __init__(self, disp, env, a1_count, a2_count, a1_rate, a2_rate):
        core.model.Simulation.__init__(self, disp, env)

        self.agents = []

        for i in range(0, a1_count):
            self.agents.append(SAgent(disp, a1_rate, "A"))

        for i in range(0, a2_count):
            self.agents.append(SAgent(disp, a2_rate, "B"))


    def notify(self, e):
        core.model.Simulation.notify(self, e)

        if isinstance(e, AgentMoveRequest):
            for a in self.agents:
                a.update(self.env)

        if isinstance(e, SimulationStartedEvent):
            # Build list of free sectors and select one
            free_sectors = []
            for l in e.sim.env.sectors:
                for c in l:
                    if c.occupant == None:
                        free_sectors.append(c)

            if len(free_sectors) == 0:
                print "Unable to place this agent, no free space!"
                sys.exit(1)

            for a in self.agents:
                choice = random.choice(free_sectors)
                a.place(choice)
                free_sectors.remove(choice)