# Built by Simon Lacoste-Julien
# started 2002/06/2
import Model

###########################
# A Model for a PoolTable
# with a magnetic field
##########################

# Inputs (parameters for now)
#   - B field (in Tesla*kg/Coulomb; renormalized for unit charge and unit mass...)
#               positive direction is x X y
# Parameters:
#   - size of table: x_length anb y_length 9in m)
# State Variables:
#   - position of ball: x and y (in m)
#   - speed of vall: v_x and v_y (in m/s)
# ODE:
#   - dx/dt = v_x
#   - dv_x/dt = B*v_y  (considering unit charge and unit mass)
#   - dy/dt = v_y
#   - dv_y/dt = -B*v_x
#
# State Event Location:
#   - checking collision with wall and doing reflection according to Law of Reflection

# building the subODE class:
class PoolODE(Model.FirstOrderDEs):
    def __init__(self, model, xvar, yvar, v_xvar, v_yvar, B = -1.0):
        Model.FirstOrderDEs.__init__(self, model) # this simply assigns model to our ODE object

        # we now assign the names used for the variables in the array varnames; the order of our labeling
        # HERE will determine the association of each variable with its corresponding derivative function
        # (giving the D.E.)
        # --> the varnames list is associated with the ODE and is used in its step function to call
        # getvars(varnames) in the model class.        
        self.varnames =  [ xvar, yvar, v_xvar, v_yvar]
        self.varderivs = [ self.xdot , self.ydot, self.v_xdot, self.v_ydot ]

        # --- assignment of parameters (they are simply stored as variables in the object)
        self.B = B

    # --- differential equations specification
    # the canonical form of the definition is: derivName(self, x,t)
    def xdot(self, x, t):
        return x[2] # v_x

    def ydot(self, x, t):
        return x[3] # v_y

    def v_xdot(self, x, t):
        return self.B * x[3]  # B*v_y

    def v_ydot(self, x, t):
        return - self.B * x[2]  # -B*v_x 

# building the subModel class:
class PoolTableModel(Model.Model):
    # x = 0, y = 0 is lower south-west corner of table
    # with those initial conditions, you have a nice pattern! ;-)
    def __init__(self, x_init = 1.0, y_init= 1.0, v_xinit = 3.0, v_yinit = -1.0, l_x = 10.0, l_y = 5.0, B = -1.0):
        # -- initial conditions
        self.x_init = x_init
        self.y_init = y_init
        self.v_xinit = v_xinit
        self.v_yinit = v_yinit
        # -- control parameters:
        self.l_x = l_x
        self.l_y = l_y
        self.B = B

        # 'modelvarnames' and 'params' are used for the simulator (GUI environment)
        # their content has to absolutely match the model structure
        self.modelvarnames = ['x', 'y', 'v_x', 'v_y']
        self.params = ['x_init', 'y_init', 'v_xinit', 'v_yinit', 'l_x', 'l_y', 'B']
        # I guess that the params can be changed before a run..., but not after the
        # the beginning of a run (I'm not sure yet, though... 02_6_1)
        
    def start(self):
        self.x = self.x_init
        self.y = self.y_init
        self.v_x= self.v_xinit
        self.v_y = self.v_yinit
        self.diffeqs = PoolODE(self, 'x', 'y', 'v_x', 'v_y', self.B)
        self.zerocrossingfuncs = [(self.testNorth, self.bounceY, '+-'),
                                  (self.testSouth, self.bounceY, '+-'),
                                  (self.testEast, self.bounceX, '+-'),
                                  (self.testWest, self.bounceX, '+-')]
    #-------------------------#
    # Zero-crossing functions #
    #-------------------------#
    def testNorth(self, t):
        return (self.l_y - self.y)  # this should be kept positive

    def testEast(self, t):
        return (self.l_x - self.x) # this should be kept positive

    def testSouth(self, t):
        return (self.y) # this should be kept positive

    def testWest(self, t):
        return (self.x) # this should be kept positive

    #---------------------------#
    # Event handlers definition #
    #---------------------------#
    def bounceY(self):
        self.v_y = - self.v_y

    def bounceX(self):
        self.v_x = - self.v_x


    
# need to declare the global variable 'model' for the simulator:
model = PoolTableModel()