#!/usr/bin/env nix-script
#!>python3
#! python3 | sympy

from sympy.physics.units import *
from sympy import symbols, solve, Eq, pi

##
## Bohr model for the hydrogen atom
##

# Constants
ε_0, e, m_e, h, c, π = symbols('ε_0 e m_e h c π')

values = {
	e  : 1.602176565e-19 * C,
	m_e: 9.10938291e-31  * kg,
  	h  : planck,
  	ε_0: electric_constant,
  	c  : speed_of_light,
  	π  : pi }

# Variables
v, r, λ, ν, n = symbols('v r λ v n')


##
## Electron velocity
##

# Coulomb force is a centripetal force
F_c = m_e * v**2/r
F   = 1/(4*π*ε_0) * e**2 / r**2

# Equate and solve for the velocity (taking the positive solution)
v = solve(Eq(F_c, F), v)[1]


##
## Energy levels
##

# Total energy of the electron
U = -e**2 / (4*π * ε_0 * r)
K = m_e * v**2 / 2
E = U + K

# De Broglie wavelength of the electron
λ = h / (m_e*v)

# Hypothesize the orbits are quantized
C_o = 2*π * r
C_q = n * λ

# Solve for the radius
r_n = solve(Eq(C_o, C_q), r)[0]

# Bohr radius is the radius of the first orbit
r_1 = r_n.subs(n, 1)

# Substitute r_n in the energy equation
E_n = E.subs(r, r_n)

# The ground state is lowest energy level
E_1 = E_n.subs(n, 1)


##
## Spectum of emission
##

# Energy of the emitted photon from Planck–Einstein relation
λ = symbols('λ')
ν = c/λ
E_γ = h*ν

# Energy radiated by the electron
# for n > m
E_n                   # initial state
E_m = E_n.subs(n, m)  # final state
E = E_n - E_m

# Equate and solve for λ
λ = solve(Eq(E_γ, E), λ)[0]

# Find the inverse wavelength
λ_1 = 1/λ

# Substitute the Rydberg constant
R_h = e**4 * m_e / (8 * c * ε_0**2 * h**3)
λ_1 = λ_1.subs(R_h, symbols('R_h'))


value = lambda x: x.subs(values).evalf(10)
results = '''
Bohr radius: {r_1} ≈ {vr_1}
Ground state / Rydberg energy: {E_1} ≈ {vE_1}
Rydberg constant: {R_h} ≈ {vR_h}

Energy of the nth level: E_n = {E_n}
Rydberg equation: 1/λ = {λ_1}
''' .format(vr_1=value(r_1),
	        vE_1=value(E_1),
	        vR_h=value(R_h),
	        **locals()
	).replace('**','^').replace('*', ' ')

print(results)