# /usr/bin/env python
import matplotlib.pyplot as plt
import matplotlib.animation as ani
from math import *

def draw_trajectory(ball):
    """
    Draws the trajectory of the ball as a line.
    
    Args:
        ball (Ball): the ball whose trajectory will be drawn
    """

    # set limits for the image
    ax = plt.axes()
    ax.set_aspect('equal')
    ax.set(xlim=[-2,10], ylim=[-2,10])

    # plot horizontal line to denote ground
    ax.plot([-2,10],[0,0],'b-')
    
    # plot trajectory
    ax.plot( ball.x_trajectory, ball.y_trajectory, 'r-' )
    plt.show()
        
    
def draw_motion_blur_image(ball):
    """
    Draws the trajectory of the ball as a series or points.
    
    Args:
        ball (Ball): the ball whose trajectory will be drawn
    """

    # set limits for the image
    ax = plt.axes()
    ax.set_aspect('equal')
    ax.set(xlim=[-2,10], ylim=[-2,10])

    # plot horizontal line to denote ground
    ax.plot([-2,10],[0,0],'b-')
    
    # plot the ball as a circle
    ax.plot( ball.x_trajectory, ball.y_trajectory, 'ro', markersize=6, alpha=0.3 )
    plt.show()


def draw_animation_frame(frame, xs, ys):
    """
    Draws the system at a certain time. Used for animation.
    
    Args:
        frame (int): index of the time step to draw
        xs (list): array of ball x-coordinates
        ys (list): array of ball y-coordinates
    """

    # clear the image
    plt.clf()
    
    # set limits for the image
    ax = plt.axes()
    ax.set_aspect('equal')
    ax.set(xlim=[-2,10], ylim=[-2,10])

    # plot horizontal line to denote ground
    ax.plot([-2,10],[0,0],'b-')
    
    # plot the ball as a circle
    ax.plot( xs[frame], ys[frame], 'ro', markersize=6 )


def animate(ball):
    """
    Animate the motion of the ball.
    
    Args:
        ball (Ball): the ball whose trajectory will be drawn
    """

    nframes = len(ball.x_trajectory)
    print("animating "+str(nframes)+" frames")
    fig = plt.figure()
    motion = ani.FuncAnimation(fig, draw_animation_frame, nframes, interval=10, 
                               fargs=(ball.x_trajectory, ball.y_trajectory) )
    plt.show(block=True)
    plt.close(fig)



class Ball:
    """
    A center-of-mass model for a ball.
    
    The model does not include rotation of the ball.
    
    Args:
        x (float): initial x-coordinate
        y (float): initial y-coordinate
        vx (float): initial velocity x-component
        vy (float): initial velocity y-component
        mass (float): ball mass
        radius (float): ball radius, needed for bouncing
    """

    def __init__(self, x, y, vx, vy, mass = 1.0, radius = 0.1):
        # position of the ball
        self.x = x 
        self.y = y
        
        # velocity of the ball
        self.vx = vx
        self.vy = vy
        
        # force acting on the ball
        self.fx = 0.0
        self.fy = 0.0
        
        # mass of the ball
        self.mass = mass
        
        # radius of the ball
        self.radius = radius
        
        # lists for storing the coordinates of the ball as time advances
        self.x_trajectory = [ self.x ]
        self.y_trajectory = [ self.y ]


    def move(self, dt):
        """
        Calculate the position of the ball after some has passed.
        
        The new position is calculated by
        
        .. math ::
             x(t + \Delta t) = x(t) + v_x(t) \Delta t + \\frac{1}{2} \\frac{F_x(t)}{m} (\Delta t)^2
             
             y(t + \Delta t) = x(t) + v_y(t) \Delta t + \\frac{1}{2} \\frac{F_y(t)}{m} (\Delta t)^2.

        The new position is saved in the properties of the ball.
        
        Args:
            dt (float): timestep :math:`\Delta t`
        """
        self.x += self.vx*dt + 0.5*self.fx/self.mass*dt*dt
        self.y += self.vy*dt + 0.5*self.fy/self.mass*dt*dt


    def accelerate(self, dt):
        """
        Calculate the velocity of the ball after some time has passed.
        
        The new velocity is calculated by
        
        .. math ::
            v_x(t + \Delta t) = v_x(t) + \\frac{F_x(t)}{m} \Delta t
            
            v_y(t + \Delta t) = v_y(t) + \\frac{F_y(t)}{m} \Delta t.
        
        The new velocity is saved in the properties of the ball.
        
        Args:
            dt (float): timestep :math:`\Delta t`
        """
        self.vx += self.fx/self.mass*dt
        self.vy += self.fy/self.mass*dt


    def reset_forces(self):
        """
        Forget the forces from a previous time.
        """
        self.fx = 0.0
        self.fy = 0.0


    def apply_gravity(self, g):
        """
        Make gravity act on the ball.
        
        Adds the force :math:`F_y = - m g`.
        
        Args:
            g (float): gravitational acceleration :math:`g`
        """
        self.fy += -self.mass * g


    def apply_air_resistance(self, drag):
        """
        Make air resistance act on the ball.
        
        Adds the force :math:`\\vec{F} = - \\gamma m \\vec{v}`.
        
        Args:
            drag (float): air drag coefficient :math:`\\gamma`
        """
        self.fx -= drag * self.mass * self.vx
        self.fy -= drag * self.mass * self.vy
        
        
    def apply_bounce(self, restitution):
        """
        Make the ball bounce from the ground (ground level is at y = 0).
        
        Args:
            restitution (float): Coefficient for the elasticity of the bounce. 
                     1 = completely elastic, 0 = completely inelastic.
        """
        if self.y < self.radius:
            self.vy = -restitution*self.vy
            self.y = -self.y + 2*self.radius
       
        
    def record_position(self):
        """
        Save the position of the ball in memory.
        """
        self.x_trajectory.append( self.x )
        self.y_trajectory.append( self.y )


