####
#Code to calculate the oscilator strength of RRLs
#by Xunchuan in 2023.01.01
#refer to  
#   van Regemorter et al. 1979 Radial transition integrals involving low or high effective
#                              quantum numbers in the Coulomb approximation
#   Andri Prozesky et al. 2018 MNRAS 478, 2766
#   Brocklehurst 1971 MNRAS 153, 471
#  
#for details
####

import astropy.units as u
import astropy.constants as c
import numpy as np

def fac1(n):
    result = 0.
    for i in range(int(n)):
        i = i+1
        result+=np.log10(i)
    return result

def my2f1(a0,b,y,x):
    assert a0<=0
    F0  = np.log10(1.); sign0 = 1; #in log scale
    Fm1 =  1.-(b/y)*x
    sign1 = 1 if (Fm1>0) else -1;  Fm1 = np.log10(np.abs(Fm1)); 
    
    if a0==-1:
        return Fm1, sign1
    if a0==0:
        return 0., 1
    
    for a in np.arange(1, -int(a0) ):
        a = -a
        
        Fm2_0      = F0 + np.log10( np.abs((1-x)*a) )
        Fm2_0_sign = sign0*(1 if ((1-x)*a)>0. else -1)
        Fm2_1      = Fm1 + np.log10( np.abs(  y-a-b+(b-a)*(1-x)  ) )
        Fm2_1_sign = sign1*(1 if  (  y-a-b+(b-a)*(1-x)  ) >0 else -1)
             
        
        if Fm2_0>(Fm2_1+20):
            Fm2_01      = Fm2_0
            Fm2_01_sign = Fm2_0_sign
        elif Fm2_1>(Fm2_0+20):
            Fm2_01      = Fm2_1
            Fm2_01_sign = Fm2_1_sign   
        else:

            if Fm2_0 > Fm2_1:
                bigger       = Fm2_0
                bigger_sign  = Fm2_0_sign
                smaller      = Fm2_1
                smaller_sign = Fm2_1_sign
            else:
                bigger       = Fm2_1
                bigger_sign  = Fm2_1_sign
                smaller      = Fm2_0
                smaller_sign = Fm2_0_sign  
            Fm2_01      = bigger + np.log10(1.+  bigger_sign*smaller_sign*10**(smaller-bigger) )
            Fm2_01_sign = bigger_sign
            
            
            
        Fm2_01      = Fm2_01+np.log10(np.abs(  (-1/(a-y)) ))
        Fm2_01_sign = Fm2_01_sign*(1 if (-1/(a-y))>0. else -1) 
        
        F0    = Fm1
        sign0 = sign1
        Fm1   = Fm2_01
        sign1 = Fm2_01_sign
    return Fm1, sign1
        
    

def rho2(n1,n2,l):
    """
    rho(n1,l-1,n2,l)
    n2 is the quantum number of the upper or lower level,
    depending on which one has a larger value of l.
    """
    n = n2*1.
    n1 = n1*1.
    result = 0.
    result = result - np.log10(4) - fac1(2*l-1)       # (-1)**(n1-l)/4/fac(2*l-1)
    result = result + 1./2*( fac1(n+l)-fac1(n-l-1)+fac1(n1+l-1)-fac1(n1-l)  )                  
                      #np.sqrt(  fac(n+l)/fac(n-l-1)*fac(n1+l-1)/fac(n1-l)  )
    result = result + (l+1)*np.log10(4*n*n1) - (n+n1)*np.log10(n+n1)
                    #* (4*n*n1)**(l+1)/(n+n1)**(n+n1)
    result = result + (n+n1-2*l-2)*np.log10( np.abs(n-n1) )
                     #* (n-n1)**(n+n1-2*l-2)
#    print(result)    
#     result = 10**result * (
#                         hyp2f1(-n+l+1, -n1+l, 2*l, -4*n*n1/(n-n1)**2)\
#                         - ((n-n1)/(n+n1))**2*hyp2f1( -n+l-1, -n1+l, 2*l, -4*n*n1/(n-n1)**2 )
#                       ) 

    value1, sign1 = my2f1(-n+l+1, -n1+l, 2*l, -4*n*n1/(n-n1)**2)
    value2, sign2 = my2f1(-n+l-1, -n1+l, 2*l, -4*n*n1/(n-n1)**2) 
    value2 = value2+np.log10( np.abs( - ((n-n1)/(n+n1))**2  ) )
    sign2  = sign2*( 1 if ( - ((n-n1)/(n+n1))**2  )>0 else -1  )
    
    if value1>value2+20:
        value = value1
        sign  = sign1
    if value2>value1+20:
        value = value2
        sign  = sign2
    else:
        if value1>value2:
            bigger      = value1
            bigger_sign = sign1
            smaller     = value2
            smaller_sign= sign2
        else:
            bigger      = value2
            bigger_sign = sign2
            smaller     = value1
            smaller_sign= sign1
        value      = bigger + np.log10(1.+  bigger_sign*smaller_sign*10**(smaller-bigger) )
        sign = bigger_sign        

    result = result+value
    result = 10**(result)
    
    result = result**2
    return result

def getAn2n1(n2,n1):
    assert n2>n1
    weight = 0.
    totalrho2 = 0.
    
    for l2 in range(n2):

        weight+=(2*l2+1)
        l1 = l2-1
        if (l1>=0) & (l1<n1):
            totalrho2= totalrho2+rho2(n1, n2,l2)*l2
            
            if np.isnan(totalrho2):
                print(n1,n2,l2)            
        l1 = l2+1
        if (l1>=0) & (l1<n1):
            totalrho2= totalrho2+rho2(n2, n1,l1)*l1
            if np.isnan(totalrho2):
                print(n2,n1,l1, rho2(n2,n1,l1) )  

    totalrho2/=weight
    return totalrho2*(2.6774E9)*(1/n1**2-1/n2**2)**3


def _freq(n,dn):
    return 3.2880512581042156E15/1E6*(1./n**2-1./(n+dn)**2)
def getaccuratefn2n1(n,dn):
    n1 = n
    n2 = n+dn
    factor = 1.
    e2 = (4.8032E-10)**2 * u.cm**3*u.g/u.s**2
    factor = e2*4*np.pi**2/c.m_e/c.c/c.h/(_freq(n,dn)*u.MHz)
    factor = factor*2*c.h*(_freq(n,dn)*u.MHz)**3/c.c**2
    factor = factor.cgs.value
    return getAn2n1(n2,n1)/factor