import numpy as np
import matplotlib.pyplot as plt

def g0(x):
    """ Eq.(6) of 1986A&A...160..195G """
    return 1./( x + np.sqrt(x*x + 4./np.pi) )

def gc(xp,x):
    """ Eq.(5) of 1986A&A...160..195G """
    # symmetry relation
    if (xp <= 0.):
        xp=abs(xp)
    if (x <=0.):
        x=abs(x)
    #----------
    if (xp <= x):
        return g0(x)
    else:
        return np.exp(x*x - xp*xp)*g0(xp)

def ierfc(x):
    """ Eq.(9) of 1986A&A...160..195G """
    return np.exp(-x*x)*(1.-2*x*g0(x))/np.sqrt(np.pi)

def gw0(xp,x):
    """ Eq.(8a) of 1986A&A...160..195G """
    return ierfc(0.5*abs(xp-x))

def gw1(xp,x):
    """ Eq.(8B) of 1986A&A...160..195G """
    return (2.-xp/x)*ierfc(0.5*abs(xp-x))

def psic(x):
    """ Doppler profile, after Eq.(10) of 1986A&A...160..195G """
    return np.exp(-x*x)/np.sqrt(np.pi) 

def psiw(a,x):
    """ Lorentzian profile, after Eq.(10) of 1986A&A...160..195G """
    return a/(a*a+x*x)/np.pi

def gtrans(apr,a,xp,x):
    """ Eq.(10) of 1986A&A...160..195G """
    rpsi=psic(x)/(psic(x)+psiw(a,x))
    if (apr==0):
        return rpsi*gc(xp,x) + (1.-rpsi)*gw0(xp,x)
    if (apr==1):
        return rpsi*gc(xp,x) + (1.-rpsi)*gw1(xp,x)

def giia(apr,a,xp,x):
    if (abs(x) < 2):
        return gc(xp,x)
    if (2 <= abs(x) <= 4):
        return gtrans(apr,a,xp,x)
    if (abs(x) > 4):
        if (apr==0):
            return gw0(xp,x)
        if (apr==1):
            return gw1(xp,x)

dx=0.1
xp=np.arange(-10.,10.+dx,dx)

# may be 0 or 1 (see Gouttebroze 1986)
apr=1

a=0.001
# for G86 figures' reproduction

plt.figure(1)
for x in [-5.,0., 2., 4., 6., 8.]:
    g2=np.zeros((len(xp)),'Float64')
    for k in np.arange(len(xp)):
        xx=xp[k]
        # how to make this vector, passing all xp's once?
        g2[k]=giia(apr,a,xx,x)
        # --------------------
    plt.plot(xp,g2)
    plt.xlabel("x'",fontsize=18)
    plt.ylabel("g_II-A(x',x)",fontsize=18)
    plt.xlim(left=-7.5,right=10.)
    plt.ylim(top=1.0,bottom=0.)
    plt.xticks(fontsize=12)
    plt.yticks(fontsize=12)

plt.figure(2)
for x in [0., 1., 1.5, 2.,2.5,3.,3.5, 4.]:
#for x in [0., -1., -1.5, -2.,-2.5,-3.,-3.5, -4.]:
    g2=np.zeros((len(xp)),'Float64')
    for k in np.arange(len(xp)):
        xx=xp[k]
        # how to make this vector, passing all xp's once?
        g2[k]=giia(apr,a,xx,x)
        # --------------------
    plt.plot(xp,g2)
    plt.xlabel("x'",fontsize=18)
    plt.ylabel("g_II-A(x',x)",fontsize=18)
    plt.xlim(left=-5.,right=5.)
    plt.ylim(top=1.0,bottom=0.)
    plt.xticks(fontsize=12)
    plt.yticks(fontsize=12)

plt.show()