# Animation of predator prey trajectories in 1D
# Version 1.0  2-19-10
# Plots the trajectory of a lion and a gazelle,
#and then their distance and speed vs. time
#
from visual import *
from visual.graph import *  # import graphing features
#
# Code adapted from Dawn Meredith's exercises at
# University of New Hampshire by Suzanne Amador Kane and Peter Love
# Haverford College 2010
#
# initialization section -- all data except reaction times and ambush distances are from
#    J. P. Elliot et al. "Prey capture by the African lion"
#    Canadian Journal of Zoology vol. 55(11)pp 1811-28 (1977)
#
###############################################################################################
#
#   Student parameters
#
#   Choose prey animal
#
#   preyanimal = 1 is a Gazelle, top speed 27 m/s, K_g = 0.17 s**-1
#   preyanimal = 2 is a Zebra, top speed 16m/s, K_z = 0.31 s**-1
#   preyanimal = 3 is a wildebeest, top speed 14.2 m/s, K_w = 0.39 s**-1
#
preyanimal = 1
#
# Choose prey reaction time
#
# 0.5 seconds is an estimate from University of New Hampshire write up on this topic--note that some gazelles "stott"
# they leap vertically before running, so this may be an underestimate
prey_reaction_time =0.0
#
# 
#
#
# Distance lions are from gazelles at moment of attack in meters
lion_ambush_distance = 30
# total number of timesteps for the simulation
Ntimesteps = 100
#
# timestep in seconds
dt = 0.1
#
#######################################################################################

#####################################################
#
# Lion parameters
# 
# lion initial speed 
v_o = 0.
# lion's maximum velocity is 14 m/s
vmax_lion = 14
#
#####################################################

#####################################################
#
# Prey animal parameters
#
# gazelle's maximum velocity is 27 m/s
vmax_gazelle = 27
# zebra's maximum velocity is 16 m/s
vmax_zebra = 16
# wildebeest's maximum velocity is 14.3 m/s
vmax_wildebeest = 14.3
# human's maximum velocity
vmax_human =10
# lion's rate constant (1/exponential decay time) is 0.68 sec**-1 
K_l = 0.68
# gazelle's rate constant (1/exponential decay time) is 0.17 sec**-1
K_g = 0.17                      
# zebra's rate constant (1/exponential decay time) is 0.31 sec**-1
K_z = 0.31
# wildebeest's rate constant (1/exponential decay time) is 0.31 sec**-1
K_w = 0.39
# humnan's rate constant is 0.666 sec*-1
K_h = 0.6666
#########################################################

if preyanimal == 1:
    #gazelle
    preystring = "Gazelle"
    K_p=K_g
    vmax_prey = vmax_gazelle
elif preyanimal == 2:
    preystring = "Zebra"
    K_p=K_z
    vmax_prey = vmax_zebra
elif preyanimal == 3:
    preystring = "Wildebeest"
    K_p=K_w
    vmax_prey = vmax_wildebeest
elif preyanimal ==4:
    preystring = "Human"
    K_p=K_h
    vmax_prey = vmax_human
# lion initial value of x, y and z
x_o_l = 0                      
y_o_l = -10
z_o_l = 0
# prey initial value of x, y and z--starts at the origin
x_o_p = lion_ambush_distance    
y_o_p = 10
z_o_p = 0
# some additional values to play with for zebra and wildebeest
     
# For plotting, the wireframe dimensions along Cartesian axes
oxa=150                         
oya=30
oza=30
# Initialize graphics scene - put the 3D box at x=50, y=0, make it 800 wide and 200 high
scene=display(width=1200,height=400,x=50, y=0)

scene.ambient=0.24
scene.background=color.white
scene.title="Lion Chasing " + preystring
scene.center=(oxa/2,oya/2,0)
scene.range=(oxa,oya,oza)
#this sets the viewing angle
scene.forward=(0,0,-1)

