import numpy as np
import matplotlib.pyplot as plt
import time
from IPython import display
from numpy import linalg as mag

#create the bird, pig, and block
Rbird=0.5  #radius of the bird
Rpig=0.5   #radius of the pig
bird = np.array([-15.0,0.0])
pig = np.array([10.0,0.0])

#Set up the scene - don't worry about this part too much!
print("Time of snapshot: 0.0")
plt.axis([-24,24,0,24])   #set x and y axis limits
plt.plot(bird[0],bird[1], marker='o',markersize=20, color='red' ) #plot ball A's position with the appropriate marker
plt.plot(pig[0],pig[1], marker='o', markersize=20,color='green') #plot ball B's position
display.clear_output(wait=True) #clear the previous display
display.display(plt.gcf())  #display what's being plotted currently
plt.clf()   #clear the plot
time.sleep(.5)  #pause updates by 0.5s

#define the initial velocities and acceleration due to gravity
g=9.8
bird_vel=np.array([0.0,0.0])
pig_vel=np.array([0.0,0.0])
acc = np.array([0.0,-g])

#masses
Mbird=1.0
Mpig=2.0

#set the initial time=0 and set the time step
t=0.0
dt = 0.005
frate=10

#damping factor for collisions with floor
dampx=0.99
dampy=0.9

#logical parameters for different cases
birdpigcollided=False

#This "while" loop repeats the commands indented after it while the condition is true
while 1 == 1: #ie. repeats forever, unless something inside the loop causes it to quit
    t=t+dt   #updates the time          

    #update velocity and position of bird    
    bird_vel = bird_vel + acc*dt    #v_f=v_i+a*dt
    bird = bird + bird_vel*dt    #y_f=y_i+v*dt

    #update position of pig
    pig_vel = pig_vel + acc*dt
    pig = pig + pig_vel*dt
    
    #take a snapshot, include "SUCCESS" if you've already hit the pig
    if (round(t,3)%(round(frate*dt,3)) <= 10e-5):
        print("Time of snapshot: "+repr(t))
        plt.axis([-24,24,0,24])
        plt.plot(bird[0],bird[1], marker='o',markersize=20, color='red' ) 
        plt.plot(pig[0],pig[1], marker='o', markersize=20,color='green')
        display.clear_output(wait=True)
        display.display(plt.gcf())
        if birdpigcollided == True:
            print("SUCCESS!")
        plt.clf()
        time.sleep(.5)
    
    #condition for pig and bird to collide
    if abs(mag.norm(bird - pig)) <= Rbird + Rpig:
        if birdpigcollided ==False: #only first time
            pig_vel=2.0*Mbird/(Mbird+Mpig)*bird_vel #elastic collision
            bird_vel=bird_vel-Mpig/Mbird*pig_vel            
            birdpigcollided=True
            print("SUCCESS!")

    #Interaction with floor causes damped bouncing
    if bird[1] <=0.0:
        bird_vel[1] =-dampy*bird_vel[1]
        bird_vel[0]=dampx*bird_vel[0]
    if pig[1] <=0.0:
        pig[1]=0.0
        pig_vel[0]=dampx*pig_vel[0]
    
    if (mag.norm(bird_vel) <= 0.1):
        break #break if the bird is nearly motionless