#!/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_complete_ellipticPi_nfixed_k.svg'
fig = mplwp.fig_standard(mpl)

xlim = 0, 1; fig.gca().set_xlim(xlim)
ylim = 0, 6; fig.gca().set_ylim(ylim)
mplwp.mark_axeszero(fig.gca())

from mpmath import ellippi
from scipy.optimize import brentq

for i, n in enumerate([0.5, 0, -1, -3, -10]):
    xmax = brentq(lambda xx: ellippi(n, xx**2) - ylim[1], 0, 1)
    x = xlim[0] + (xmax - xlim[0]) * (1-(np.linspace(1, 0, 5001))**2)
    y = [ellippi(n, xx**2) for xx in x]
    if y[-1] < ylim[1] - 1e-6:
        y[-1] = ylim[1]
    plt.plot(x, y, zorder=-i,
        label=r'$\Pi(n={:g},\ k)$'.format(n).replace('-', u'\u2212'))

plt.xlabel('k')
plt.gca().xaxis.set_label_coords(1.05, 0.06)
plt.legend(loc='upper left').get_frame().set_alpha(0.9)
plt.savefig(name)
mplwp.postprocess(name)