#!/usr/bin/env python

"""
Speed Dependent Voigt (SDV) Profiles.

The 'main functions' of this module, sdvProfile and voigtProfile, evaluate the corresponding profiles
(normalized to one) for given pressure & temperature, and line position & widths.

The module can be directly used from the (Unix etc.) shell with line position as the only mandatory argument
(ifs its not executable, say e.g. 'python sdv.py -h').
The defaults for all optional arguments are set as in the Boone, Walker, Bernath (JQSRT 2011) supplement code;
with position=1350cm-1 the profiles of the two codes should be identical (except for the wavenumber grid).

Most other functions defined here provide implementations of the SDV function based on various codes
(and combinations therof) for the complex error function.
(Several rational approximations due to Humlicek (JQSRT 1979, 1982) and Weideman (SIAM 1994)
are provided in the accompanying files.)

Hint:  the "axy" in the "sdv_axy_...." function names indicates the three arguments alpha, x, y
"""

import sys
from string import join

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

try:                      from scipy.special import wofz
except ImportError, msg:  raise SystemExit (str(msg) + '\nimport scipy failed!')

from ir import *
from weideman import *
from humlicek import *

temperature0 = 296.0  # Hitran/Geisa reference temperature

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

def sdv_axy_naive (x, y, alfa=0.0, cerf=weideman32a, verbose=False):
	""" Evaluate speed-dependent Voigt function  (ignoring pressure induced shift).
	    (The equivalent of the Voigt function K(x,y), i.e. given alfa and the Voigt variables x,y).

	    This is the 'naive' version without check if the two arguments zPlus, zMinus are in the same region.
	    """

	# sdv effective parameters
	sqrtDelta = 0.5*(alfa+1.5)/y
	delta     = sqrtDelta**2
	beta      = 2.0*sqrtDelta*x

	# arguments for the complex error function
	r4  = np.sqrt((delta+alfa)**2+beta**2)
	izz = recSqrt2 * (-np.sign(beta)*np.sqrt(np.maximum(r4 - delta - alfa,0.0)) +1j*np.sqrt(r4 + delta + alfa))
	izp = izz + 1j*sqrtDelta
	izm = izz - 1j*sqrtDelta

	# evaluate the difference terms
	wPlus  = cerf(izp).real
	wMinus = cerf(izm).real
	# and return normalized profile
	if verbose:  return wMinus, wPlus, wMinus-wPlus
	else:        return                wMinus-wPlus


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

def sdv_axy_hum2wei32c (xxx, yyy, alfa=0.0, sCut=10.0):
	""" Evaluate speed-dependent Voigt function  (ignoring pressure induced shift).
	    (The equivalent of the Voigt function K(x,y), i.e. given alfa and the Voigt variables x,y).

	    This version uses the Humlicek-Weideman combination and ensures that both cerf are done with the same approximation.
	    (In contrast to the function below this is vectorized.)
	    """

	if not xxx.shape==yyy.shape:  raise SystemExit ("ERROR -- sdv_axy_huwe:  inconsistent shapes")

	# sdv effective parameters
	sqrtDelta = 0.5*(alfa+1.5)/yyy
	delta     = sqrtDelta**2
	beta      = 2.0*sqrtDelta*xxx

	# arguments for the complex error function
	r4  = np.sqrt((delta+alfa)**2+beta**2)
	r4m = r4 - delta - alfa
	r4p = r4 + delta + alfa
	izz = recSqrt2 * (-np.sign(beta)*np.sqrt(np.maximum(r4m,0.0)) +1j*np.sqrt(r4p))
	izp = izz + 1j*sqrtDelta
	izm = izz - 1j*sqrtDelta

	# vectorized version
	maskPlus  = abs(izp.real)+izp.imag>sCut
	maskMinus = abs(izm.real)+izm.imag>sCut
	www       = np.where(np.logical_and(maskPlus, maskMinus),  humlicek2(izm)-humlicek2(izp),
	                                                           weideman32a(izm)-weideman32a(izp))
	return www.real


