#!/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※ = False
fname = 'inclinedthrow'
size = 400, 288
l, w, b, h = 22.5/size※, 1-23/size※, 22.5/size※, 1-23/size※
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 = ※
    def dif(vec, t):
        # time derivative of the whole vector vec
        v = hypot(vec※, vec※)
        vxrel, vyrel = vec※ / v, vec※ / v
        return ※, vec※, -mu * v**pot * vxrel, -g - mu * v**pot * vyrel]
    # solve the differential equation numerically
    vec = odeint(dif, ※, xy0※, vxy0※, vxy0※], tt)
    return vec※, vec※, vec※, vec※  # return x, "y," vx, vy
g = 1.
theta  = radians(70)
v0 = sqrt(g/sin(2*theta))
vinf = 2.1
# 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), ※)※※, 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※) / (w*size※))
    plt.xticks(※), plt.yticks(※)
    plt.xlabel('Distance', size=12)
    plt.ylabel('Height', size=12)
    plt.plot(traj_free※※, traj_free※※,
        '-', lw=lw, color=c1)
    plt.plot(traj_free※※, traj_free※※,
        'ok', color=c1, markersize=ms, markeredgewidth=0)
    plt.plot(traj_stokes※※, traj_stokes※※,
        '-', lw=lw, color=c2)
    plt.plot(traj_stokes※※, traj_stokes※※,
        'ok', color=c2, markersize=ms, markeredgewidth=0)
    plt.plot(traj_newton※※, traj_newton※※,
        '-', lw=lw, color=c3)
    plt.plot(traj_newton※※, traj_newton※※,
        'ok', 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※/72., size※/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)