p1 = (oxa,oya,oza)
p2 = (-0.1*oxa,oya,oza)
p3 = (-0.1*oxa,-oya,oza)
p4 = (oxa,-oya,oza)
p5 = (oxa,oya,-oza)
p6 = (-0.1*oxa,oya,-oza)
p7 = (-0.1*oxa,-oya,-oza)
p8 = (oxa,-oya,-oza)
square1 = curve(pos=[p1,p2,p3,p4,p1],color=color.black)
square2 = curve(pos=[p5,p6,p7,p8,p5],color=color.black)
square3 = curve(pos=[p1,p5,p8,p4,p1],color=color.black)
square4 = curve(pos=[p2,p6,p7,p3,p2],color=color.black)

# initialize lion's position at (xo,yo,zo), pointing along +y 
lion = sphere (pos=(x_o_l,y_o_l,z_o_l), radius=5, color=color.yellow)
lionlabel = label(pos=lion.pos - vector(0,10,0),text='Lion')
# initialize gazelle's position at (xo,yo,zo), pointing along +y
prey = sphere (pos=(x_o_p,y_o_p,z_o_p), radius=5, color=color.red)       
preylabel = label(pos=prey.pos+ vector(0,10,0),text=preystring)

lion_position=[]
lion_speed=[]
prey_position=[]
prey_speed=[]
#
# plots the lion's trajectory
#
c = curve(pos=(x_o_l,y_o_l,z_o_l),color=color.yellow)
#
# plots the prey's trajectory
#
g = curve(pos=(x_o_p,y_o_p,z_o_p),color=color.red)
#
#initialize counter to zero
n=0

#set up the gdisplay options, with width = 1200 and height=600 and x=0 and y=500
#the graph basically covers the bottom half of the screen
graph1 = gdisplay(x=50, y=700, width=1200, height=300,
          title='Graph of position vs. time for Lion chasing '+preystring, xtitle='t (s)', ytitle='Position (m)',
          foreground=color.black, background=color.white)

funct1 = gcurve(gdisplay=graph1,color=color.black) 
funct2 = gcurve(gdisplay=graph1,delta=0.05, color=color.red)


graph2 = gdisplay(x=50, y=400, width=1200, height=300,
          title='Graph of speed vs. time for Lion chasing '+preystring, xtitle='t (s)', ytitle='Speed (m/s)',
          foreground=color.black, background=color.white)

funct3 = gcurve(gdisplay=graph2,color=color.black) 
funct4 = gcurve(gdisplay=graph2,delta=0.05, color=color.red)

while n <= Ntimesteps:
    # screen refresh rate 
    rate(5)                                                             
    
    t = n*dt
    # lion's velocity starts at zero, asymptotes rapidly to vmax
    # first update velocity, then update position second - Euler-Cromer method
    lion.velocity = vector(vmax_lion,0,0)*(1 - exp(-t*K_l)) +vector(v_o,0,0)         
    lion.pos = lion.pos + lion.velocity*dt
    #Only plot the x component
    lion_position.append(lion.pos[0])
    lion_speed.append(mag(lion.velocity))
    c.append(pos=lion.pos)
    lionlabel.pos=lion.pos - vector(0,10,0)

    print lion.pos,prey.pos

    if t < prey_reaction_time:
        prey.velocity = vector(0,0,0)
    else:
        # gazelle's velocity starts at zero, asymptotes rapidly to vmax
        prey.velocity = vector(vmax_prey,0,0)*(1 - exp(-(t-prey_reaction_time)*K_p))
    
    prey.pos = prey.pos + prey.velocity*dt
    #Only plot the x component
    prey_position.append(prey.pos[0])
    prey_speed.append(mag(prey.velocity))
    g.append(pos=prey.pos)
    preylabel.pos = prey.pos + vector(0,10,0)
    
    funct1.plot(pos=(n*dt,lion_position[n]))   
    funct2.plot(pos=(n*dt,prey_position[n]))
    funct3.plot(pos=(n*dt,lion_speed[n]))   
    funct4.plot(pos=(n*dt,prey_speed[n]))

    # step forward in time
    n=n+1                                                   
    
#Now output various relevent facts from the plots
#print "Minimum value of (Gazelle position - Lion position) = ",min(array(gazelle_position)-array(lion_position))