def sdv_axy_huwe (xxx, yyy, alfa=0.0, sCut=10.0):
	""" Evaluate speed-dependent Voigt function  (ignoring pressure induced shift).
	    (The equivalent of the Voigt function K(x,y), i.e. given alfa and the Voigt variables x,y).

	    This version uses the Humlicek-Weideman combination and ensures that both cerf are done with the same approximation.
	    (In contrast to the function above this uses an explicit loop (slow)!)
	    """

	if not xxx.shape==yyy.shape:  raise SystemExit ("ERROR -- sdv_axy_huwe:  inconsistent shapes")

	# sdv effective parameters
	sqrtDelta = 0.5*(alfa+1.5)/yyy
	delta     = sqrtDelta**2
	beta      = 2.0*sqrtDelta*xxx

	# arguments for the complex error function
	r4  = np.sqrt((delta+alfa)**2+beta**2)
	r4m = r4 - delta - alfa
	r4p = r4 + delta + alfa
	izz = recSqrt2 * (-np.sign(beta)*np.sqrt(np.maximum(r4m,0.0)) +1j*np.sqrt(r4p))
	izp = izz + 1j*sqrtDelta
	izm = izz - 1j*sqrtDelta

	# scalar version
	www = np.zeros_like(izz.flat)
	for i in range(www.size):
		zm, zp    = izm.flat[i], izp.flat[i]
		maskPlus  = abs(zp.real)+zp.imag>sCut
		maskMinus = abs(zm.real)+zm.imag>sCut
		if maskPlus and maskMinus:  wm, wp = humlicek2(zm), humlicek2(zp)
		else:                       wm, wp = weideman32a(zm), weideman32a(zp)
		www[i] = (wm-wp).real
	return www.reshape(xxx.shape)


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

def sdv_axy_w4 (xxx, yyy, alfa=0.0, sCut=20., verbose=False):
	""" Evaluate speed-dependent Voigt function  (ignoring pressure induced shift).
	    (The equivalent of the Voigt function K(x,y), i.e. given alfa and the Voigt variables x,y).

	    This version uses the humlicek w4 code and ensures that both cerf are done with the same approximation.
	    To avoid "RuntimeWarning: overflow encountered in exp", humlicek4 has been replaced by weideman32a.
	    """

	if not xxx.shape==yyy.shape:  raise SystemExit ("ERROR -- sdv_axy_w4:  inconsistent shapes")

	# sdv effective parameters
	sqrtDelta = 0.5*(alfa+1.5)/yyy
	delta     = sqrtDelta**2
	beta      = 2.0*sqrtDelta*xxx

	# arguments for the complex error function
	r4  = np.sqrt((delta+alfa)**2+beta**2)
	r4m = r4 - delta - alfa
	r4p = r4 + delta + alfa
	izz = recSqrt2 * (-np.sign(beta)*np.sqrt(np.maximum(r4m,0.0)) +1j*np.sqrt(r4p))
	izp = izz + 1j*sqrtDelta
	izm = izz - 1j*sqrtDelta

	# initialize
	www = np.zeros_like(izz.flat)
	# and evaluate point by point (not very efficient, but ...)
	for i in range(www.size):
		zm, zp    = izm.flat[i], izp.flat[i]
		if abs(zp.real)+zp.imag>sCut and abs(zm.real)+zm.imag>sCut:
			wm, wp = humlicek1(zm).real, humlicek1(zp).real
		elif abs(zp.real)+zp.imag>5.5 and abs(zm.real)+zm.imag>5.5:
			wm, wp = humlicek2(zm).real, humlicek2(zp).real
		else:
			if zp.imag>=0.195*abs(zp.real)-0.176 and zm.imag>=0.195*abs(zm.real)-0.176:
				wm, wp = humlicek3(zm).real, humlicek3(zp).real
			else:
				wm, wp = weideman32a(zm).real, weideman32a(zp).real
		www[i] = wm-wp

	if verbose:  return izm, izp, www.reshape(xxx.shape)
	else:        return www.reshape(xxx.shape)


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

