from Numeric import *
from visual import vector

class Agent:
    #Initialize the agent with specified alpha and beta. betaRewardMatrix eliminates a scalar-matrix multiplication at every integration step
    def __init__(self,alpha,beta,rewardMatrix):
        self.alpha = alpha
        self.beta = beta
        self.betaRewardMatrix = beta*rewardMatrix
    #Vectorized version of the agent's equation of motion given by u'=-betaX*Ay-alphaX*u (or v'= -betaY*Bx-alphaY*v)
    def allDot(self,Ay,u):
        return vector(-Ay-self.alpha*u)
    #Iterate one step of 4th order Runge-Kutta
    def RK4(self,y,u,dt):
        Ay = dot(self.betaRewardMatrix,y)
        #print Ay
        k1 = dt * self.allDot(Ay,u)
        k2 = dt * self.allDot(Ay,u + k1 / 2)
        k3 = dt * self.allDot(Ay,u + k2 / 2)
        k4 = dt * self.allDot(Ay,u + k3)
        return (u + (k1 + 2 * k2 + 2 * k3 + k4) / 6)