"""
Execution of this script inside the IPython interpreter should produce contour and xy plots of relative errors

Notes:
* The script has bee tested with Python 3 (and probably works with Python 2.7, too?)
"""

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

from cpfX import cpf12a, cpf12b, cpf12w, cpfX, zpf16p, hum1zpf16m, hum2zpf16m, hum1zpf16w, hum2zpf16w
from racef import weideman24a, hum1wei24, hum2wei32, hum1wei24w, hum2wei32w


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
rc('figure.subplot', top=0.95, bottom=0.07, left=0.08, right=0.975, wspace=0.15)


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:  # plot error of imaginary part
 		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.02

 	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


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

xx=np.linspace(0.,25.,251)
yy=np.logspace(-8,2,101)

# kind of 'matrices' required for the contour plots (note Python is case sensitive!)
XX, YY = np.meshgrid(xx, yy)
ZZ = XX +1j*YY

# the reference code
wRef  = wofz(ZZ)
kRef  = wRef.real

w_zpf16 = zpf16p(ZZ)
w_hum1zpf16 = hum1zpf16w(ZZ)
w_hum2zpf16 = hum2zpf16w(ZZ)
w_hum1wei24 = hum1wei24w(XX,YY)  # note the two arguments!
w_hum2wei32 = hum2wei32w(XX,YY)


# plot the real part of the complex error function for various y
for j,y in enumerate(yy[::10]):  semilogy (xx, w_zpf16[10*j,:].real, label='%i %g' % (j, y))
xlabel('$x$');  ylabel(r'$K(x,y)$');  legend()

# plot relative errors as contour plot and for selected y
wTest = w_hum1wei24
kTest = wTest.real
subplot(121)
for j,y in enumerate(yy[::10]):
    kDelta = abs(kTest[10*j,:]-kRef[10*j,:])/kRef[10*j,:]
    semilogy (xx, kDelta, label = '%.1e %.3g' % (y, max(kDelta)))
xlabel('$x$');  ylabel(r'$\Delta K/K_{\rm ref}$')
legend()
subplot(122)
relErrorContour(XX,YY, wTest, wRef)
suptitle('humlicek1 weideman24', fontsize='x-large')

# compare relative errors for several approximations
for j,y in enumerate(yy[::-20]):
	subplot(6,1,j+1)
	kRef = wofz(xx+1j*y).real
	semilogy (xx, abs(zpf16p(xx+1j*y).real-kRef)/kRef, label='zpf16')
	# milogy (xx, abs(hum1zpf16w(xx+1j*y).real-kRef)/kRef, label='hum1zpf16w')
	semilogy (xx, abs(hum1zpf16m(xx,   y).real-kRef)/kRef, label='hum1zpf16m') 
	# milogy (xx, abs(hum1wei24w(xx,   y).real-kRef)/kRef, label='hum1wei24w') 
	semilogy (xx, abs(hum1wei24 (xx,   y).real-kRef)/kRef, label='hum1wei24m')  # note the missing "m" of the function name
	semilogy (xx, abs(hum2wei32w(xx,   y).real-kRef)/kRef, label='hum2wei32') 
	title('$y$ = %.1e' % y)
xlabel('$x$');  ylabel(r'$\Delta K/K_{\rm ref}$');  legend()