"""
Humlicek's rational approximations for the complex error/probability function:
w4    -- Optimized computation of the Voigt and complex probability function (JQSRT 1982).
cpf12 -- An efficient method for evaluation of the complex probability function: the Voigt function and its derivatives (JQSRT 1979)
"""

from math import pi, sqrt
recSqrtPi = 1./sqrt(pi)

try:                      import numpy as np
except ImportError, msg:  raise SystemExit (str(msg) + '\nhumlicek.py:  import numpy failed!')

# select the functions to be imported  (thus 'from humlicek import *' will not import the constants!)
__all__ = 'humlicek1 humlicek2 humlicek3 humlicek4 humlicek  cpf12a cpf12b cpf12'.split()

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

def humlicek1 (z):
	""" Humlicek (1982) rational approximation:  region I. """
	#t = z.imag-1j*z.real;   w = t * recSqrtPi / (0.5 + t*t)
	w = 1j*recSqrtPi *z / (z*z-0.5)
	return w

def humlicek2 (z):
	""" Humlicek (1982) rational approximation:  region II. """
	# this implementation corresponds to the Fortran code
	x, y = z.real, z.imag
	t = y-1j*x
	u = t*t
	w = (t * (1.4104739589 + u*recSqrtPi))/ (.75 + (u *(3.+u)))
	return w

def Humlicek2 (z):
	""" Humlicek (1982) rational approximation:  region II. """
	# this implementation corresponds to Eq. (12) of the manuscript
	zz = z*z
	w = 1j* (z * (zz*recSqrtPi-1.4104739589))/ (0.75 + zz*(zz-3.0))
	return w

def humlicek3 (z):
	""" Humlicek (1982) rational approximation:  region III. """
	x, y = z.real, z.imag
	t = y-1j*x
	w =  (16.4955 + t * (20.20933 + t * (11.96482 + t * (3.778987 + 0.5642236*t)))) \
	     / (16.4955 + t * (38.82363 + t * (39.27121 + t * (21.69274 + t * (6.699398 + t)))))
	return w

def humlicek4 (z):
	""" Humlicek (1982) rational approximation:  region IV. """
	x, y = z.real, z.imag
	t = y-1j*x
	u = t*t
	nom = t*(36183.31-u*(3321.99-u*(1540.787-u*(219.031-u*(35.7668-u*(1.320522-u*.56419))))))
	den = 32066.6-u*(24322.8-u*(9022.23-u*(2186.18-u*(364.219-u*(61.5704-u*(1.84144-u))))))
	w  = np.exp(u) - nom/den
	return w


def _humlicek_(z):
	""" Humlicek (1982) complex probability function:  w4 rational approximations.
	    Scalar version;  use humlicek for arrays. """
	s = abs(z.real) + z.imag
	if s>15.0:
		w = humlicek1(z)
	elif s>5.5:
		w = humlicek2(z)
	else:
		if z.imag>=0.195*abs(z.real)-0.176:  w = humlicek3(z)
		else:                                w = humlicek4(z)
	return w


# define a vectorized function
humlicek = np.vectorize(_humlicek_)


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

cr=1.5; crr=2.25  #  y0 and y0**2 of the original
ct = np.array([0., .3142403762544,  .9477883912402, 1.5976826351526, 2.2795070805011, 3.0206370251209, 3.88972489786978])
ca = np.array([0., -1.393236997981977,    -0.2311524061886763,   +0.1553514656420944,
                   -0.006218366236965554, -9.190829861057117e-5, +6.275259577e-7])
cb = np.array([0.,  1.011728045548831,    -0.7519714696746353, 0.01255772699323164,
                    0.01002200814515897, -2.420681348155727e-4,  5.008480613664576e-7])

def cpf12a (z):
	""" Humlicek (1979) complex probability function:  rational approximation for y>0.85 OR |x|<18.1*y+1.65. """
	x, y = z.real, z.imag
	ry = cr+y
	ryry = ry**2
	wk =  (ca[1]*(x-ct[1]) + cb[1]*ry) / ((x-ct[1])**2 + ryry) - (ca[1]*(x+ct[1]) - cb[1]*ry) / ((x+ct[1])**2 + ryry) \
	    + (ca[2]*(x-ct[2]) + cb[2]*ry) / ((x-ct[2])**2 + ryry) - (ca[2]*(x+ct[2]) - cb[2]*ry) / ((x+ct[2])**2 + ryry) \
	    + (ca[3]*(x-ct[3]) + cb[3]*ry) / ((x-ct[3])**2 + ryry) - (ca[3]*(x+ct[3]) - cb[3]*ry) / ((x+ct[3])**2 + ryry) \
	    + (ca[4]*(x-ct[4]) + cb[4]*ry) / ((x-ct[4])**2 + ryry) - (ca[4]*(x+ct[4]) - cb[4]*ry) / ((x+ct[4])**2 + ryry) \
	    + (ca[5]*(x-ct[5]) + cb[5]*ry) / ((x-ct[5])**2 + ryry) - (ca[5]*(x+ct[5]) - cb[5]*ry) / ((x+ct[5])**2 + ryry) \
	    + (ca[6]*(x-ct[6]) + cb[6]*ry) / ((x-ct[6])**2 + ryry) - (ca[6]*(x+ct[6]) - cb[6]*ry) / ((x+ct[6])**2 + ryry)
	wl =  (cb[1]*(x-ct[1]) - ca[1]*ry) / ((x-ct[1])**2 + ryry) + (cb[1]*(x+ct[1]) + ca[1]*ry) / ((x+ct[1])**2 + ryry) \
	    + (cb[2]*(x-ct[2]) - ca[2]*ry) / ((x-ct[2])**2 + ryry) + (cb[2]*(x+ct[2]) + ca[2]*ry) / ((x+ct[2])**2 + ryry) \
	    + (cb[3]*(x-ct[3]) - ca[3]*ry) / ((x-ct[3])**2 + ryry) + (cb[3]*(x+ct[3]) + ca[3]*ry) / ((x+ct[3])**2 + ryry) \
	    + (cb[4]*(x-ct[4]) - ca[4]*ry) / ((x-ct[4])**2 + ryry) + (cb[4]*(x+ct[4]) + ca[4]*ry) / ((x+ct[4])**2 + ryry) \
	    + (cb[5]*(x-ct[5]) - ca[5]*ry) / ((x-ct[5])**2 + ryry) + (cb[5]*(x+ct[5]) + ca[5]*ry) / ((x+ct[5])**2 + ryry) \
	    + (cb[6]*(x-ct[6]) - ca[6]*ry) / ((x-ct[6])**2 + ryry) + (cb[6]*(x+ct[6]) + ca[6]*ry) / ((x+ct[6])**2 + ryry)
	return wk+1j*wl   # wk, wl

