# Function for creating an embedded video given a filename
function html_video(filename)
	base64_video = base64(open(readbytes, filename))
	"""<video controls src="data:video/x-m4v;base64,$base64_video">"""
end

display("text/html", html_video("4Body3D.mp4"))

using PyPlot
using PyCall
using ODE
@pyimport matplotlib.animation as anim
@pyimport mpl_toolkits

function F(t,y)
    
    #Number of Planets (note we're solving the n-body problem, rather than just a 3-body problem)
    n = int(length(y)/6)
    
    #Extract current position and velocity
    r = zeros(n,3)
    v = zeros(n,3)
    for i=1:n
        r[i,:] = y[(i-1)*6+1:(i-1)*6+3]
        v[i,:] = y[(i-1)*6+4:(i-1)*6+6]
    end
        
    #Calculate spatial derivatives
    drdt = v
        
    #Work out velocity derivatives (ie accelerations)
    dvdt = zeros(n,3)
    for i = 1:n
        for j = 1:n
            if i != j
                dvdt[i,:] += -G*m[j]*(r[i,:]-r[j,:])/(norm(r[i,:]-r[j,:])^3)
            end
        end
    end
    
    #Combine back together into dydt vector
    dydt = zeros(6*n)
    for i = 1:n
        dydt[6(i-1)+1] = drdt[i,1]
        dydt[6(i-1)+2] = drdt[i,2]
        dydt[6(i-1)+3] = drdt[i,3]
        dydt[6(i-1)+4] = dvdt[i,1]
        dydt[6(i-1)+5] = dvdt[i,2]
        dydt[6(i-1)+6] = dvdt[i,3]
    end
    return dydt
end

#Set the planetary masses
m = [5,4,3,5]
n = length(m)

#Set the gravitational field strength
G = 1

#Set initial positions and velocities 
r0 = zeros(n,3)
r0[1,:] = [1.0,-1.0,1.0]
r0[2,:] = [1,3,0.0]
r0[3,:] = [-1,-2,0.0]
r0[4,:] = [0,0,0.0]

v0 = zeros(n,3)

#Define y0
y0 = zeros(6*n)
for i = 1:n
    y0[6(i-1)+1] = r0[i,1]
    y0[6(i-1)+2] = r0[i,2]
    y0[6(i-1)+3] = r0[i,3]
    y0[6(i-1)+4] = v0[i,1]
    y0[6(i-1)+5] = v0[i,2]
    y0[6(i-1)+6] = v0[i,3]
end
    
#Solve the system
tf = 50
stepsPerUnitTime = 5000
tspan = linspace(0,tf,tf*stepsPerUnitTime)
t,y = ode23s(F, y0, tspan; points=:specified);

#Extract the data into a useful form
ymat= hcat(y...)
rcoords = sort([[6(i-1)+1 for i = 1:n],[6(i-1)+2 for i = 1:n],[6(i-1)+3 for i = 1:n]])
rcoords = convert(Array{Int64,1}, rcoords)
r = ymat[rcoords,:];

#Find Axis Limits
xmin = minimum(r[1,:])
for i = 2:n
    xmin = min(xmin, minimum(r[3(i-1)+1,:]))
end
xmax = maximum(r[1,:])
for i = 2:n
    xmax = max(xmax, maximum(r[3(i-1)+1,:]))
end
xmin, xmax = xmin - 0.1(xmax-xmin), xmax+ 0.1*(xmax-xmin)

ymin = minimum(r[2,:])
for i = 2:n
    ymin = min(ymin, minimum(r[3(i-1)+2,:]))
end
ymax = maximum(r[2,:])
for i = 2:n
    ymax = max(ymax, maximum(r[3(i-1)+2,:]))
end
ymin, ymax = ymin - 0.1(ymax-ymin), ymax+ 0.1*(ymax-ymin)

zmin = minimum(r[3,:])
for i = 2:n
    zmin = min(zmin, minimum(r[3(i-1)+3,:]))
