import math
import matplotlib.pyplot as plt

fig, ax=plt.subplots()
dt=0.001 
t=0

x1=[]
y1=[]
x1n=1 
y1n=3 
vx1=0            
vy1=0
ax1=0
ay1=0

x2=[]
y2=[]
x2n=-2 
y2n=-1 
vx2=0
vy2=0
ax2=0
ay2=0

x3=[]
y3=[]
x3n=1 
y3n=-1 
vx3=0
vy3=0
ax3=0
ay3=0

g=1 
m1=3   
m2=4
m3=5
r1=0
r2=0
r3=0


# Integrate for sixty seconds
while t < 60:  
  t=t+dt
  
  x1.append(x1n)
  y1.append(y1n)
  x2.append(x2n)
  y2.append(y2n)
  x3.append(x3n)
  y3.append(y3n)

  #distances
  r1=((x2n-x1n)**2 + (y2n-y1n)**2)**0.5
  r2=((x3n-x2n)**2 + (y3n-y2n)**2)**0.5
  r3=((x3n-x1n)**2 + (y3n-y1n)**2)**0.5

  # first accelerations
  ax1=g*((m3*(x3n-x1n)/r3**3) + (m2*(x2n-x1n)/r1**3))
  ay1=g*((m3*(y3n-y1n)/r3**3) + (m2*(y2n-y1n)/r1**3))

  ax2=g*((m1*(x1n-x2n)/r1**3) + (m3*(x3n-x2n)/r2**3))
  ay2=g*((m1*(y1n-y2n)/r1**3) + (m3*(y3n-y2n)/r2**3))

  ax3=g*((m1*(x1n-x3n)/r3**3) + (m2*(x2n-x3n)/r2**3))
  ay3=g*((m1*(y1n-y3n)/r3**3) + (m2*(y2n-y3n)/r2**3))

  # Choose a suitable time step based on the acceleration
  a = math.sqrt(max(ax1*ax1 + ay1*ay1, ax2*ax2 + ay2*ay2, ax3*ax3 + ay3*ay3))
  dt = 0.001/a

  # coordinates
  x1n=x1n+vx1*dt+ax1*(dt**2)*0.5
  y1n=y1n+vy1*dt+ay1*(dt**2)*0.5

  x2n=x2n+vx2*dt+ax2*(dt**2)*0.5
  y2n=y2n+vy2*dt+ay2*(dt**2)*0.5

  x3n=x3n+vx3*dt+ax3*(dt**2)*0.5
  y3n=y3n+vy3*dt+ay3*(dt**2)*0.5

  #new distances
  r1=((x2n-x1n)**2 + (y2n-y1n)**2)**0.5
  r2=((x3n-x2n)**2 + (y3n-y2n)**2)**0.5
  r3=((x3n-x1n)**2 + (y3n-y1n)**2)**0.5

  # second accelerations
  ax1n=g*((m3*(x3n-x1n)/r3**3) + (m2*(x2n-x1n)/r1**3))
  vx1=vx1+(ax1+ax1n)*dt*0.5

  ay1n=g*((m3*(y3n-y1n)/r3**3) + (m2*(y2n-y1n)/r1**3))
  vy1=vy1+(ay1+ay1n)*dt*0.5

  ax2n=g*((m1*(x1n-x2n)/r1**3) + (m3*(x3n-x2n)/r2**3))
  vx2=vx2+(ax2+ax2n)*dt*0.5

  ay2n=g*((m1*(y1n-y2n)/r1**3) + (m3*(y3n-y2n)/r2**3))
  vy2=vy2+(ay2+ay2n)*dt*0.5

  ax3n=g*((m1*(x1n-x3n)/r3**3) + (m2*(x2n-x3n)/r2**3))
  vx3=vx3+(ax3+ax3n)*dt*0.5

  ay3n=g*((m1*(y1n-y3n)/r3**3) + (m2*(y2n-y3n)/r2**3))
  vy3=vy3+(ay3+ay3n)*dt*0.5

  # Calculate the total energy
  # KE = 0.5*m1*(vx1**2 + vy1**2) + 0.5*m2*(vx2**2 + vy2**2) + 0.5*m3*(vx3**2 + vy3**2)
  # PE = -m1*m2/r1 - m1*m3/r3 - m2*m3/r2
  # print(f"t = {t}, E = {KE + PE}")

ax.plot(x1,y1, c="black")
ax.plot(x2,y2, c="red")
ax.plot(x3,y3, c="green")
ax.set(xlim=(-6,6), ylim=(-6,6))
ax.set_aspect('equal')
plt.show()