import bpy
import numpy as np


def RK4(x, v, n, h, F):

    for i in range(n):  # written for readability, not speed

        kv1 = F(x[:, i])
        kx1 = v[:, i]

        kv2 = F(x[:, i] + kx1 * (h / 2.0))
        kx2 = v[:, i] + kv1 * (h / 2.0)

        kv3 = F(x[:, i] + kx2 * (h / 2.0))
        kx3 = v[:, i] + kv2 * (h / 2.0)

        kv4 = F(x[:, i] + kx3 * h)
        kx4 = v[:, i] + kv3 * h

        v[:, i + 1] = v[:, i] + (h / 6.0) * (kv1 + 2. * (kv2 + kv3) + kv4)
        x[:, i + 1] = x[:, i] + (h / 6.0) * (kx1 + 2. * (kx2 + kx3) + kx4)


def acc(x):
    """ acceleration due to the sun's gravity (NumPy version) """
    return -Gm * x / (((x**2).sum(axis=1)**1.5)[:, None])


def make_from_pydata(named, verts, edges):
    profile_mesh = bpy.data.meshes.new(named)
    profile_mesh.from_pydata(verts, edges, [])
    profile_mesh.update()
    profile_object = bpy.data.objects.new(named, profile_mesh)
    bpy.context.scene.objects.link(profile_object)


# m^3 s^-2   (Wikipedia Standard Gravitational Parameter)
Gm = 1.3271244002E+20
t_year = 31558464.         # s  (roughly)
scale = 4.0E-11    # blender units per meter

n_frames = 150
Dt = t_year / float(n_frames)   # time step

n = 8

X = np.zeros((n, 1000, 3))
V = np.zeros((n, 1000, 3))
T = np.zeros((n, 1000))

tilt = 30. * (np.pi / 180.)  # radians
sin_tilt, cos_tilt = np.sin(tilt), np.cos(tilt)

# start in the same place...
X[:, 0] = np.array([1.5E+11, 0.0, 0.0])[None, :]

V[:, 0] = 29300. * (np.linspace(0.5, 1.2, n)[:, None] *
                    np.array([0.0, cos_tilt, sin_tilt])[None, :])  # ...but different initial velocities

# NOTE!! This is just for quickie demos only.
# Will give wrong result if step size too big.
RK4(X, V, n_frames, Dt, acc)  # pre-calculate orbits

p_size, s_size = 0.2, 0.5

# Create the Universe
ok = bpy.ops.mesh.primitive_ico_sphere_add(size=s_size, location=(0.0, 0.0, 0.0))
sun = bpy.context.active_object
sun.name = "Sun"

n_echos, frames_per_echo = 20, 1
e_sizes = np.linspace(p_size, 0, n_echos + 1)[:-1]

paths = {}


for i in range(n):

    paths[i] = []
    for i_frame in range(n_frames):
        location = scale * X[i, i_frame]
        paths[i].append(location)
    for iecho in range(n_echos):
        i_frame_echo = max(0, i_frame - frames_per_echo * (iecho + 1))
        location = scale * X[i, i_frame_echo]
        paths[i].append(location)

for p in sorted(paths.keys()):
    verts = paths[p]
    edges = [(i, i + 1) for i in range(len(verts) - 1)]
    make_from_pydata("p_" + str(p), verts, edges)