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

import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
from math import *

code_website = 'http://commons.wikimedia.org/wiki/User:Geek3/mplwp'
try:
    import mplwp
except ImportError, er:
    print 'ImportError:', er
    print 'You need to download mplwp.py from', code_website
    exit(1)

name = 'mplwp_skydive_trajectory.svg'
fig = mplwp.fig_standard(mpl)
fig.set_size_inches(480. / fig.dpi, 500. / fig.dpi)
mplwp.move_axes(fig, 24, 6)

xlim = -10,320; fig.gca().set_xlim(xlim)
ylim = -340,20; fig.gca().set_ylim(ylim)
fig.gca().set_aspect('equal')
fig.gca().xaxis.set_major_locator(mpl.ticker.MultipleLocator(100))
fig.gca().yaxis.set_major_locator(mpl.ticker.MultipleLocator(100))

from scipy.integrate import odeint
from scipy.optimize import brentq
def ballistic(g, mu, xy0, v0, alpha0, tt):
    # use a four-dimensional vector function vec = [x, y, vx, vy]
    def dif(vec, t):
        v = sqrt(vec[2]**2 + vec[3]**2)
        return [vec[2], vec[3], -mu*v*vec[2], -g -mu*v*vec[3]]
    
    # solve the differenctial equation numerically
    vec = odeint(dif, [xy0[0], xy0[1], v0*cos(alpha0), v0*sin(alpha0)], tt)
    return vec[:,0], vec[:,1] # return x(tt) and y(tt)

g = 9.81
vmax = 50.
mu = g / vmax**2
alpha0 = 0.

for iv, v0 in enumerate((0, 17, 33, 50)):
    t1 = brentq(lambda t: ballistic(g,mu,[0,0],v0,alpha0,[0,t])[1][1]-ylim[0],0,20)
    t = np.linspace(0, t1, 5001)
    x, y = ballistic(g, mu, [0, 0], v0, alpha0, t)
    plt.plot(x, y, label='$v_0 = {:.0f}\,m/s$'.format(v0), zorder=-iv)
    col = fig.gca().lines[-1].get_color()
    xmax = ballistic(g, mu, [0, 0], v0, alpha0, [0, 100*vmax/g])[0][-1]
    plt.axvline(xmax, 0, 0.33, linestyle='--', dashes=[10,10],
        color=col, zorder=-5, lw=1.5)
    times = np.arange(t1)
    xt, yt = ballistic(g, mu, [0, 0], v0, alpha0, times)
    plt.plot(xt, yt, '.', color=col, label=None, zorder=-iv-0.5)
    
    if v0 == 50.:
        for ti, xi, yi in zip(times, xt, yt):
            plt.text(xi+8*sin(pi/2*ti/10), yi+8*cos(pi/2*ti/10),
                '{:.0f}s'.format(ti), ha='left', va='center', fontsize=10)

plt.xlabel('x [m]')
plt.ylabel('y [m]')
fig.gca().xaxis.labelpad = 0
fig.gca().yaxis.set_label_coords(-0.08, 0.95)
mpl.rc('legend', borderaxespad=0.4)
plt.legend(loc='upper right')

plt.savefig(name)
mplwp.postprocess(name)