""" Speed dependent Voigt (SDV) function K(x,y,q)

    The generalization of the Voigt function K(x,y) with one additional arguments q.

    ARGUMENTS:
    x = sqrtLn2*(vGrid-vLine)/gammaG         (vGrid = wavenumber (grid);   vLine = actual line position (incl. shift))
    y = sqrtLn2*gammaL/gammaG                (gammaL = gamma0 speed averaged broadening at actual p,T)
    q = sqrtLn2*gamma2/gammaG                (gamma2 = quadratic dependence of broadening parameter at p,T)

    RETURNS:
    sdv = the complex "generalization" of the complex error function w

    REFERENCES:
    BWB = Boone, Walker, Bernath:
    An efficient analytical approach for calculating line mixing in atmospheric remote sensing applications.
    JQSRT 112(6), pp. 8=980-089, 2011

    TNH = Tran, Ngo, Hartmann: Efficient computation of some speed-dependent isolated line profiles.
    JQSRT 129, pp. 199–203, 2013. (Erratum: JQSRT 134, 104 (2014))

    FS: Computational aspects of speed-dependent Voigt profiles. JQSRT 187, pp. 44-53, 2017
"""

import numpy as np
from scipy.special import wofz
recSqrt2 = 1.0/np.sqrt(2.)

import mpmath as mp

# Py4CAtS functions
from aeiou import awrite

# public functions (i.e. do not import the constants):
__all__ = 'sdv_real_bwb sdv_real_fgs sdv_naive sdv_safe sdv_tnh sdv_mpm'.split()


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

def sdv_real_bwb (xx, y, q, cerf=wofz, verbose=False):
	""" Speed dependent Voigt: naiv computation of complex square root difference.
	    (The equivalent of the Voigt function K(x,y), i.e. given q and the Voigt variables x,y).

	    This is the 'naive' version without check if the two arguments zPlus, zMinus are in the same region.
	    Real arithmetic with three square roots and the sign(beta) as in the BWB paper. """

	# sdv effective parameters --- bwb notation
	alfa  = y/q - 1.5
	sqrtDelta = 0.5*(alfa+1.5)/y
	delta     = sqrtDelta**2

	if verbose:
		print ('\nsdv    y', y, '   alfa', alfa, '   sqrtDelta', sqrtDelta, '   delta', delta,
		       '   a/2*sd',  0.5*alfa/sqrtDelta, 'a/2sd + 2sd', 2.0*sqrtDelta+0.5*alfa/sqrtDelta, '\n')

	if isinstance(xx,(int,float)):  xx = np.array([xx])

	kSDV = np.empty_like(xx)

	beta    = 2.0*sqrtDelta*xx
	r4      = np.sqrt((delta+alfa)**2+beta**2)
	sqrtXY  = recSqrt2 * (-np.sign(beta)*np.sqrt(np.maximum(r4 - delta - alfa,0.0)) +1j*np.sqrt(r4 + delta + alfa))
	izMinus = sqrtXY - 1j*sqrtDelta
	izPlus  = sqrtXY + 1j*sqrtDelta
	wMinus  = cerf(izMinus).real
	wPlus   = cerf(izPlus).real
	kSDV    = wMinus-wPlus

	if verbose:  awrite ((xx, alfa+1j*beta, sqrtXY, -1j*izMinus, -1j*izPlus, wMinus, wPlus, kSDV), format='%s')

	if verbose:   return kSDV, izMinus, wMinus, izPlus, wPlus
	else:         return kSDV


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

