from visual import *

Balls = [sphere(color = color.blue, vel = vector(0, 0, 0), mass = 1.0) for i in range(13)]

Balls[0].color = color.red
Balls[0].vel = vector(5, 0, 0)
Balls[0].mass = 10.0
Balls[0].radius = 1.5
Balls[0].pos = (-3, 0, 0)

Balls[1].pos = (0, 2, 0)
Balls[2].pos = (3, 2, 0)
Balls[3].pos = (6, 2, 0)
Balls[4].pos = (0, -2, 0)
Balls[5].pos = (3, -2, 0)
Balls[6].pos = (6, -2, 0)
Balls[7].pos = (0, 6, 0)
Balls[8].pos = (3, 6, 0)
Balls[9].pos = (6, 6, 0)
Balls[10].pos = (0, -6, 0)
Balls[11].pos = (3, -6, 0)
Balls[12].pos = (6, -6, 0)

freq = 100
t = 0
dt = 1.0/freq

while t < 20:
    rate(freq)
    for i in range(13):
        for j in range(13):
            if i == j: continue
            if mag(Balls[i].pos - Balls[j].pos) < Balls[i].radius + Balls[j].radius:
                a = mag(Balls[i].vel - Balls[j].vel) ** 2
                b = -2 * dot((Balls[i].pos - Balls[j].pos), (Balls[i].vel - Balls[j].vel))
                c = mag(Balls[i].pos - Balls[j].pos) ** 2 - (Balls[i].radius + Balls[j].radius) ** 2
                delta = b**2 - 4 * a * c
                if a != 0 and delta >= 0:
                    dt_prim = (-b + sqrt(delta)) / (2 * a)
                else:
                    dt_prim = 0
                Balls[i].pos = Balls[i].pos - Balls[i].vel * dt_prim
                Balls[j].pos = Balls[j].pos - Balls[j].vel * dt_prim
                Balls[i].vel = Balls[i].vel - 2 * (Balls[j].mass / (Balls[j].mass + Balls[i].mass)) * dot((Balls[i].vel - Balls[j].vel), ((Balls[i].pos - Balls[j].pos)/mag(Balls[i].pos - Balls[j].pos))) * ((Balls[i].pos - Balls[j].pos)/mag(Balls[i].pos - Balls[j].pos))
                Balls[j].vel = Balls[j].vel + 2 * (Balls[i].mass / (Balls[i].mass + Balls[j].mass)) * dot((Balls[i].vel - Balls[j].vel), ((Balls[i].pos - Balls[j].pos)/mag(Balls[i].pos - Balls[j].pos))) * ((Balls[i].pos - Balls[j].pos)/mag(Balls[i].pos - Balls[j].pos))
                Balls[i].pos = Balls[i].pos + Balls[i].vel * dt_prim
                Balls[j].pos = Balls[j].pos + Balls[j].vel * dt_prim
        Balls[i].pos += Balls[i].vel*dt

    t += dt