import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import matplotlib.patches as mpatches

def add_arrow_to_line2D(
    axes, line, arrow_locs=[0.2, 0.4, 0.6, 0.8],
    arrowstyle='-|>', arrowsize=1, transform=None):
    """
    Add arrows to a matplotlib.lines.Line2D at selected locations.

    Parameters:
    -----------
    axes: 
    line: list of 1 Line2D obbject as returned by plot command
    arrow_locs: list of locations where to insert arrows, % of total length
    arrowstyle: style of the arrow:
    -         None
    ->         head_length=0.4,head_width=0.2
    -[         widthB=1.0,lengthB=0.2,angleB=None
    |-| widthA=1.0,widthB=1.0
    -|> head_length=0.4,head_width=0.2
    <-         head_length=0.4,head_width=0.2
    <-> head_length=0.4,head_width=0.2
    <|- head_length=0.4,head_width=0.2
    <|-|> head_length=0.4,head_width=0.2
    fancy head_length=0.4,head_width=0.4,tail_width=0.4
    simple head_length=0.5,head_width=0.5,tail_width=0.2
    wedge tail_width=0.3,shrink_factor=0.5

    arrowsize: size of the arrow
    transform: a matplotlib transform instance, default to data coordinates

    Returns:
    --------
    arrows: list of arrows
    """
    if (not(isinstance(line, list)) or not(isinstance(line[0], 
                                           mlines.Line2D))):
        raise ValueError("expected a matplotlib.lines.Line2D object")
    x, y = line[0].get_xdata(), line[0].get_ydata()

    arrow_kw = dict(arrowstyle=arrowstyle, mutation_scale=10 * arrowsize)

    color = line[0].get_color()
    use_multicolor_lines = isinstance(color, np.ndarray)
    if use_multicolor_lines:
        raise NotImplementedError("multicolor lines not supported")
    else:
        arrow_kw['color'] = color

    linewidth = line[0].get_linewidth()
    if isinstance(linewidth, np.ndarray):
        raise NotImplementedError("multiwidth lines not supported")
    else:
        arrow_kw['linewidth'] = linewidth

    if transform is None:
        transform = axes.transData

    arrows = []
    for loc in arrow_locs:
        s = np.cumsum(np.sqrt(np.diff(x) ** 2 + np.diff(y) ** 2))
        n = np.searchsorted(s, s[-1] * loc)
        arrow_tail = (x[n], y[n])
        arrow_head = (np.mean(x[n:n + 2]), np.mean(y[n:n + 2]))
        p = mpatches.FancyArrowPatch(
            arrow_tail, arrow_head, transform=transform,
            **arrow_kw)
        axes.add_patch(p)
        arrows.append(p)
    return arrows

A=1.
B=2.
C=3.
D=4.
E=0.
def fun(u, t):
    x, y = u
    dudt = [x*(A-B*y-E*x),y*(C*x-D)]
    return dudt

t = np.linspace(0, 25, 1001)

fig, (ax,ax2) = plt.subplots(2,1)

for y0 in [[D/C,.1],[D/C,.5],[D/C,1], [D/C,2],[D/C,3]]:
    sol = odeint(fun, y0, t)
    x = sol[:,0]
    y = sol[:,1]
    line = ax.plot(x, y, 'k-')
    add_arrow_to_line2D(ax, line, arrow_locs=[.2,.45],
                        arrowstyle='->', arrowsize=2)

ax2.plot(t,x,'k')
ax2.plot(t,y,'r')

plt.show()