def sdv_real_fgs (xx, y, q, cerf=wofz, verbose=False):
	""" Speed dependent Voigt: naiv computation of complex square root difference.
	    (The equivalent of the Voigt function K(x,y), i.e. given q and the Voigt variables x,y).

	    This is the 'naive' version without check if the two arguments zPlus, zMinus are in the same region.
	    Real arithmetic with two square roots, but one more division. """

	# sdv effective parameters --- bwb notation
	alfa  = y/q - 1.5
	sqrtDelta = 0.5*(alfa+1.5)/y
	delta     = sqrtDelta**2

	if verbose:
		print ('\nsdv    y', y, '   alfa', alfa, '   sqrtDelta', sqrtDelta, '   delta', delta,
		       '   a/2*sd',  0.5*alfa/sqrtDelta, 'a/2sd + 2sd', 2.0*sqrtDelta+0.5*alfa/sqrtDelta, '\n')

	if isinstance(xx,(int,float)):  xx = np.array([xx])

	kSDV = np.empty_like(xx)

	beta    = 2.0*sqrtDelta*xx
	r4      = np.sqrt((delta+alfa)**2+beta**2)
	r2      = np.sqrt(r4 + delta + alfa)
	sqrtXY  = recSqrt2 * (1j*r2 - beta/r2)
	izMinus = sqrtXY - 1j*sqrtDelta
	izPlus  = sqrtXY + 1j*sqrtDelta
	wMinus  = cerf(izMinus).real
	wPlus   = cerf(izPlus).real
	kSDV    = wMinus-wPlus

	if verbose:  awrite ((xx, alfa+1j*beta, sqrtXY, -1j*izMinus, -1j*izPlus, wMinus, wPlus, kSDV), format='%s')

	if verbose:   return kSDV, izMinus, wMinus, izPlus, wPlus
	else:         return kSDV


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

def sdv_naive (xx, y, q, cerf=wofz, verbose=False):
	""" Speed dependent Voigt: naiv computation of complex square root difference.
	    (The equivalent of the Voigt function K(x,y), i.e. given q and the Voigt variables x,y).

	    This is the 'naive' version without check if the two arguments zPlus, zMinus are in the same region.
	    Complex square root computation similar to TNH. """

	# sdv effective parameters --- bwb notation
	alfa  = y/q - 1.5
	sqrtDelta = 0.5*(alfa+1.5)/y
	delta     = sqrtDelta**2

	if verbose:
		print ('\nsdv    y', y, '   alfa', alfa, '   sqrtDelta', sqrtDelta, '   delta', delta,
		       '   a/2*sd',  0.5*alfa/sqrtDelta, 'a/2sd + 2sd', 2.0*sqrtDelta+0.5*alfa/sqrtDelta, '\n')

	if isinstance(xx,(int,float)):  xx = np.array([xx])

	kSDV = np.empty_like(xx);   zMinus = np.empty_like(xx, dtype=complex)

	X      = alfa + 2j*sqrtDelta*xx  # alfa+1j*beta
	sqrtXY = np.sqrt(X+delta)
	zMinus = sqrtXY - sqrtDelta
	zPlus  = sqrtXY + sqrtDelta
	wMinus = cerf(1j*zMinus).real
	wPlus  = cerf(1j*zPlus).real
	kSDV   = wMinus-wPlus

	if verbose:  awrite ((xx, X, sqrtXY, zMinus, zPlus, wMinus, wPlus, kSDV))

	if verbose:   return kSDV, 1j*zMinus, wMinus, 1j*zPlus, wPlus
	else:         return kSDV


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

def sdv_safe (xx, y, q, cerf=wofz, verbose=False):
	""" Speed dependent Voigt: naiv computation of complex square root difference.
	    (The equivalent of the Voigt function K(x,y), i.e. given q and the Voigt variables x,y).

	    This is the 'naive' version without check if the two arguments zPlus, zMinus are in the same region.
	    Computation of complex square root,
	    cancellation safe computation of complex square root difference. """

	if isinstance(q,(int,float)) and q<0:  q=y*abs(q)

	# sdv effective parameters --- bwb notation
	alfa  = y/q - 1.5
	sqrtDelta = 0.5*(alfa+1.5)/y
	delta     = sqrtDelta**2

	if verbose:
		print ('\nsdv    y', y, '   alfa', alfa, '   sqrtDelta', sqrtDelta, '   delta', delta,
		       '   a/2*sd',  0.5*alfa/sqrtDelta, 'a/2sd + 2sd', 2.0*sqrtDelta+0.5*alfa/sqrtDelta, '\n')

	if isinstance(xx,(int,float)):  xx = np.array([xx])

	kSDV = np.empty_like(xx);   zMinus = np.empty_like(xx, dtype=complex)

	X      = alfa + 2j*sqrtDelta*xx  # alfa+1j*beta
	sqrtXY = np.sqrt(X+delta)
	zPlus  = sqrtXY + sqrtDelta
	zMinus = X / zPlus
	wMinus = cerf(1j*zMinus).real
	wPlus  = cerf(1j*zPlus).real
	kSDV   = wMinus-wPlus

	if verbose:  awrite ((xx, X, sqrtXY, zMinus, zPlus, wMinus, wPlus, kSDV))

	if verbose:   return kSDV, 1j*zMinus, wMinus, 1j*zPlus, wPlus
	else:         return kSDV


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

