# matrix elements calulation in dipole approximation for transition from
# state |n1,l1,m1> to state |n2,l2,m2> 
import numpy as np
from scipy.special import sph_harm, genlaguerre, factorial
from scipy.integrate import simps
import warnings
warnings.filterwarnings("ignore")


a0 = 1.0       # Bohr-radius is equal to 1 (dimensionless r-vector in a0 units)
print ("    ") 
print ("----------------------------------------------------------------------------------------") 
print ("|        Progrm calculate the matrix elements <n1,l1,m1|x,y,z|n2,l2,m2>                |") 
print ("|        for transition from state |n1,l1,m1> to |n2,l2,m2> in hydrogen atom.          |")
print ("|                                                                                      |")
print ("| Enter three quantum numbers n1,l1,m1 for start state and n2,l2,m2 for end statе,     |")
print ("| please enter data separated by spaces, for example: 1 2 3                            |")
print ("----------------------------------------------------------------------------------------") 
n1, l1, m1 = input("Enter quantum numbers n1  l1  m1 for start state: ").split()
n2, l2, m2 = input("Enter quantum numbers n2  l2  m2 for end state: ").split()
def check(n,l,m,k):
 if n<1:
  print ("ERROR: n",k,"=",n,"is wrong")
  exit()
 if l<0 or l>n-1:
  print ("ERROR: l",k,"=",l,"is wrong")
  exit()
 if m<-l or m>l:
  print ("ERROR: m",k,"=",m,"is wrong")
  exit()
 return
n1=int(n1)
l1=int(l1)
m1=int(m1)
check(n1,l1,m1,1)
n2=int(n2)
l2=int(l2)
m2=int(m2)
check(n2,l2,m2,2)

# --- radial part of wavefunction  ---
def R_nl(r, n, l):
    rho = 2.0 * r / (n * a0)
    pref = np.sqrt((2.0 / (n * a0))**3 *
                   factorial(n - l - 1) /
                   (2.0 * n * factorial(n + l)))
    L = genlaguerre(n - l - 1, 2*l + 1)(rho)
    return pref * np.exp(-rho/2) * rho**l * L

# --- Radial integral  ---
def radial_integral(n1,l1,n2,l2):
    r = np.linspace(0, 100, 10000)
    R1 = R_nl(r, n1, l1)
    R2 = R_nl(r, n2, l2)
    return simps(R1 * R2 * r**3, r)


# --- Spherical integral ---
def angular_integral(l1,m1,l2,m2):
    theta = np.linspace(0, np.pi, 500)
    phi = np.linspace(0, 2*np.pi, 500)

    Theta, Phi = np.meshgrid(theta, phi, indexing='ij')

    Y1 = sph_harm(m1, l1, Phi, Theta)
    Y2 = sph_harm(m2, l2, Phi, Theta)

    integrand_z = np.conj(Y1) * np.cos(Theta) * Y2 * np.sin(Theta)
    integrand_x = np.conj(Y1) * np.sin(Theta) *np.cos(phi) * Y2 * np.sin(Theta)
    integrand_y = np.conj(Y1) * np.sin(Theta) *np.sin(phi) * Y2 * np.sin(Theta)

    # integrals calculation for different coordinate of dipole approximation 
    int_phi_x = simps(integrand_x, phi, axis=1)
    int_phi_y = simps(integrand_y, phi, axis=1)
    int_phi_z = simps(integrand_z, phi, axis=1)
    return simps(int_phi_x, theta),simps(int_phi_y, theta),simps(int_phi_z, theta)

# --- full matrix elements for x,y and z coordinate ---
def dipole_matrix(n1,l1,m1,n2,l2,m2):
    R = radial_integral(n1,l1,n2,l2)
    A = angular_integral(l1,m1,l2,m2)
    return R * A[0],R * A[1],R * A[2]
dx,dy,dz=np.abs(dipole_matrix(n1,l1,m1,n2,l2,m2))
if dx<1.e-8:
 select_x="forbidden"
else:
 select_x="allowed"

if dy<1.e-8:
 select_y="forbidden"
else:
 select_y="allowed"

if dz<1.e-8:
 select_z="forbidden"
else:
 select_z="allowed"
 
#print ("    ") 
print ("    ") 
print("The perturbation matrix elements calculation in dipol approximation")
print("============n1==l1==m1=====n2==l2==m2===============================")
print("transition ",n1," ",l1," ",m1,"--->",n2," ",l2," ",m2)
print("======================================== transition is: ============")

print("<",n1,l1,m1,"| x |",n2,l2,m2,"> =",'{0:10.5e}'.format(dx),"  =>",select_x)
print("<",n1,l1,m1,"| y |",n2,l2,m2,"> =",'{0:10.5e}'.format(dy),"  =>",select_y)
print("<",n1,l1,m1,"| z |",n2,l2,m2,"> =",'{0:10.5e}'.format(dz),"  =>",select_z)