
"""PKTron Chemistry — Quantum Chemistry helpers."""
import numpy as np
_R = np.float32

h_eV   = _R(4.1357e-15)   # eV·s
hbar   = _R(1.0546e-34)   # J·s
me     = _R(9.109e-31)    # kg
a0     = _R(5.292e-11)    # Bohr radius m
E_h    = _R(27.211)       # Hartree in eV

def hydrogen_energy_level(n):
    """E_n = -13.6 / n^2 eV"""
    return float(-_R(13.6)/_R(n)**2)

def photon_energy_eV(wavelength_nm):
    """E = hc/lambda  in eV"""
    hc_eV_nm = _R(1239.8)
    return float(hc_eV_nm/_R(wavelength_nm))

def transition_wavelength_nm(n1, n2):
    """Wavelength of H-atom transition n1->n2 (nm)."""
    E1 = hydrogen_energy_level(n1)
    E2 = hydrogen_energy_level(n2)
    dE = abs(E2-E1)
    return float(_R(1239.8)/dE)

def hydrogen_wavefunction_1s(r_bohr):
    """Psi_1s(r) = (1/sqrt(pi))*exp(-r) in atomic units."""
    r = np.asarray(r_bohr, dtype=_R)
    return np.exp(-r)/np.sqrt(_R(np.pi))

def radial_probability_1s(r_bohr):
    """P(r) = 4*pi*r^2 * |psi|^2"""
    r   = np.asarray(r_bohr, dtype=_R)
    psi = hydrogen_wavefunction_1s(r)
    return 4*_R(np.pi)*r**2*psi**2

def molecular_orbital_energy_huckel(alpha, beta, k, n):
    """Huckel MO energy: E_k = alpha + 2*beta*cos(k*pi/(n+1))"""
    return float(_R(alpha)+2*_R(beta)*np.cos(_R(k)*_R(np.pi)/(_R(n)+1)))

def lennard_jones(r, epsilon, sigma):
    """LJ potential V = 4*eps*[(sigma/r)^12 - (sigma/r)^6]"""
    r = np.asarray(r, dtype=_R)
    sr6 = (_R(sigma)/r)**6
    return 4*_R(epsilon)*(sr6**2 - sr6)