epsTNH = 3e-8

def sdv_tnh (xx, y, q, cerf=wofz, verbose=False):
	""" Speed dependent Voigt: naiv computation of complex square root difference.
	    (The equivalent of the Voigt function K(x,y), i.e. given q and the Voigt variables x,y).

	    This is the 'naive' version without check if the two arguments zPlus, zMinus are in the same region.
	    Computation of complex square root and
	    computation of complex square root difference using Taylor for Y>>|X| as in TNH. """

	if isinstance(q,(int,float)) and q<0:  q=y*abs(q)

	# sdv effective parameters --- bwb notation
	alfa  = y/q - 1.5
	sqrtDelta = 0.5*(alfa+1.5)/y
	delta     = sqrtDelta**2

	if verbose:
		print ('\nsdv    y', y, '   alfa', alfa, '   sqrtDelta', sqrtDelta, '   delta', delta,
		       '   a/2*sd',  0.5*alfa/sqrtDelta, 'a/2sd + 2sd', 2.0*sqrtDelta+0.5*alfa/sqrtDelta, '\n')

	if isinstance(xx,(int,float)):  xx = np.array([xx])
	kSDV = np.empty_like(xx)

	X    = alfa + 2j*sqrtDelta*xx  # alfa+1j*beta
	sqrtXY = np.sqrt(X+delta)
	zMinus = np.where(abs(X)<=epsTNH*delta, 0.5*X/sqrtDelta, sqrtXY - sqrtDelta)
	zPlus  = sqrtXY + sqrtDelta
	wMinus = cerf(1j*zMinus).real
	wPlus  = cerf(1j*zPlus).real
	kSDV   = wMinus-wPlus

	if verbose:  awrite ((xx, X, sqrtXY, zMinus, zPlus, wMinus, wPlus, kSDV))

	if verbose:   return kSDV, 1j*zMinus, wMinus, 1j*zPlus, wPlus
	else:         return kSDV


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

mpsqrt = mp.sqrt

def sdv_mpm (xx, y, q, verbose=False):
	""" Speed dependent Voigt: naiv computation of complex square root difference.
	    (The equivalent of the Voigt function K(x,y), i.e. given q and the Voigt variables x,y).

	    Computation with multiprecision math. """

	def cerf(z):  return mp.exp(-z*z)*mp.erfc(-1j*z)

	if isinstance(q,(int,float)) and q<0:  q=y*abs(q)

	# sdv effective parameters --- bwb notation
	alfa  = y/q - 1.5
	sqrtDelta = 0.5*(alfa+1.5)/y
	delta     = sqrtDelta**2

	if verbose:
		print ('\nsdv    y', y, '   alfa', alfa, '   sqrtDelta', sqrtDelta, '   delta', delta,
		       '   a/2*sd',  0.5*alfa/sqrtDelta, 'a/2sd + 2sd', 2.0*sqrtDelta+0.5*alfa/sqrtDelta, '\n')

	if isinstance(xx,(int,float)):  xx = np.array([xx])
	kSDV = np.empty_like(xx);   zMinus_list = []

	for i,x in enumerate(xx):
		X      = alfa + 2j*sqrtDelta*x  # alfa+1j*beta
		sqrtXY =  mpsqrt(mp.fadd(X,delta))
		zPlus  = sqrtXY + sqrtDelta
		zMinus =  X / zPlus;            zMinus_list.append(zMinus)

		wMinus = cerf(1j*zMinus).real
		wPlus  = cerf(1j*zPlus).real
		kSDV[i]   = wMinus-wPlus

	if verbose:  awrite ((xx, X, sqrtXY, zMinus, zPlus, wMinus, wPlus, kSDV))

	if verbose:   return kSDV, np.array([zms.real for zms in zMinus_list])
	else:         return kSDV