def run_simulation(ball, g, drag, restitution, dt, simulation_time, recording_dt):
    """
    Run a dynamic simulation.
    
    The simulation runs repeating the following steps:
    
        * Forces are applied on the ball using e.g. :meth:`Ball.apply_gravity`.
        * The ball is moved with :meth:`Ball.move`.
        * The velocity is updated with :meth:`Ball.accelerate`.
        
    The trajectory is saved in the :class:`Ball` object.
    
    .. note ::
        This function is incomplete!
    
    Args:
        ball (Ball): the moving ball
        g (float): acceleration due to gravity
        drag (float): air resistance factor (between 0 and 1)
        restitution (float): impact elasticity factor (between 0 and 1)
        dt (float): time step :math:`\\Delta t`
        simulation_time (float): total simulation time
        recording_dt (float): time between trajectory recording
    """

    steps = int(simulation_time / dt)
    recording_interval = int(recording_dt / dt)

    for step in range(1,steps):
    
        #
        # ToDo: implement the simulation
        #
        ball.reset_forces()

    
        # record the position every so often
        if step%recording_interval == 0:
            ball.record_position()
    

def main(v_start, angle):
    """
    Main program.
    
    Creates a ball, throws it at an angle and simulates the throw.
    
    Args:
        v_start (float): initial speed
        angle (float): throwing angle in degrees, measured from the horizontal
    """

    # starting position
    x_start = 0.0
    y_start = 1.5

    # starting velocity
    degree = pi/180
    throw_angle = angle*degree
    vx_start = v_start*cos(throw_angle)
    vy_start = v_start*sin(throw_angle)

    # physical constants:
    # changing these will change the physics of the system

    # gravity
    g = 9.8

    # air resistance (between 0 and 1)
    drag = 0.2

    # bounce elasticity (between 0 and 1)
    restitution = 0.7

    # simulation time step
    dt = 0.01

    # create the ball
    ball = Ball(x_start, y_start, vx_start, vy_start)

    # how often do you want to save data for plotting?
    recording_dt = 0.1

    # how long do you want to simulate?
    simulation_time = 20.0

    # run the simulation
    run_simulation(ball, g, drag, restitution, dt, simulation_time, recording_dt)
    
    # visualize the simulation
    draw_trajectory(ball)
    draw_motion_blur_image(ball)
    animate(ball)


# Run this code if the file is executed as a program 
# but not if it's loaded as a module.
if __name__ == "__main__":
    v = 8.0
    angle = 75 # degrees
    main(v, angle)
    