import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt
from matplotlib.animation import ArtistAnimation

# Решение задачи многих тел

# Определяем переменную величину
seconds_in_year = 365 * 24 * 60 * 60
seconds= 60*60*12
years = 12
t = np.arange(0, years*seconds_in_year, seconds)

# Определяем функцию для системы диф. уравнений
def move_func(s, t):
    (xsun, v_xsun, ysun, v_ysun,
     xJupiter, v_xJupiter, yJupiter, v_yJupiter,
     x1, v_x1, y1, v_y1,
     x2, v_x2, y2, v_y2,
     x3, v_x3, y3, v_y3,
     x4, v_x4, y4, v_y4,
     x5, v_x5, y5, v_y5) = s

    # Динамика Солнца
    dxdtsun = v_xsun
    dv_xdtsun = (- G * mJupiter * (xsun - xJupiter) / ((xsun - xJupiter)**2 + (ysun - yJupiter)**2)**1.5)

    dydtsun = v_ysun
    dv_ydtsun = (- G * mJupiter * (ysun - yJupiter) / ((xsun - xJupiter)**2 + (ysun - yJupiter)**2)**1.5)

    # Динамика Земли
    dxdtJupiter = v_xJupiter
    dv_xdtJupiter = (- G * msun * (xJupiter-xsun) / ((xJupiter-xsun)**2 + (yJupiter-ysun)**2)**1.5)

    dydtJupiter = v_yJupiter
    dv_ydtJupiter = (- G * msun * (yJupiter-ysun) / ((xJupiter-xsun)**2 + (yJupiter-ysun)**2)**1.5)

    # Динамика первого тела под влиянием Солнца и Земли
    dxdt1 = v_x1
    dv_xdt1 = (- G * msun * (x1-xsun) / ((x1-xsun)**2 + (y1-ysun)**2)**1.5
            - G * mJupiter * (x1 - xJupiter) / ((x1 - xJupiter)**2 + (y1 - yJupiter)**2)**1.5)

    dydt1 = v_y1
    dv_ydt1 = (- G * msun * (y1-ysun) / ((x1-xsun)**2 + (y1-ysun)**2)**1.5
            - G * mJupiter * (y1 - yJupiter) / ((x1 - xJupiter)**2 + (y1 - yJupiter)**2)**1.5)

    # Динамика второго тела под влиянием Солнца и Земли
    dxdt2 = v_x2
    dv_xdt2 = (- G * msun * (x2-xsun) / ((x2-xsun)**2 + (y2-ysun)**2)**1.5
            - G * mJupiter * (x2 - xJupiter) / ((x2 - xJupiter)**2 + (y2 - yJupiter)**2)**1.5)

    dydt2 = v_y2
    dv_ydt2 = (- G * msun * (y2-ysun) / ((x2-xsun)**2 + (y2-ysun)**2)**1.5
            - G * mJupiter * (y2 - yJupiter) / ((x2 - xJupiter)**2 + (y2 - yJupiter)**2)**1.5)

    # Динамика третьего тела под влиянием Солнца и Земли
    dxdt3 = v_x3
    dv_xdt3 = (- G * msun * (x3-xsun) / ((x3-xsun)**2 + (y3-ysun)**2)**1.5
            - G * mJupiter * (x3 - xJupiter) / ((x3 - xJupiter)**2 + (y3 - yJupiter)**2)**1.5)

    dydt3 = v_y3
    dv_ydt3 = (- G * msun * (y3-ysun) / ((x3-xsun)**2 + (y3-ysun)**2)**1.5
            - G * mJupiter * (y3 - yJupiter) / ((x3 - xJupiter)**2 + (y3 - yJupiter)**2)**1.5)

    # Динамика четвёртого тела под влиянием Солнца и Земли
    dxdt4 = v_x4
    dv_xdt4 = (- G * msun * (x4-xsun) / ((x4-xsun)**2 + (y4-ysun)**2)**1.5
            - G * mJupiter * (x4 - xJupiter) / ((x4 - xJupiter)**2 + (y4 - yJupiter)**2)**1.5)

    dydt4 = v_y4
    dv_ydt4 = (- G * msun * (y4-ysun) / ((x4-xsun)**2 + (y4-ysun)**2)**1.5
            - G * mJupiter * (y4 - yJupiter) / ((x4 - xJupiter)**2 + (y4 - yJupiter)**2)**1.5)

    # Динамика пятого тела под влиянием Солнца и Земли
    dxdt5 = v_x5
    dv_xdt5 = (- G * msun * (x5-xsun) / ((x5-xsun)**2 + (y5-ysun)**2)**1.5
            - G * mJupiter * (x5 - xJupiter) / ((x5 - xJupiter)**2 + (y5 - yJupiter)**2)**1.5)

    dydt5 = v_y5
    dv_ydt5 = (- G * msun * (y5-ysun) / ((x5-xsun)**2 + (y5-ysun)**2)**1.5
            - G * mJupiter * (y5 - yJupiter) / ((x5 - xJupiter)**2 + (y5 - yJupiter)**2)**1.5)

    return (dxdtsun, dv_xdtsun, dydtsun, dv_ydtsun,
            dxdtJupiter, dv_xdtJupiter, dydtJupiter, dv_ydtJupiter,
            dxdt1, dv_xdt1, dydt1, dv_ydt1,
            dxdt2, dv_xdt2, dydt2, dv_ydt2,
            dxdt3, dv_xdt3, dydt3, dv_ydt3,
            dxdt4, dv_xdt4, dydt4, dv_ydt4,
            dxdt5, dv_xdt5, dydt5, dv_ydt5)

