"""
Execution of this script inside the IPython interpreter should produce the Figures 1, 2, and 4 of the manuscript.

Notes:
* The script should work both with Python 2.7 and Python 3
* Older versions of scipy do not provide special.roots_hermite, so tests with cpfX cannot be done. 
"""

from scipy.special import wofz
from scipy import pi
recPi = 1/pi

from cpfX import cpf12a, cpf12b, cpf12w, cpfX, zpf16p

from matplotlib.pyplot import *
from matplotlib import ticker   # possibly its already imported automtically
contourLevels = np.logspace(-12,4,17)
set_cmap('RdBu_r')  # if you have already closed the figure this will give a RuntimeError, ignore it


####################################################################################################################################
#####    Contour plot of relative difference of real (or imaginary) complex error functions                                    #####
####################################################################################################################################

def relErrorContour (X, Y, wTest, wRef, iml=False, errorY=False, cBar=True, yLabel=True):
	""" Color contour plot of relative function difference (default real part of complex functions). """
	if iml:
		relError         = abs(wTest.imag-wRef.imag)/wRef.imag
		maxRelErrorOther = np.max(abs(wTest.real-wRef.real)/wRef.real)
		cBarLabel =  r'$\Delta L/L_{\rm ref}$' 
	else:
		relError = abs(wTest.real-wRef.real)/wRef.real
		ijMax = np.argmax(relError);  iMax=ijMax%relError.shape[1];  jMax=int(ijMax/relError.shape[1])
		print ('%s %g @ %i %i %i     x = %.3f   y = %.1g' % \
		       ('max rel error:', np.max(relError), ijMax, iMax, jMax, X[0,iMax], Y[jMax,0]))
		cBarLabel = r'$\Delta K/K_{\rm ref}$' 
		# max rel. error of imaginary part
		relErrorIm       = np.where(abs(wRef.imag)>0, abs(wTest.imag-wRef.imag)/wRef.imag, 0.0)
		maxRelErrorOther = np.max(relErrorIm)

	if errorY:
		axes([0.05,0.1,0.99,0.85])  # axes([0.05,0.1,0.75,0.85])
		cPad=0.25
	else:
		cPad=0.05

	contourf(X,np.log10(Y), relError, contourLevels, locator=ticker.LogLocator())
	xlabel('$x$')
	if yLabel:  ylabel(r'$\log_{10}(y)$')
	if cBar:    colorbar(label=cBarLabel, pad=cPad, extend='both')  # pad=0.01
	title( 'max error  %.3e  (%.2e)' % (np.max(relError), maxRelErrorOther))

	if errorY:
		maxRelErrorY = np.max(relError,1)
		logY = np.log10(Y[:,0])
		axes([0.67,0.1,0.19,0.85])  # axes([0.78,0.1,0.2,0.85])
		semilogx (maxRelErrorY,logY);  ylim(logY.min(),logY.max());  yticks(visible=0);  xlabel(cBarLabel)
		return maxRelErrorY


####################################################################################################################################
####################################################################################################################################

x=np.linspace(0.,25.,251)
y=np.logspace(-8,2,101)
# kind of 'matrices' required for the contour plots (note Python is case sensitive!)
X, Y = np.meshgrid(x, y)
Z = X +1j*Y

# the reference code
wRef  = wofz(Z)

# Humlicek (1979) 12 term rational approximations
w12w = cpf12w(Z)
w12a = cpf12a(Z)

figure('fig1_cpf12', figsize=(12,5))
axes([0.05,0.09,0.4,0.85]); relErrorContour(X,Y, w12a,wRef, cBar=0)
axes([0.49,0.09,0.5,0.85]); relErrorContour(X,Y, w12w,wRef, yLabel=0)


# Humlicek (1979) 16 term rational approximations
w16a = cpfX(Z, 16, 1.31183)
w16b = cpfX(Z, 16, 1.35)

figure('fig2_cpf16', figsize=(12,5))
axes([0.05,0.09,0.4,0.85]); relErrorContour(X,Y, w16a,wRef, cBar=0)
axes([0.49,0.09,0.5,0.85]); relErrorContour(X,Y, w16b,wRef, yLabel=0)


# Humlicek (1979) 18 and 20 term rational approximations
w18 = cpfX(Z, 18, 1.45)
w20 = cpfX(Z, 20, 1.55)

figure('fig4_cpf18_20', figsize=(12,5))
axes([0.05,0.09,0.4,0.85]); relErrorContour(X,Y, w18,wRef, cBar=0)
axes([0.49,0.09,0.5,0.85]); relErrorContour(X,Y, w20,wRef, yLabel=0)