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

xlim = -8,5; fig.gca().set_xlim(xlim)
ylim = -4,4; fig.gca().set_ylim(ylim)
mplwp.mark_axeszero(fig.gca())

import mpmath

for s in range(-2, 3+1):
    def f(s, x, part='real'):
        if (part == 'real' or s <= 1) and x == 1.0:
            return float('nan')
        if part == 'imag':
            return float(mpmath.polylog(s, x).imag)
        else:
            return float(mpmath.polylog(s, x).real)
    
    x = np.linspace(xlim[0], xlim[1], 6501)
    if s <= 0 and int(s) == s:
        x = np.linspace(xlim[0], xlim[1], 6501)
        y = [f(s, xx, 'real') for xx in x]
        plt.plot(x, y, label='Li' + mplwp.unicode_sub(s) + '(x)')
    else:
        yre = [f(s, xx, 'real') for xx in x]
        plt.plot(x, yre, label='Re(Li' + mplwp.unicode_sub(s) + '(x))')
        
        yim = [f(s, xx, 'imag') for xx in x]
        plt.plot(x, yim, label='Im(Li' + mplwp.unicode_sub(s) + '(x))',
            color=fig.gca().lines[-1].get_color(), ls='--')

mpl.rc('legend', borderaxespad=0.5)
plt.legend(loc='upper left', ncol=2, columnspacing=1.1, handletextpad=0.4
    ).get_frame().set_alpha(0.9)

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