#!/usr/bin/env python3
"""
generate_radius_table_data.py

Generate exact data for the radius approaches comparison table.
"""

import numpy as np

# Constants
HBAR = 1.054571817e-34
ME = 9.1093837015e-31
E = 1.602176634e-19
K = 8.9875517923e9
A0 = 5.29177210903e-11

def calculate_force_approaches(element_name, shells_data):
    """Calculate forces using different radius approaches"""
    
    results = {}
    
    # Extract shell information
    radii = []
    electrons = []
    z_effs = []
    
    for n, l, num_e, z_eff in shells_data:
        r = n * A0 / z_eff
        if l == 2:  # d-orbital
            r *= 0.35
        radii.append(r)
        electrons.append(num_e)
        z_effs.append(z_eff)
    
    # Total electrons
    total_e = sum(electrons)
    
    # 1. Outermost approach (our method)
    r_outer = radii[-1]
    z_eff_outer = z_effs[-1]
    s = 1 if len(radii) == 1 else 1  # simplified
    F_outer = HBAR**2 * s**2 / (ME * r_outer**3)
    F_coulomb_outer = K * z_eff_outer * E**2 / r_outer**2
    
    # 2. Mean radius
    r_mean = sum(r * n_e for r, n_e in zip(radii, electrons)) / total_e
    # Use average Z_eff
    z_eff_mean = sum(z * n_e for z, n_e in zip(z_effs, electrons)) / total_e
    F_mean = HBAR**2 * s**2 / (ME * r_mean**3)
    F_coulomb_mean = K * z_eff_mean * E**2 / r_mean**2
    
    # 3. RMS radius
    r_rms = np.sqrt(sum(r**2 * n_e for r, n_e in zip(radii, electrons)) / total_e)
    F_rms = HBAR**2 * s**2 / (ME * r_rms**3)
    
    # 4. Innermost (1s)
    r_inner = radii[0]
    z_eff_inner = z_effs[0]
    F_inner = HBAR**2 * s**2 / (ME * r_inner**3)
    F_coulomb_inner = K * z_eff_inner * E**2 / r_inner**2
    
    results = {
        'element': element_name,
        'F_outer': F_outer,
        'F_mean': F_mean,
        'F_rms': F_rms,
        'F_inner': F_inner,
        'F_coulomb': F_coulomb_outer,  # Use outer for comparison
        'agreement_outer': (F_outer/F_coulomb_outer) * 100,
        'agreement_mean': (F_mean/F_coulomb_outer) * 100,
        'agreement_rms': (F_rms/F_coulomb_outer) * 100,
        'agreement_inner': (F_inner/F_coulomb_outer) * 100
    }
    
    return results

# Define test elements
elements = {
    'Hydrogen': [(1, 0, 1, 1.00)],
    'Helium': [(1, 0, 2, 1.69)],
    'Carbon': [(1, 0, 2, 5.67), (2, 0, 2, 3.22), (2, 1, 2, 3.14)],
    'Neon': [(1, 0, 2, 9.64), (2, 0, 2, 6.52), (2, 1, 6, 6.52)],
    'Iron': [(1, 0, 2, 25.38), (2, 0, 2, 22.36), (2, 1, 6, 22.36),
             (3, 0, 2, 13.18), (3, 1, 6, 13.18), (3, 2, 6, 9.1), (4, 0, 2, 3.75)]
}

# Calculate and print results
print("LaTeX Table Data for Radius Approaches Comparison")
print("=" * 70)

for elem_name, shells in elements.items():
    results = calculate_force_approaches(elem_name, shells)
    
    print(f"\n{elem_name}:")
    print(f"  Forces (N):")
    print(f"    Outermost:  {results['F_outer']:.2e}")
    print(f"    Mean:       {results['F_mean']:.2e}")
    print(f"    RMS:        {results['F_rms']:.2e}")
    print(f"    Innermost:  {results['F_inner']:.2e}")
    print(f"    Coulomb:    {results['F_coulomb']:.2e}")
    print(f"  Agreement (%):")
    print(f"    Outermost:  {results['agreement_outer']:.2f}")
    print(f"    Mean:       {results['agreement_mean']:.2f}")
    print(f"    RMS:        {results['agreement_rms']:.2f}")
    print(f"    Innermost:  {results['agreement_inner']:.2f}")

print("\n" + "=" * 70)
print("LaTeX formatted (forces):")
print("Element & Outermost & Mean & RMS & Innermost & Coulomb \\\\")
for elem_name, shells in elements.items():
    r = calculate_force_approaches(elem_name, shells)
    print(f"{elem_name:9} & ${r['F_outer']:.2e}$ & ${r['F_mean']:.2e}$ & "
          f"${r['F_rms']:.2e}$ & ${r['F_inner']:.2e}$ & ${r['F_coulomb']:.2e}$ \\\\")

print("\nLaTeX formatted (agreement %):")
print("Element & Outermost & Mean & RMS & Innermost \\\\")
for elem_name, shells in elements.items():
    r = calculate_force_approaches(elem_name, shells)
    print(f"{elem_name:9} & {r['agreement_outer']:.2f} & {r['agreement_mean']:.2f} & "
          f"{r['agreement_rms']:.2f} & {r['agreement_inner']:.2f} \\\\")