from lyamc import *
import matplotlib.pyplot as plt
import numpy as np
from scipy import integrate
from scipy.special import wofz
from scipy.special import erfc
import time
from numba import jit
from scipy.special import lambertw



def Voigt(a, x):
   """returns Voigt function using the Faddeeva function
   ! does not work wtih numba !
   """
   z = x + 1.0j*a
   return wofz(z).real


def get_xcw_exact(a):
    """returns xcw that separates core from wing
    """
    return np.abs(np.sqrt(lambertw(-a / np.sqrt(np.pi), k=-1)))


def pdf_parallel(u, x, a):
    H = Voigt(a, x)
    return (a / (np.pi * H)) * np.exp(-u ** 2) / ((x - u)**2 + a**2)


def R(xf, x, mu, a):
    """
    Eq. (71) from Dijkstra (2017)
    """
    f = lambda u: np.exp(-u ** 2) / ((x - u)**2 + a**2) * np.exp(-((xf - x - u * (mu - 1)) ** 2) / (1. - mu**2)) 
    I = integrate.quad(f, -np.inf, np.inf, limit=1000)[0]
    return I * a / (np.pi ** 1.5 * np.sqrt(1. - mu**2) * Voigt(a, x))


def RA(x, xf, a):
    """redistirbution probability for isotropic case
    ! needs to be norbalized !
    ! slow !
    """
    f = lambda mu: R(xf, x, mu, a)
    I = integrate.quad(f, -1., 1., limit=1000)[0]
    return I


def RB(x, xf, a):
    """redistribution probability for dipole Rayleigh case
    ! needs to be norbalized !
    ! slow !
    """
    f = lambda mu: 0.75 * (1 + mu**2) * R(xf, x, mu, a)
    I = integrate.quad(f, -1., 1., limit=1000)[0]
    return I


def J(x, atau):
    """returns Neufeld solution for photons emitted from a single point at line center x=0
    as defined in Laursen (2009) arxiv.org/abs/0805.3153
    """
    return np.sqrt(6 / np.pi) * (x ** 2 / (24 * atau)) / np.cosh(np.sqrt(np.pi ** 3 / 54) * np.abs(x ** 3) / atau)


def test_upar(x=3.0, a=0.00047, N=int(1e5)):
    """tests function get_upar by comparing pdf generated by it with the exact pdf for u parallel
    compare with Laursen (2010) Fig 8.1
    """
    get_upar(x, a)
    t = time.process_time()
    upar = np.linspace(-6, 6, 999)
    hist1 = np.zeros(N)
    for i in range(N):   
        hist1[i] = get_upar(x, a)
    elapsed_time = time.process_time() - t
    print("elapsed time =", elapsed_time)
    plt.hist(hist1, 200, density=True, histtype='step')
    plt.plot(upar, pdf_parallel(upar, x, a))
    plt.xlim(-2, 5)
    plt.ylim(0, 1)
    #plt.yscale('log')
    plt.show()


def test_redistribution(x=3.0, T=1.e4, N=int(1e5)):
    """tests function atom_scattering by comparing redistribution function from x to xout generated by it 
    with the exact function from Hummer (1962) doi=10.1093/mnras/125.1.21
    compare with Laursen (2010) Fig 8.2
    """
    xfs = np.linspace(-4, 8, 200)
    a = get_a(T)
    hist1 = np.zeros(N)
    k = np.array([0., 0., 1.]) #get_random_k()
    vbulk = np.array([0., 0., 0.])
    for i in range(N):
        hist1[i] = atom_scattering(x, k, T, vbulk)[0]
    plt.hist(hist1, 200, density=True, histtype='step')
    # qsA = []
    # qsB = []
    # for xf in xfs:
    #     print("xf =", xf)
    #     qsA.append(RA(x, xf, a))
    #     qsB.append(RB(x, xf, a))
    # qsA /= np.trapz(qsA, dx=dx)
    # qsB /= np.trapz(qsB, dx=dx)
    # np.savetxt("qsA.csv", qsA)
    # np.savetxt("qsB.csv", qsB)
    qsA = np.genfromtxt("data/qsA.csv")
    qsB = np.genfromtxt("data/qsB.csv")
    plt.plot(xfs, qsA, label='isotropic')
    plt.plot(xfs, qsB, label='Rayleigh')
    plt.legend()
    plt.ylim(0, 1)
    plt.show()


def test_Neufeld():
    """tests against Neufeld solution
    compare with Laursen (2010) Fig 7.2
    """
    T = 1.e4
    a = get_a(T)
    x = np.linspace(-17, 17, 1000)
    eta = 0.71
    tau0 = 1.e5
    xs = np.genfromtxt('data/xs_data(tau=1.e5).csv')
    Ns = np.genfromtxt('data/Ns_data(tau=1.e5).csv')
    print("Ns_avr / tau0 =", np.average(Ns) / tau0)
    plt.hist(xs, 50, density=True, histtype='step', color='k')
    plt.plot(x, 4 * np.pi * J(x, eta * a * tau0), label='tau = 1.e5')
    tau0 = 1.e6
    xs = np.genfromtxt('data/xs_data(tau=1.e6).csv')
    Ns = np.genfromtxt('data/Ns_data(tau=1.e6).csv')
    print("Ns_avr / tau0 =", np.average(Ns) / tau0)
    plt.hist(xs, 50, density=True, histtype='step', color='k')
    plt.plot(x, 4 * np.pi * J(x, eta * a * tau0), label='tau = 1.e6')
    plt.legend()
    plt.show()


def test_expansion():
    T = 1.e4
    a = get_a(T)
    x = np.linspace(-50, 50, 1000)
    eta = 0.71
    tau0 = 1.e7
    plt.plot(x, 4 * np.pi * J(x, eta * a * tau0))
    xs = np.genfromtxt('data/xs_data(tau=1.e7, vmax=200).csv')
    plt.hist(xs, 50, density=True, histtype='step', color='r', label='vmax=200')
    xs = np.genfromtxt('data/xs_data(tau=1.e7, vmax=-200).csv')
    plt.hist(xs, 50, density=True, histtype='step', color='b', label='vmax=-200')
    plt.legend()
    plt.title('tau=1.e7')
    plt.show()


# test_upar()
# test_redistribution()
# test_Neufeld()
# test_expansion()

#run 
t = time.process_time()
xs, Ns = runner()
elapsed_time = time.process_time() - t
print("elapsed_time =", elapsed_time)
np.savetxt("xs_data.csv", xs)  
np.savetxt("Ns_data.csv", Ns)
print("Nscattering_avr =", np.average(Ns))


