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

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

code_website = 'https://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_earth-magnetic-field.svg'
fig = mplwp.fig_standard(mpl)

xlim = -90, 90; fig.gca().set_xlim(xlim)
ylim = -65, 65; fig.gca().set_ylim(ylim)
fig.gca().xaxis.set_major_locator(mpl.ticker.MultipleLocator(30))
mplwp.mark_axeszero(fig.gca())

# add degrees to xaxis labels
def flabel(x, i):
    return u'{}\u00B0'.format(int(x)).replace('-', u'\u2212')
fig.gca().xaxis.set_major_formatter(mpl.ticker.FuncFormatter(flabel))

m = 7.746e22 # earth's magnetic moment
R = 6.368e6 # earth's radius
mu0 = 4 * pi * 1e-7

def Br(lat):
    return mu0 / (4*pi) * m / R**3 * 2 * -np.sin(lat)

def Bphi(lat):
    return mu0 / (4*pi) * m / R**3 * np.cos(lat)

def Babs(lat):
    return mu0 / (4*pi) * m / R**3 * np.sqrt(1. + 3. * np.sin(lat)**2)

latitudes = np.linspace(-pi/2, pi/2, 5001)
Br_array = Br(latitudes)
Bphi_array = Bphi(latitudes)
Babs_array = Babs(latitudes)

plt.plot(np.degrees(latitudes), 1e6 * Babs_array, label=r'$\vert B\vert$', zorder=-1)
plt.plot(np.degrees(latitudes), 1e6 * Br_array, label=r'$B_r$', zorder=-2)
plt.plot(np.degrees(latitudes), 1e6 * Bphi_array, label=r'$B_\varphi$', zorder=-3)

mplwp.set_bordersize(fig, 70.5, 18.5, 18.5, 48.5)
plt.xlabel(r'$latitude$')
plt.ylabel(r'$B\ [\mathrm{\mu T}]$')
plt.legend(loc='lower left')
plt.savefig(name)
mplwp.postprocess(name)