#!/usr/bin/python3
# -*- coding: utf8 -*-

import os
import inspect
from math import *
import numpy as np
from scipy.integrate import odeint
from scipy.optimize import newton
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib import animation

# settings
mpl.rcParams['path.snap'] = False
fname = 'inclinedthrow2'
size = 400, 288
l, w, b, h = 22.5/size[0], 1-23/size[0], 22.5/size[1], 1-23/size[1]
nframes = 102
delay = 8
lw = 1.
ms = 6
c1, c2, c3 = "#000000", "#0000ff", "#007100"

def projectile_motion(g, mu, pot, xy0, vxy0, tt):
    # use a four-dimensional vector function vec = [x, y, vx, vy]
    def dif(vec, t):
        # time derivative of the whole vector vec
        v = hypot(vec[2], vec[3])
        vxrel, vyrel = vec[2] / v, vec[3] / v
        return [vec[2], vec[3], -mu * v**pot * vxrel, -g - mu * v**pot * vyrel]

    # solve the differential equation numerically
    vec = odeint(dif, [xy0[0], xy0[1], vxy0[0], vxy0[1]], tt)
    return vec[:, 0], vec[:, 1], vec[:, 2], vec[:, 3]  # return x, y, vx, vy

g = 1.
theta  = radians(70)
v0 = sqrt(g/sin(2*theta))
vinf = 0.65
# use identical terminal velocity vinf for both types of friction
mu_stokes = g / vinf**1
mu_newton = g / vinf**2
x0, y0 = 0.0, 0.0
vx0, vy0 = v0 * cos(theta), v0 * sin(theta)

T = newton(lambda t: projectile_motion(g, 0, 0, (x0, y0), (vx0, vy0), [0, t])[1][1], 2*vy0/g)
nsub = 10
tt = np.linspace(0, T * nframes / (nframes - 1), (nframes - 1) * nsub + 1)

traj_free = projectile_motion(g, 0, 0, (x0, y0), (vx0, vy0), tt)
traj_stokes = projectile_motion(g, mu_stokes, 1, (x0, y0), (vx0, vy0), tt)
traj_newton = projectile_motion(g, mu_newton, 2, (x0, y0), (vx0, vy0), tt)

def animate(nframe, saveframes=False):
    print(nframe, '/', nframes)
    t = T * float(nframe) / nframes
    
    plt.clf()
    fig.gca().set_position((l, b, w, h))
    fig.gca().set_aspect("equal")
    plt.xlim(0, 1)
    plt.ylim(0, (h*size[1]) / (w*size[0]))
    plt.xticks([]), plt.yticks([])
    plt.xlabel('Distance', size=12)
    plt.ylabel('Height', size=12)
    
    plt.plot(traj_free[0][:nframe*nsub+1], traj_free[1][:nframe*nsub+1],
        '-', lw=lw, color=c1)
    plt.plot(traj_free[0][nframe*nsub], traj_free[1][nframe*nsub],
        'o', color=c1, markersize=ms, markeredgewidth=0)
    
    plt.plot(traj_stokes[0][:nframe*nsub+1], traj_stokes[1][:nframe*nsub+1],
        '-', lw=lw, color=c2)
    plt.plot(traj_stokes[0][nframe*nsub], traj_stokes[1][nframe*nsub],
        'o', color=c2, markersize=ms, markeredgewidth=0)
    
    plt.plot(traj_newton[0][:nframe*nsub+1], traj_newton[1][:nframe*nsub+1],
        '-', lw=lw, color=c3)
    plt.plot(traj_newton[0][nframe*nsub], traj_newton[1][nframe*nsub],
        'o', color=c3, markersize=ms, markeredgewidth=0)
    
    if saveframes:
        # export frame
        dig = int(ceil(log10(nframes)))
        fsavename = ('frame{:0' + str(dig) + '}.svg').format(nframe)
        fig.savefig(fsavename)
        with open(fsavename) as f: content = f.read()
        content = content.replace('pt"', 'px"').replace('pt"', 'px"')
        with open(fsavename, 'w') as f: f.write(content)

fig = plt.figure(figsize=(size[0]/72., size[1]/72.))

os.chdir(os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))))
for i in range(nframes):
    animate(i, True)
os.system('convert -loop 0 -delay ' + str(delay) + ' frame*.svg +dither ' + fname + '.gif')
# keep last frame for two seconds
os.system('gifsicle -k32 --color-method blend-diversity -b ' + fname + '.gif -d' + str(delay) + ' "#0-' + str(nframes-2) + '" -d200 "#' + str(nframes-1) + '"')
for i in os.listdir('.'):
    if i.startswith('frame') and i.endswith('.svg'):
        os.remove(i)