# Определяем начальные значения и параметры, входящие в систему диф. уравнений
xsun0 = 0
v_xsun0 = 0
ysun0 = 0
v_ysun0 = 0

xJupiter0 = 5.2028*1.496*10**11
v_xJupiter0 = 0
yJupiter0 = 0
v_yJupiter0 = 13070

x10 = 2.386693824*10**11
v_x10 = 0
y10 = 0
v_y10 = 23608.55591

x20 = 1.3178008377*10**12
v_x20 = 0
y20 = 0
v_y20 = 10046.35897

x30 = -1.102645773*10**12
v_x30 = 0
y30 = 0
v_y30 = -10983.72799

x40 = 3.889839700*10**11
v_x40 = -2*17535.24508/3
y40 = 6.737399995*10**11
v_y40 = 17535.24508/3

x50 = 3.889839700*10**11
v_x50 = 2*17535.24508/3
y50 = -6.737399995*10**11
v_y50 = 17535.24508/3

s0 = (xsun0, v_xsun0, ysun0, v_ysun0,
      xJupiter0, v_xJupiter0, yJupiter0, v_yJupiter0,
      x10, v_x10, y10, v_y10,
      x20, v_x20, y20, v_y20,
      x30, v_x30, y30, v_y30,
      x40, v_x40, y40, v_y40,
      x50, v_x50, y50, v_y50)

G = 6.67 * 10**(-11)

msun = 1.989*10**30
mJupiter = 5.972*10**24
m1 = 1
m2 = 1
m3 = 1
m4 = 1
m5 = 1

# Решаем систему диф. уравнений
sol = odeint(move_func, s0, t)

# Строим решение в виде графика и анимируем
fig = plt.figure()
plt.style.use(['dark_background'])
fig = plt.figure(figsize=(12,12))

bodys = []

for i in range(0, len(t), 100):
    bodysun_line, = plt.plot(sol[i, 0], sol[i, 2], 'o', color='yellow',markersize=10)

    bodyJupiter, = plt.plot(sol[:i, 4], sol[:i, 6], '-', color='orange')
    bodyJupiter_line, = plt.plot(sol[i, 4], sol[i, 6], 'o', color='orange', markersize=2)

    body1_line, = plt.plot(sol[i, 8], sol[i, 10], 'o', color='white',markersize=2)

    body2_line, = plt.plot(sol[i, 12], sol[i, 14], 'o', color='white',markersize=2)

    body3_line, = plt.plot(sol[i, 16], sol[i, 18], 'o', color='white',markersize=2)

    body4_line, = plt.plot(sol[i, 20], sol[i, 22], 'o', color='white', markersize=2)

    body5_line, = plt.plot(sol[i, 24], sol[i, 26], 'o', color='white',markersize=2)


    bodys.append([bodysun_line, bodyJupiter, bodyJupiter_line, body1_line, body2_line, body3_line, body4_line, body5_line])

ani = ArtistAnimation(fig, bodys, interval=100)

plt.axis('off')
ani.save('point lagr.gif')
plt.show()