%matplotlib inline
import numpy as np
import pylab as plt
from matplotlib import animation
from JSAnimation import IPython_display

swarm_params = { 'Ncarbon': 6, 'Nsulfur': 8, 'Niron': 5,'Nseed': 1 }
color_dict = { 'carbon': 'k', 'sulfur' : 'y', 'iron' : 'g', 'seed' :'r' }
xlims = [-5,5]
ylims = [-5,5]

class agent():
    def __init__(self,id_number,pars):
        self.id = id_number
        self.xpos = pars['x_init']
        self.ypos = pars['y_init']
        self.cat = pars['cat']
        self.fixed = pars['fixed']
        return None
    
    def update_position(self, newx, newy):
        self.xpos = newx
        self.ypos = newy
        return None
    
    def checkBCs(self):
        if self.xpos > xlims[1]:
            self.update_position(-5,self.ypos)
        if self.xpos < xlims[0]:
            self.update_position(5,self.ypos)
        if self.ypos > ylims[1]:
            self.update_position(self.xpos,-5)
        if self.ypos < ylims[0]:
            self.update_position(self.xpos, 5)
    
    def moveRight(self):
        newx = self.xpos + 1
        newy = self.ypos
        self.update_position(newx,newy)
        
    def moveLeft(self):
        newx = self.xpos - 1
        newy = self.ypos
        self.update_position(newx,newy)
        
    def moveUp(self):
        newx = self.xpos
        newy = self.ypos + 1
        self.update_position(newx,newy)
        
    def moveDown(self):
        newx = self.xpos 
        newy = self.ypos - 1
        self.update_position(newx,newy)
    
    
def agentSep(agent1,agent2):
    return np.sqrt((agent1.xpos-agent2.xpos)**2+(agent1.ypos-agent2.ypos)**2)

class swarm():
    
    def __init__(self, pars):
        
        self.members = []
        self.fig = plt.figure()
        self.ax = self.fig.gca()
        
        for i in np.arange(pars['Ncarbon']):
            temp = agent(i,{'x_init':i, 'y_init':0, 'cat': 'carbon', 'fixed': False})
            self.members.append(temp)
        """     
        for i in np.arange(pars['Nsulfur']):
            temp = agent(i,{'x_init':i, 'y_init':1, 'cat': 'sulfur', 'fixed': False})
            self.members.append(temp)
            
        for i in np.arange(pars['Niron']):
            temp = agent(i,{'x_init':i, 'y_init':2, 'cat': 'iron', 'fixed': False})
            self.members.append(temp)
        """    
        for i in np.arange(pars['Nseed']):
            temp = agent(i,{'x_init':i, 'y_init':3, 'cat': 'seed', 'fixed': True})
            self.members.append(temp)
    
    def checkMove(self,current,newx,newy):
        for other in self.members:
            if newx == other.xpos and newy == other.ypos:
                return 1
        return 0
    
    def timeEvolve(self, N):
        counter = 0
        while counter < N:
            for member in self.members:
                if not member.fixed:
                    test = np.random.rand(1)
                    if test >= 0 and test < 0.25:
                        member.moveLeft()
                    if test >= 0.25 and test < 0.5:
                        member.moveRight()
                    if test >= 0.5 and test < 0.75:
                        member.moveUp()
                    if test >= 0.75 and test <= 1:
                        member.moveDown()
                    member.checkBCs()
                    member.fixed = self.shouldFixQ(member)
            counter += 1
                    
                
    def shouldFixQ(self,member):
        for other in self.members:
            if agentSep(member,other) == 1 and other.fixed:
                return True
        return False
            
    def showCurrentState(self):
        self.ax.cla()
        self.ax.set_xlim(xlims[0],xlims[1])
        self.ax.set_ylim(ylims[0],ylims[1])
        xvec = []
        yvec = []
        colvec = []
        for member in self.members:
            xvec.append(member.xpos)
            yvec.append(member.ypos)
            colvec.append(color_dict[member.cat])
        scat = self.ax.scatter(xvec,yvec,c=colvec,)
        return scat
        
    def animInit(self):
        self.ax.grid(b=True, which='major', color='0.65',linestyle='-')
        
    def timeEvolveAndShowCurrentState(self):
        print '1'
        self.timeEvolve(1)
        print '2'
        scat = self.showCurrentState()
        return scat
        
    def animateSwarm(self):
        anim = animation.FuncAnimation(self.fig, self.timeEvolveAndShowCurrentState, init_func=self.animInit,
                               frames=xrange(200))

classroom = swarm(swarm_params)
classroom.timeEvolve(300)
classroom.showCurrentState()

classroom = swarm(swarm_params)
classroom.animateSwarm()

def main():
    numframes = 100
    numpoints = 10
    color_data = np.random.random((numframes, numpoints))
    x, y, c = np.random.random((3, numpoints))

    fig = plt.figure()
    scat = plt.scatter(x, y, c=c, s=100)

    ani = animation.FuncAnimation(fig, update_plot, frames=xrange(numframes),
                                  fargs=(color_data, scat))
    plt.show()

def update_plot(i, data, scat):
    scat.set_array(data[i])
    return scat,

main()