def sdv_axy_wofz (x, y, alfa=0.0, verbose=False):
	""" Evaluate speed-dependent Voigt function  (ignoring pressure induced shift).
	    (The equivalent of the Voigt function K(x,y), i.e. given alfa=gammaL/gamma2-1.5 and the Voigt variables x,y).

	    This is the reference version using wofz with claimed accuracy of at least 13 significant digits.
	    [Note:  Using this routine is equivalent to sdv_axy_naive(x, y, alfa, wofz)]
	    """

	# sdv effective parameters
	sqrtDelta = 0.5*(alfa+1.5)/y
	delta     = sqrtDelta**2
	beta      = 2.0*sqrtDelta*x

	# arguments for the complex error function
	r4  = np.sqrt((delta+alfa)**2+beta**2)
	r4m = r4 - delta - alfa
	r4p = r4 + delta + alfa
	izz = recSqrt2 * (-np.sign(beta)*np.sqrt(np.maximum(r4m,0.0)) +1j*np.sqrt(r4p))
	izp = izz + 1j*sqrtDelta
	izm = izz - 1j*sqrtDelta

	# evaluate the difference terms and return normalized profile
	if verbose:  return izm, izp, wofz(izm).real, wofz(izp).real, (wofz(izm)-wofz(izp)).real
	else:        return (wofz(izm)-wofz(izp)).real


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

def wavenumber_grid (position, GammaL0=0.1, tExp=0.5, pressure=1.0, temperature=296.0, mass=18.0, sampling=10.0, extension=20.0):
	""" Setup a uniform (equdistant) wavenumber grid centered around `position` appropriate for a Voigt line profile.

	position    center wavenumber
	gammaL0     reference air broadening half width parameter
	tExp        n temperature exponent Lorentz width
	pressure    in units of atm
	temperature Kelvin
	mass        in units of amu
	"""

	# effective widths
	gammaL = pressure * GammaL0 * (temperature0/temperature)**tExp
	gammaD = position * sqrt(2.0*ln2*k*temperature / (amu*mass*c*c))
	gammaV = 0.5*(gammaL+sqrt(gammaL**2+4.0*gammaD**2))  # Voigt width (Whiting approximation)

	# setting up the wavenumber grid
	dv   = gammaV / sampling
	vMax = extension*gammaV

	# evaluate grid and return
	return  np.arange(position-vMax, position+vMax+dv, dv)


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

def voigtProfile (vGrid, position, GammaL0=0.1, tExp=0.5, pressure=1.0, temperature=296.0, mass=18.0):
	""" Evaluate standard Voigt profile  (ignoring pressure induced shift).

	vGrid       wavenumber grid
	position    center wavenumber
	GammaL0     reference air broadening half width parameter
	tExp        n temperature exponent Lorentz width
	pressure    in units of atm
	temperature Kelvin
	mass        in units of amu
	"""

	# effective widths
	gammaL = pressure * GammaL0 * (temperature0/temperature)**tExp
	gammaD = position * sqrt(2.0*ln2*k*temperature / (amu*mass*c*c))

	# Voigt effective parameters
	x = sqrtLn2 * (vGrid-position)/gammaD
	y = sqrtLn2 * gammaL/gammaD
	print 'voigtProfile --- y: ', y

	# evaluate complex error function
	w = wofz(x+1j*y)
	# ... and return normalized profile
	return sqrtLn2*recSqrtPi*w.real/gammaD


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

def sdvProfile (vGrid, position, GammaL0=0.1, GammaL2=0.0, tExp=0.5, pressure=1.0, temperature=296.0, mass=18.0):
	""" Evaluate speed-dependent Voigt profile  (ignoring pressure induced shift).

	vGrid       wavenumber grid
	position    center wavenumber
	GammaL0     reference air broadening half width parameter
	GammaL2     reference speed dependence parameter
	tExp        n temperature exponent Lorentz width
	pressure    in units of atm
	temperature Kelvin
	mass        in units of amu
	"""

	# effective widths
	gammaL = pressure * GammaL0 * (temperature0/temperature)**tExp
	gamma2 = pressure * GammaL2 * (temperature0/temperature)**tExp
	gammaD = position * sqrt(2.0*ln2*k*temperature / (amu*mass*c*c))

	# sdv effective parameters
	alfa = gammaL/gamma2 - 1.5
	beta = (vGrid - position)/gamma2
	sqrtDelta = 0.5*recSqrtLn2 * (gammaD/gamma2)
	delta = sqrtDelta**2
	print 'sdvProfile --- alfa, delta: ', alfa, delta

	# arguments for the complex error function
	r4  = np.sqrt((delta+alfa)**2+beta**2)
	izz = recSqrt2 * (np.sign(beta)*np.sqrt(np.maximum(r4-delta-alfa,0.0)) +1j*np.sqrt(r4+delta+alfa))
	izp = izz + 1j*sqrtDelta
	izm = izz - 1j*sqrtDelta

	# evaluate the difference terms
	wPlus  = wofz(izp)
	wMinus = wofz(izm)
	# compute difference of real parts and normalize
	profile = sqrtLn2*recSqrtPi*(wMinus-wPlus).real/gammaD

	# verbose output for 'debugging'
	#print sqrtDelta, min(izp.imag), max(izp.imag),  min(izm.imag), max(izm.imag), max(sqrtDelta/abs(izz))
	#np.savetxt ('/tmp/abd.xy', np.column_stack([vGrid, beta, r4, r4-delta-alfa, r4+delta+alfa]),'%10.4f  %10g  %10g  %10g %10g',
	            #header ='alfa %g\ndelta %g' % (alfa,delta))
	#np.savetxt ('/tmp/sdv.xy', np.column_stack([vGrid, izp.real,izp.imag, wPlus.real, izm.real,izm.imag, wMinus.real, profile]),
	                                        #'%10.4f  %10g %10g %12g  %10g %10g %12g  %15g')
	return profile


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

