def collide2d(a, b, e):
    """Returns the velocities of balls a and b after collision. 

    e is the *coefficient of restitution*, a parameter between 0 and 1 
    which determines how much the objects bounce off each other. 
    When e is 0, the two objects completely stick together. 
    When e is 1, the two objects bounce off each other completely. 

    This is a 2-dimensional collision because a and b each have a 2d 
    velocity vector, like two balls moving on a plane. We can separate
    each velocity vector into two components: a normal vector (pointed 
    straight at the other ball) and a tangent vector (perpendicular to
    the normal). The tangent vectors don't participate in the collision
    at all, so we just need to calculate how the normal vectors change
    after the collision. This is easier, since it's a 1-dimensional 
    collision.

    Afterwards we combine the updated normal vector with the tangent 
    vector to get the resulting velocity vector for each ball.
    """
    a_to_b = b.position - a.position
    av_normal = a.velocity.project(a_to_b)
    av_tangent = a.velocity - av_normal
    bv_normal = b.velocity.project(-a_to_b)
    bv_tangent = b.velocity - bv_normal
    av_normal_after, bv_normal_after = collide1d(av_normal, a.mass, bv_normal, b.mass, e)
    av_after = av_normal_after + av_tangent
    bv_after = bv_normal_after + bv_tangent
    return av_after, bv_after

def collide1d(av, am, bv, bm, e):
    """Returns the velicoties of a and b after a 1-dimensional collision, 
    where a and b are headed directly at each other. This formula is from Wikipedia: 
    https://en.wikipedia.org/wiki/Elastic_collision#One-dimensional_Newtonian
    """
    velocity_center_of_mass = (av * am + bv * bm) / (am + bm)
    av_after = velocity_center_of_mass * (1 + e) - av * e
    bv_after = velocity_center_of_mass * (1 + e) - bv * e
    return av_after, bv_after