import arts_math
from Numeric import *


def LaguerrePoly(x,alpha,n):
    if n==0:
        return 1.0,0.0
    elif n==1:
        return -x + alpha + 1.0,-1.0
    else:
        Lnminus1=LaguerrePoly(x,alpha,n-1)[0]
        Lnminus2=LaguerrePoly(x,alpha,n-2)[0]
        Ln=((-x+2*(n-1)+alpha+1.0)*Lnminus1-(n-1.0+alpha)*Lnminus2)/n
        dLnlx=(n*Ln-(n+alpha)*Lnminus1)/x
        return Ln,dLnlx

def LegendrePoly(x,n):
    if n==-1:
        return 0.0,0.0
    if n==0:
        return 1.0,0.0
    Pnminus1,dPnminus1=LegendrePoly(x,n-1)
    Pnminus2,dPnminus2=LegendrePoly(x,n-2)
    Pn=(1.0/n)*((2*n-1)*x*Pnminus1-(n-1.0)*Pnminus2)
    dPn=(1.0/n)*((2*n-1)*(Pnminus1+x*dPnminus1)-(n-1.0)*dPnminus2)
    return Pn,dPn

def WigertPoly(x,sigma,n):
    
    if n==-1:
        return 0.0,0.0
    if n==0:
        return 1.0,0.0
    E=exp(sigma**2/2.0)
    if n==1:
        a=E
        b=0.0
    else:
        a=((E**2+1)*E**(2*(n-1))-1)*E**(2*n-3)
        b=(E**(2*(n-1))-1)*E**(6*n-10)
    pminus1,dpminus1=WigertPoly(x,sigma,n-1)
    pminus2,dpminus2=WigertPoly(x,sigma,n-2)
    p=(x-a)*pminus1-b*pminus2
    dp=(x-a)*dpminus1+pminus1-b*dpminus2
    return p,dp
    

def newton(func,xstart,delta_x):
    max_iter=20
    x=xstart
    for i in range(max_iter):
        f,df=func(x)
        dx=f/df
        x-=dx
        if abs(dx)<delta_x:
            return x
    raise 'Root not found'


def laggausdata(n,alpha):
    abscissa=zeros(n,Float)
    weights=zeros(n,Float)
    def Ln(x):
        return LaguerrePoly(x,alpha,n)
    for i in range(n):
        if i==0:
            x=(1.0+alpha)*(3.0+0.92*alpha)/(1.0+2.4*n+1.8*alpha)
        elif i==1:
            x+=(15.0+6.25*alpha)/(1.0+0.9*alpha+2.5*n)
        else:
            ai=i-1
            x+=((1.0+2.55*ai)/(1.9*ai)+1.26*ai*alpha/(1.0+3.5*ai))*\
                (x-abscissa[i-2])/(1.0+0.3*alpha)
        x=newton(Ln,x,1e-5)
        abscissa[i]=x
        
#        weights[i]=-exp(arts_math.gosper(alpha+n-1.0)-arts_math.gosper(n-1.0))/\
#                    (LaguerrePoly(x,alpha,n)[1]*n*LaguerrePoly(x,alpha,n-1)[0])
        if remainder(alpha,1)==0:
            weights[i]=fact(alpha+n)*x/\
                (fact(n)*(n+1)**2*LaguerrePoly(x,alpha,n+1)[0]**2)
        else:
            weights[i]=arts_math.gosper(alpha+n)*x/\
                (fact(n)*(n+1)**2*LaguerrePoly(x,alpha,n+1)[0]**2)

    return abscissa,weights

def gausslegdata(a,b,n):
    m=(n+1)/2
    abscissa=zeros(n,Float)
    weights=zeros(n,Float)
    xm=0.5*(a+b)
    xl=0.5*(b-a)
    def Pn(x):
        return LegendrePoly(x,n)
    
    for i in range(m):
        x=cos(pi*(i+0.75)/(n+0.5))
        x=newton(Pn,x,1e-5)
        abscissa[i]=xm-xl*x
        abscissa[-i-1]=xm+xl*x
        weights[i]=2.0*xl/((1.0-x*x)*Pn(x)[1]**2)
        weights[-1-i]=weights[i]
    return abscissa,weights

def gaussleg(func,a,b,n):
    integral=0.0
    abscissa,weights=gausslegdata(a,b,n)
    for i in range(n):
        integral+=weights[i]*func(abscissa[i])
    return integral


def lag_gauss(func,alpha,n):
    """performs Laguerre-Gauss quadrature i.e estimates the integral
    \int_0^\infty\exp(-x)f(x)dx"""
    integral=0.0
    abscissa,weights=laggausdata(n,alpha)
    for i in range(n):
        integral+=weights[i]*func(abscissa[i])
    return integral

def fact(n):
    f=1
    if n==0:
        return f
    else:
        for i in range(n):
            f*=i+1
        return f