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

xlim = -0.05, 4.0; fig.gca().set_xlim(xlim)
ylim = 0, 2.5; fig.gca().set_ylim(ylim)
mplwp.mark_axeszero(fig.gca())

def agm(x, y):
    a, g = x, y
    for i in range(13):
        a, g = 0.5 * (a + g), sqrt(a * g)
        if fabs(a - g) / (a + g) < 1e-16:
            break
    return a

x = np.linspace(0, 4, 4001)

y0 = [xx for xx in x]
plt.plot(x, y0, ':', label='x', color='#00bb00', dashes=[3,3], zorder=1)

y1 = [0.5 * (1.0 + xx) for xx in x]
plt.plot(x, y1, '--', color='#00bbcc', dashes=[12,3], zorder=2,
    label=r'$\frac{1}{2}(1+x)$')

y2 = [agm(1.0, xx) for xx in x]
plt.plot(x, y2, '-', label=r'$agm(1, x)$', color='#0000cc', zorder=4)

y3 = [sqrt(1.0 * xx) for xx in x]
plt.plot(x, y3, '--', color='#dd00aa', dashes=[12,3], zorder=3,
    label=r'$\sqrt{1\cdot x}$')

y4 = [1.0 for xx in x]
plt.plot(x, y4, ':', label='1', color='#999900', dashes=[3,3], zorder=0)

plt.legend(loc='lower right')
plt.savefig(name)
mplwp.postprocess(name)