end
zmax = maximum(r[3,:])
for i = 2:n
    zmax = max(zmax, maximum(r[3(i-1)+3,:]))
end
zmin, zmax = zmin - 0.1(zmax-zmin), zmax+ 0.1*(zmax-zmin)

#Define some random colours
colors = ["r", "g", "b", "k", "y", "c","m","0.75","0.5","0.25","#eeefff","#123456","#a125f6","#ff5868","#ef555e","#98745e","#12ff56","#789463","#8dd56d","#ddee77","#698851","#5544bb"]

#Construct Figure and Plot Data
fig = figure()
ax = plt.axes(xlim = (xmin,xmax),ylim=(ymin,ymax),zlim=(zmin,zmax), projection="3d")

for i = 1:n
    line = ax[:plot]([],[],[],colors[i])[1]
    x = r[3(i-1)+1,:]
    x = reshape(x,length(x))
    y = r[3(i-1)+2,:]
    y =reshape(y,length(y))    
    z = r[3(i-1)+3,:] 
    z =reshape(z,length(z))   
    line[:set_data](x,y)
    line[:set_3d_properties](z)
end
plt.show()

#Use external backend
pygui(true)

#Construct Figure and Plot Data
fig = figure()

#Optional: Redefine plotting range so it is centered at the origin
xmin = -3; xmax = 3; ymin = -3; ymax = 3; zmin = -3; zmax = 3

#Create 3d axes
ax = plt.axes(xlim = (xmin,xmax),ylim=(ymin,ymax),zlim = (zmin,zmax),projection="3d")

#Make cartesian axes invisible
ax[:axis]("off")

#Add lines and planet points. 
global line = Array(Any,n)
global planets = Array(Any,n)
for i = 1:n
    line[i]   = ax[:plot]([],[],[],color = colors[i])[1]
    planets[i]= ax[:plot]([],[],[],"o", color = colors[i])[1]
end

function init()
    global line
    global planets
    for i = 1:n
        line[i][:set_data]([],[])
        line[i][:set_3d_properties]([])
        planets[i][:set_data]([],[])
        planets[i][:set_3d_properties]([])
    end
    result = tuple(tuple([line[i] for i = 1:n]...)...,tuple([planets[i] for i = 1:n]...)...,None)
    return result
end

#To speed up the animation, we will only plot every 'step'th data point
step = 125

function animate(i)
    global line
    global planets
    for k = 1:n
        #Get Planetary Data
        x = r[3(k-1)+1,:]
        x = reshape(x,length(x))
        y = r[3(k-1)+2,:]
        y = reshape(y,length(y)) 
        z = r[3(k-1)+3,:]
        z =reshape(z,length(z))  
        
        #Plot Data in 3D - both the planet position and the past trajectory
        line[k][:set_data](x[max(1,step*(i-200)):(step*i+1)],y[max(1,step*(i-200)):(step*i+1)])
        line[k][:set_3d_properties](z[max(1,step*(i-200)):(step*i+1)])
        planets[k][:set_data](x[step*i+1],y[step*i+1])
        planets[k][:set_3d_properties](z[step*i+1])
        
        #Rotate axes
        ax[:view_init](30, 0.3 * i)
    end
    result = tuple(tuple([line[i] for i = 1:n]...)...,tuple([planets[i] for i = 1:n]...)...,None)
    return result
end

# call the animator.  
myanim = anim.FuncAnimation(fig, animate, init_func=init, frames=ifloor(length(t)/step), interval=20)

#ffmpeg or some equialent will be required to save this animation
#myanim[:save]("4Body3D.mp4", bitrate=-1, extra_args=["-vcodec", "libx264", "-pix_fmt", "yuv420p"])

# Function for creating an embedded video given a filename
function html_video(filename)
	base64_video = base64(open(readbytes, filename))
	"""<video controls src="data:video/x-m4v;base64,$base64_video">"""
end

display("text/html", html_video("4Body3D.mp4"))