def cpf12b (z):
	""" Humlicek (1979) complex probability function:  rational approximation for y<0.85 AND |x|>18.1*y+1.65. """
	x, y = z.real, z.imag
	ry   = cr+y      # yy0   = y+1.5
	y2r  = y +2.*cr  # y2y0  = y+3.0
	rry  = cr*ry     # y0yy0 = 1.5*(y+1.5)
	ryry = ry**2     # yy0sq = (y+1.5)**2
	wk =  ( cb[1]*((x-ct[1])**2-rry) - ca[1]*(x-ct[1])*y2r ) / (((x-ct[1])**2+crr)*((x-ct[1])**2+ryry)) \
            + ( cb[1]*((x+ct[1])**2-rry) + ca[1]*(x+ct[1])*y2r ) / (((x+ct[1])**2+crr)*((x+ct[1])**2+ryry)) \
	    + ( cb[2]*((x-ct[2])**2-rry) - ca[2]*(x-ct[2])*y2r ) / (((x-ct[2])**2+crr)*((x-ct[2])**2+ryry)) \
            + ( cb[2]*((x+ct[2])**2-rry) + ca[2]*(x+ct[2])*y2r ) / (((x+ct[2])**2+crr)*((x+ct[2])**2+ryry)) \
	    + ( cb[3]*((x-ct[3])**2-rry) - ca[3]*(x-ct[3])*y2r ) / (((x-ct[3])**2+crr)*((x-ct[3])**2+ryry)) \
            + ( cb[3]*((x+ct[3])**2-rry) + ca[3]*(x+ct[3])*y2r ) / (((x+ct[3])**2+crr)*((x+ct[3])**2+ryry)) \
	    + ( cb[4]*((x-ct[4])**2-rry) - ca[4]*(x-ct[4])*y2r ) / (((x-ct[4])**2+crr)*((x-ct[4])**2+ryry)) \
            + ( cb[4]*((x+ct[4])**2-rry) + ca[4]*(x+ct[4])*y2r ) / (((x+ct[4])**2+crr)*((x+ct[4])**2+ryry)) \
	    + ( cb[5]*((x-ct[5])**2-rry) - ca[5]*(x-ct[5])*y2r ) / (((x-ct[5])**2+crr)*((x-ct[5])**2+ryry)) \
            + ( cb[5]*((x+ct[5])**2-rry) + ca[5]*(x+ct[5])*y2r ) / (((x+ct[5])**2+crr)*((x+ct[5])**2+ryry)) \
	    + ( cb[6]*((x-ct[6])**2-rry) - ca[6]*(x-ct[6])*y2r ) / (((x-ct[6])**2+crr)*((x-ct[6])**2+ryry)) \
            + ( cb[6]*((x+ct[6])**2-rry) + ca[6]*(x+ct[6])*y2r ) / (((x+ct[6])**2+crr)*((x+ct[6])**2+ryry))
	wl =  (cb[1]*(x-ct[1]) - ca[1]*ry) / ((x-ct[1])**2 + ryry) + (cb[1]*(x+ct[1]) + ca[1]*ry) / ((x+ct[1])**2 + ryry) \
	    + (cb[2]*(x-ct[2]) - ca[2]*ry) / ((x-ct[2])**2 + ryry) + (cb[2]*(x+ct[2]) + ca[2]*ry) / ((x+ct[2])**2 + ryry) \
	    + (cb[3]*(x-ct[3]) - ca[3]*ry) / ((x-ct[3])**2 + ryry) + (cb[3]*(x+ct[3]) + ca[3]*ry) / ((x+ct[3])**2 + ryry) \
	    + (cb[4]*(x-ct[4]) - ca[4]*ry) / ((x-ct[4])**2 + ryry) + (cb[4]*(x+ct[4]) + ca[4]*ry) / ((x+ct[4])**2 + ryry) \
	    + (cb[5]*(x-ct[5]) - ca[5]*ry) / ((x-ct[5])**2 + ryry) + (cb[5]*(x+ct[5]) + ca[5]*ry) / ((x+ct[5])**2 + ryry) \
	    + (cb[6]*(x-ct[6]) - ca[6]*ry) / ((x-ct[6])**2 + ryry) + (cb[6]*(x+ct[6]) + ca[6]*ry) / ((x+ct[6])**2 + ryry)
	return np.exp(-x*x)+y*wk+1j*wl   # wk, wl

def cpf12 (z):
	""" Humlicek (1979) complex probability function: a single rational approximation. """
	xMask = abs(z.real)<18.1*z.imag+1.65
	yMask = z.imag>0.85
	mask  = np.logical_or(xMask,yMask)
	return np.where (mask, cpf12a(z), cpf12b(z))