#!/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_dispersion curves.svg'
fig = mplwp.fig_standard(mpl)
xlim = 0.3, 1.6; fig.gca().set_xlim(xlim)
ylim = 1.4, 1.9; fig.gca().set_ylim(ylim)
mplwp.set_bordersize(fig, 60.5, 26.5, 24.5, 48.5)
plt.gca().axvspan(0.38, 0.78, alpha=0.5, color='#aaaaaa')

# functions: https://refractiveindex.info
def n(l, coeff):
    n21 = 0.
    for a,b in coeff:
        n21 += a * l**2 / (l**2 - b)
    return sqrt(1. + max(0, n21))

# coefficients for the https://en.wikipedia.org/wiki/Sellmeier_equation
LASF9 = ((2.00030, 0.01214), (0.29893, 0.05387), (1.80692, 156.53083))
SF10 = ((1.62154, 0.01222), (0.25629, 0.05957), (1.64448, 147.46879))
F2 = ((1.34533, 0.00998), (0.20907, 0.04705), (0.93736, 111.88676))
BaK4 = ((1.28835, 0.00780), (0.13282, 0.03156), (0.94540, 105.96588))
BK7 = ((1.18319, 0.00722), (0.08718, 0.02682), (1.03134, 101.70236))
FK51A = ((0.97125, 0.00472), (0.21690, 0.01536), (0.90465, 168.68133))

x = np.linspace(xlim[0], xlim[1], 5001)
y = [n(xx, LASF9) for xx in x]
plt.plot(x, y, color='#0000cc')
plt.text(1.5, 1.84, 'Lanthanum dense flint LaSF9', ha='right')

y = [n(xx, SF10) for xx in x]
plt.plot(x, y, color='#0000cc')
plt.text(1.5, 1.71, 'Dense flint SF10', ha='right')

y = [n(xx, F2) for xx in x]
plt.plot(x, y, color='#0000cc')
plt.text(1.5, 1.608, 'Flint F2', ha='right')

y = [n(xx, BaK4) for xx in x]
plt.plot(x, y, color='#0000cc')
plt.text(1.5, 1.565, 'Barium crown BaK4', ha='right')

y = [n(xx, BK7) for xx in x]
plt.plot(x, y, color='#0000cc')
plt.text(1.5, 1.515, 'Borosilicate crown BK7', ha='right')

y = [n(xx, FK51A) for xx in x]
plt.plot(x, y, color='#0000cc')
plt.text(1.5, 1.45, 'Fluorite crown FK51A', ha='right')

plt.xlabel(r'wavelength $\lambda$ [$\mu m$]')
plt.ylabel(r'refractive index n')

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