"""
outputs the value of the associated Laguerre function L_p,qmp(x) at a given point.
p and qmp must be non-negative integers.
"""
#==============================================================================
# There are 3 main functions here.
# The first returns the Laguerre function as a poly1d objectfor a given q.
# The second returns a function of one variable that is the associated Laguerre polynomial
# The third returns the value of the normalized radial wave function at r for any values of n, l
#==============================================================================
from math import sqrt, factorial
import numpy as np
import scipy.constants as c

def laguerre(q):
    """returns the Laguerre polynomial as a nump[y poly1d object for a given q"""
    def pascal(rows): 
        """returns one row of Pascal's triangle"""
        for rownum in range (rows):
            newValue = 1
            psc = [newValue]
            for iteration in range (rownum):
                newValue = newValue * (rownum - iteration) * 1 / (iteration + 1)
                psc.append(int(newValue))
        return psc

    coef = []
    for term in xrange(q + 1):
        c = ((-1)**(q + term))*(pascal(q + 1)[term])*(factorial(q)/factorial(q - term))
        coef.append(c)
    return np.poly1d(coef) #returns a poly1d object

def assoc_laguerre(p, qmp):
    """returns the associated Laguerre function for p, qmp"""
    q = qmp + p
    lag = laguerre(q)
    lag_der = lag.deriv(p)

    def getassoc(x):
        value = (-1**p) * lag_der(x)
        return value

    return getassoc

def radial_wave_func(n, l, r):
    """value of the radial wave function corresponding to n, l at a point r"""
    a=c.physical_constants['Bohr radius'][0]
    p = 2*l +1
    qmp = n - l -1
    x = (2*r)/(n*a)
    lfunc= assoc_laguerre(p,qmp)
    y = lfunc(x)
    norm_rad_sol = (sqrt(((2/(n*a))**3)*((factorial(qmp)/(2.0*n*(factorial(n +l))**3))))*np.exp(-r/(n*a))*x**l*y)/a**(-1.5)
    return np.round(norm_rad_sol,5)

#test cases
a=c.physical_constants['Bohr radius'][0]
print 'radial_wave_func(1,0,a)'
ans=radial_wave_func(1,0,a)
print ans
print 'radial_wave_func(2,1,a)'
ans=radial_wave_func(2,1,a)
print ans
print 'radial_wave_func(2,1,2*a)'
ans=radial_wave_func(2,1,2*a)
print ans
print 'radial_wave_func(3,1,2*a)'
ans=radial_wave_func(3,1,2*a)
print ans

#don't mess with this
def test(p,qmp,x):
	f=assoc_laguerre(p,qmp)
	return f(x)