from visual import *

class Agent:
    #each agent carries its own location (some vector),
    #vector field (some vector valued function), integrator (some function), 
    #neighbor (some other agent), and coupling which is a function that
    #takes a position and returns a vector that possibly uses neighbor 
    #information
    def __init__(self,pos,field,integrator,neighbor=None):
	#we need this to keep a copy of the previous position
	#because it is used in the coupled lorenz system
        self.next_pos = pos
        self.pos = pos
        self.field = field
        #integrator needs to take account of
        #the neighbor's field
        self.integrator = integrator
        #neighbor is some other agent
        self.neighbor = neighbor
    #move updates the position of the agent according
    #to the field and its neighbors interference using
    #its integrator and dt
    def move(self,dt):
        self.next_pos = self.integrator(dt,self.pos,self.field,self.neighbor)
        return self
    def update(self):
        self.pos = self.next_pos
        return self

#------------------------Integrators(4th order runge-kutta and euler)----------------------------

#integrator methods for agent class
#all functions are assumed to be vector valued
#so that scalar multiplication and 
#vector addition is well defined

#Runge-Kutta integrator where field
#is a function from vectors to vectors,
#pos is a vector and neighbor is another
#agent with its own field and position
from visual import mag,norm

#this needs to be fixed
def RKIntegrator(dt,pos,field,agent=None):
    if agent == None:
        k1 = dt*field(pos)
        k2 = dt*field(pos+0.5*k1)
        k3 = dt*field(pos+0.5*k2)
        k4 = dt*field(pos+k3)
    else:
        k1 = dt*(field(pos)+someFunc(pos))
        k2 = dt*(field(pos+0.5*k1)+someFunc(pos+0.5*k1))
        k3 = dt*(field(pos+0.5*k2)+someFunc(pos+0.5*k2))
        k4 = dt*(field(pos+k3)+someFunc(pos+k3))
    return pos+1.0/6*(k1+2*k2+2*k3+k4)

#simple euler integrator that implicity assumes
#field returns a vector and accepts vectors as input
#so dt*field(pos) makes sense
def EulerIntegrator(dt,pos,field,coupling,agent=None):
    if agent == None:
        return pos+dt*field(pos)
    else:
        return pos+dt*(field(pos)-coupling*agent.field(agent.pos))

#I need this one because I need to take care of how to change the coupling strength
def NewEulerIntegrator(coupling):
	return lambda dt,pos,field,agent=None: EulerIntegrator(dt,pos,field,coupling,agent)

#---------------------------------helper functions-----------------

#define a function that returns a lorenz
#field with sigma, R and b fixed
def lorenzField(sigma,R,b):
	return lambda v: vector(sigma*(v[1]-v[0]),R*v[0]-v[0]*v[2]-v[1],b*v[2]+v[0]*v[1])