if (__name__ == "__main__"):
    import argparse
    parser = argparse.ArgumentParser('sdv', usage='%(prog)s [options] position',
                                     description=__doc__%globals(),
                                     formatter_class=argparse.RawDescriptionHelpFormatter)

    parser.add_argument("position", type=float, help="line position [cm-1]")
    parser.add_argument('-c', "--commentChar",  type=str, default='#', metavar='', help="comment character, default='#'")
    parser.add_argument('-o', "--outFile",  type=argparse.FileType('w'), metavar='', help="output file to save profile")
    parser.add_argument('-q', "--GammaL2", type=float, default=0.009,  metavar='', help="GammaL2")  # speed dependence quadratic term
    parser.add_argument('-a', "--GammaL0", type=float, default=0.08, metavar='', help="GammaL0")    # air broadening coefficient (Lorentz width)
    parser.add_argument('-m', "--mass", type=float, default=18.01, metavar='', help="mass [amu]")
    parser.add_argument('-n', "--tExp", type=float, default=0.59, metavar='', help="n = temperature exponent")
    parser.add_argument('-p', "--pressure", type=float, default=0.2, metavar='', help="pressure")
    parser.add_argument('-T', "--temperature", type=float, default=215.0, metavar='', help="temperature")
    parser.add_argument('-r', "--relative", action="store_true", default=False, help="relative wavenumber grid")
    parser.add_argument('-s', "--sampling", type=float, metavar='', default=10.0, help="number of grid points per (Voigt) hwhm")
    parser.add_argument('-x', "--extension", type=float, metavar='', default=20.0, help="number of (Voigt) width left/right of center")

    args = parser.parse_args()

    if args.outFile: outFile    = args.outFile
    else:            outFile    = sys.stdout

    # set up uniform wavenumber grid
    vGrid = wavenumber_grid (args.position, args.GammaL0, args.tExp, args.pressure, args.temperature, args.mass, args.sampling, args.extension)

    # evaluate speed-dependent Voigt profile (normalized to 1.0)
    sp = sdvProfile (vGrid, args.position, args.GammaL0, args.GammaL2, args.tExp, args.pressure, args.temperature, args.mass)

    # evaluate Voigt profile for reference (normalized to 1.0)
    vp = voigtProfile (vGrid, args.position, args.GammaL0, args.tExp, args.pressure, args.temperature, args.mass)

    listOfComments=['GammaL0  [cm-1] %g' % args.GammaL0,
                    'GammaL2  [cm-1] %g' % args.GammaL2,
                    'tExp            %g' % args.tExp,
                    'mass      [amu] %g' % args.mass,
                    'pressure  [atm] %g' % args.pressure,
                    'temperature [K] %.2f' % args.temperature,
                    'wavenumber    voigt         speed-voigt']
    if args.relative:
        np.savetxt (outFile, np.column_stack([vGrid-args.position, vp, sp]), '%12.7f %12g %12g', header=join(listOfComments,'\n'))
    else:
        np.savetxt (outFile, np.column_stack([vGrid, vp, sp]), '%12.6f %12g %12g', header=join(listOfComments,'\n'))