# This is the simulation engine for the replicator equations.
# Specifically, I define the ODEs, the integrator, and several
# useful functions.


from Numeric import *

# The replicator equations
class ODE:
    def __init__(self,beta,alpha,R1,R2,R3):
        self.beta = beta
        self.alpha = alpha
        self.R1 = R1
        self.R2 = R2
        self.R3 = R3
    def xdot(self,x,y,z):
        return -1*(self.beta*(self.R1) + self.alpha*x)
    def ydot(self,x,y,z):
        return -1*(self.beta*(self.R2) + self.alpha*y)
    def zdot(self,x,y,z):
        return -1*(self.beta*(self.R3) + self.alpha*z)

# Fourth order Runge-Kutta returns a vector with 3 components
def RK4_3D(f,dt,x,y,z):
    k1x = dt * f.xdot(x,y,z)
    k1y = dt * f.ydot(x,y,z)
    k1z = dt * f.zdot(x,y,z)
    
    k2x = dt * f.xdot(x + k1x/2.0, y + k1y/2.0, z + k1z/2.0)
    k2y = dt * f.ydot(x + k1x/2.0, y + k1y/2.0, z + k1z/2.0)
    k2z = dt * f.zdot(x + k1x/2.0, y + k1y/2.0, z + k1z/2.0)

    k3x = dt * f.xdot(x + k2x/2.0, y + k2y/2.0, z + k2z/2.0)
    k3y = dt * f.ydot(x + k2x/2.0, y + k2y/2.0, z + k2z/2.0)
    k3z = dt * f.zdot(x + k2x/2.0, y + k2y/2.0, z + k2z/2.0)

    k4x = dt * f.xdot(x + k3x, y + k3y, z + k3z)
    k4y = dt * f.ydot(x + k3x, y + k3y, z + k3z)
    k4z = dt * f.zdot(x + k3x, y + k3y, z + k3z)

    x = x + ( k1x + 2.0*k2x + 2.0*k3x + k4x )/6.0
    y = y + ( k1y + 2.0*k2y + 2.0*k3y + k4y )/6.0
    z = z + ( k1z + 2.0*k2z + 2.0*k3z + k4z )/6.0
    
    return array([x,y,z])

# A function used in updating the parameter matrix
# of the replicator equations.

def ParameterUpdate(A,B,x,y): 

        # Find each dot product of the reward with x and y 
        rx = dot(A,y)       # Ay
        ry = dot(B,x)       # Bx

        Rx = dot(rx,array([1,1,1]))     # sum elements together
        Ry = dot(ry,array([1,1,1]))     # sum elements together
                 
        # Find the difference of the rewards with the average
        # From Equation 32
        return rx - Rx,ry - Ry        # Ay - sum(Ay), Bx - sum(Bx)

# This function is used to project the solutions onto the plane

def CoorTrans(x):
    Rot = array([[ -1./2.,1./2.,0.],
	   [ 0,0,sqrt(3.)/2.]])
    return dot(Rot,x)

# This function is the actual integraion. It passes the 
# replicator equations through the integrator.

def MADS_Simulation(ex,ey,alphaX,alphaY,betaX,betaY,u,v,nIterates,dt):

    # Initialize arrays with dummy variables
    x = array([0.,0.,0.])
    y = array([0.,0.,0.])

    XAgent = []
    YAgent = [] 

    # Rewards Matrices
    A = array( [ [2/3*ex, -1 - 1/3*ex, 1 - 1/3*ex],
      [1 - 1/3*ex,2/3*ex,-1 - 1/3*ex],
      [-1 - 1/3*ex,1 - 1/3*ex,2/3*ex] ])

    B = array( [ [2/3*ey, -1 - 1/3*ey, 1 - 1/3*ey],
      [1 - 1/3*ey,2/3*ey,-1 - 1/3*ey],
      [-1 - 1/3*ey,1 - 1/3*ey,2/3*ey] ])
   
    for i in range(nIterates):

        tempRx,tempRy = ParameterUpdate(A,B,x,y)       

        # Make ODE for udot and vdot
        f = ODE(betaX,alphaX,tempRx[0],tempRx[1],tempRx[2])
        g = ODE(betaY,alphaY,tempRy[0],tempRy[1],tempRy[2])
        
        # Integrate for one time step       
        u = RK4_3D(f,dt,u[0],u[1],u[2])
        v = RK4_3D(g,dt,v[0],v[1],v[2])

        # Transform u and v back to x and y to display
        tempx = dot(exp(-u),array([1,1,1]))
        tempy = dot(exp(-v),array([1,1,1]))
        x = exp(-u)/tempx
        y = exp(-v)/tempy

	# Append the points to the solutions
	XAgent.append(x)
	YAgent.append(y)

    return XAgent,YAgent