"""

Copyright Malcolm J. Macleod (malcolm@theprogrammergod.com)

6. Natural Planck units MTP and the Fine Structure Constant
https://ssrn.com/abstract=4346640

 * This software is free to use, modify, and distribute under creative commons
 * https://creativecommons.org/licenses/by-nc-sa/4.0/ (CC CC BY-NC-SA 4.0).
 * Any derivative works or redistributions must include appropriate credit to
 * Malcolm Macleod and this permission notice shall
 * be included in all copies or substantial portions of the Software.

"""

import mpmath as mp
import numpy as np
import matplotlib.pyplot as plt
import time

mp.mp.dps = 24

# ==========================================================
# Choose which parameters to fit.
# Only the ones set True are computed at every (a, r).
# ==========================================================
USE = {
    'R':     False,
    'e':     False,
    'mu0':   False,
    'me':    True,
    'le':    True,
    'ye/ge': True,
    'h':     False,
    'kB':    False,     # <-- only kB is active
}

# ==========================================================
# Fixed constants
# ==========================================================
FIXED_OMEGA = mp.mpf("2.007134954324946198")
C_SPEED     = mp.mpf("299792458.0")
TWO         = mp.mpf(2.0)
THREE       = mp.mpf(3.0)
FOUR        = mp.mpf(4.0)
PI          = mp.pi

# ==========================================================
# CODATA loader
# ==========================================================
def load_constants(codata_2014):
    if codata_2014:
        return {
            'e':        mp.mpf("1.6021766208e-19"),
            'h':        mp.mpf("6.626070040e-34"),
            'lambda_e': mp.mpf("2.4263102367e-12"),
            'R':        mp.mpf("10973731.568508"),
            'ge':       mp.mpf("2.00231930436182"),
            'ye':       mp.mpf("28024951640"),
            'me':       mp.mpf("9.10938356e-31"),
            'mu0':      4 * PI * mp.mpf("1e-7"),
            'kB':       mp.mpf("1.38064852e-23"),
            'tag': '2014',
        }
    else:
        return {
            'e':        mp.mpf("1.602176634e-19"),
            'h':        mp.mpf("6.62607015e-34"),
            'lambda_e': mp.mpf("2.42631023538e-12"),
            'R':        mp.mpf("10973731.568157"),
            'ge':       mp.mpf("2.00231930436092"),
            'ye':       mp.mpf("28024951386.1"),
            'me':       mp.mpf("9.1093837139e-31"),
            'mu0':      mp.mpf("1.25663706212e-6"),
            'kB':       mp.mpf("1.380649e-23"),
            'tag': '2022',
        }

# ==========================================================
# Build prefactors AND evaluation lambdas ONLY for active params.
#
#   Each entry is (prefactor, f(a, r) -> ai)
#   The prefactor is folded into a closure so we compute it once.
# ==========================================================
def build_active(C):
    c = C_SPEED
    O = FIXED_OMEGA
    prefs = {}          # name -> prefactor (for reporting)
    evals = {}          # name -> callable(a, r) -> ai

    if USE['R']:
        pref = c**5 / (TWO**28 * THREE**3 * PI**16 * O**27)
        prefs['R'] = pref
        evals['R'] = lambda a, r, p=pref: C['R'] * a**5 * r**9 / p

    if USE['e']:
        pref = TWO**10 * PI**7 * O**9 / c**3
        prefs['e'] = pref
        evals['e'] = lambda a, r, p=pref: C['e'] * a / (p * r**3)

    if USE['mu0']:
        pref = 1 / (TWO**11 * PI**5 * O**4)
        prefs['mu0'] = pref
        evals['mu0'] = lambda a, r, p=pref: C['mu0'] / (p * a * r**7)

    if USE['me']:
        pref = 1 / (TWO**19 * PI**7 * THREE**3 * O**13 * c)
        prefs['me'] = pref
        evals['me'] = lambda a, r, p=pref: C['me'] * a**3 / (p * r**4)

    if USE['le']:
        pref = TWO**27 * PI**16 * THREE**3 * O**27 / c**5
        prefs['le'] = pref
        evals['le'] = lambda a, r, p=pref: C['lambda_e'] / (p * a**3 * r**9)

    if USE['ye/ge']:
        pref = TWO**27 * PI**13 * THREE**3 * O**22 / c**2
        prefs['ye/ge'] = pref
        evals['ye/ge'] = lambda a, r, p=pref: C['ye'] * r / (C['ge'] * p * a**2)

    if USE['h']:
        pref = TWO**8 * PI**9 * O**14 / c**5
        prefs['h'] = pref
        evals['h'] = lambda a, r, p=pref: C['h'] / (p * r**13)

    if USE['kB']:
        # kB* = a * pi^2 * Omega^5 * r^10 / (4 c^3)
        pref = PI**2 * O**5 / (FOUR * c**3)
        prefs['kB'] = pref
        evals['kB'] = lambda a, r, p=pref: C['kB'] / (p * a * r**10)

    return prefs, evals

# ==========================================================
# Residual: sum |ai - 1| over active params only
# ==========================================================
def residual(a, r, evals):
    total = mp.mpf(0)
    for fn in evals.values():
        total += abs(fn(a, r) - 1)
    return total

# ==========================================================
# Sweep
# ==========================================================
def run_sweep(use_2014):
    C = load_constants(use_2014)
    prefs, evals = build_active(C)
    active = list(evals.keys())
    tag = C['tag']

    print(f"\n{'='*60}")
    print(f"  CODATA {tag}   |   active: {active}")
    print(f"{'='*60}")

    A_START = mp.mpf("137.0359890")
    A_END   = mp.mpf("137.036000")
    A_STEP  = mp.mpf("1e-8")

    R_START = mp.mpf("0.71256212")
    R_END   = mp.mpf("0.71256292")
    R_STEP  = mp.mpf("1e-10")

    t0 = time.time()
    a_vals, best_r_vals, best_res_vals = [], [], []
    gb_res, gb_a, gb_r = mp.inf, None, None

    a = A_START
    while a <= A_END + A_STEP / 2:
        best_r, best_res = None, mp.inf
        r = R_START
        while r <= R_END + R_STEP / 2:
            res = residual(a, r, evals)
            if res < best_res:
                best_res, best_r = res, r
            r += R_STEP

        a_vals.append(float(a))
        best_r_vals.append(float(best_r))
        best_res_vals.append(float(best_res))

        if best_res < gb_res:
            gb_res, gb_a, gb_r = best_res, a, best_r
        a += A_STEP

    print(f"Sweep done in {time.time()-t0:.1f}s")
    print(f"Global best: a = {float(gb_a):.12f}  "
          f"r = {float(gb_r):.15f}  res = {float(gb_res):.3e}")

    # Report each ai at the global best
    print("\n  ai values at global minimum:")
    for name, fn in evals.items():
        ai = fn(gb_a, gb_r)
        print(f"    a_{name:>5s} = {float(ai):.15f}  "
              f"(dev = {float(abs(ai-1)):.3e})")

    a_arr       = np.array(a_vals)
    best_r_arr  = np.array(best_r_vals)
    best_res_arr= np.array(best_res_vals)

    # Local minima
    idx_min = []
    for i in range(1, len(best_res_arr) - 1):
        if best_res_arr[i] < best_res_arr[i-1] and best_res_arr[i] < best_res_arr[i+1]:
            idx_min.append(i)
    if len(best_res_arr) > 1:
        if best_res_arr[0] < best_res_arr[1]:
            idx_min.insert(0, 0)
        if best_res_arr[-1] < best_res_arr[-2]:
            idx_min.append(len(best_res_arr) - 1)

    minima_info = sorted(
        [(a_arr[i], best_r_arr[i], best_res_arr[i], i) for i in idx_min],
        key=lambda x: x[2],
    )

    return {
        'tag': tag,
        'a_arr': a_arr,
        'best_r_arr': best_r_arr,
        'best_res_arr': best_res_arr,
        'minima_info': minima_info,
        'global_best_a': float(gb_a),
        'global_best_r': float(gb_r),
        'evals': evals,
    }

# ==========================================================
# Run
# ==========================================================
data_2014 = run_sweep(use_2014=True)
data_2022 = run_sweep(use_2014=False)

# ==========================================================
# Plot
# ==========================================================
color_2014 = 'blue'
color_2022 = 'darkorange'
active_str = ", ".join(k for k, v in USE.items() if v)

plt.figure(figsize=(16, 6))

plt.subplot(1, 2, 1)
plt.plot(data_2014['a_arr'], np.log10(data_2014['best_res_arr']),
         color=color_2014, lw=1.5, label=data_2014['tag'])
plt.plot(data_2022['a_arr'], np.log10(data_2022['best_res_arr']),
         color=color_2022, lw=1.5, label=data_2022['tag'])
for (av, rv, resv, _) in data_2014['minima_info']:
    plt.plot(av, np.log10(resv), 'o', color=color_2014, ms=6, zorder=5)
for (av, rv, resv, _) in data_2022['minima_info']:
    plt.plot(av, np.log10(resv), 'o', color=color_2022, ms=6, zorder=5)
plt.xlabel(r'$\alpha^{-1}$ (a)', fontsize=14)
plt.ylabel(r'$\log_{10}$(min residual)', fontsize=14)
plt.title(f'Min residual vs a  |  active: {active_str}', fontsize=13)
plt.grid(True, ls=':')
plt.legend(fontsize=12)

plt.subplot(1, 2, 2)
def do_fine_sweep(d, color):
    a_val = mp.mpf(str(d['global_best_a']))
    r_val = mp.mpf(str(d['global_best_r']))
    evals = d['evals']
    r0 = r_val - mp.mpf("2e-6")
    r1 = r_val + mp.mpf("2e-6")
    dr = mp.mpf("2e-8")
    rs, es = [], []
    r = r0
    while r <= r1 + dr/2:
        rs.append(float(r))
        es.append(float(residual(a_val, r, evals)))
        r += dr
    plt.plot(rs, np.log10(es), color=color, lw=2, label=f"min {d['tag']}")

do_fine_sweep(data_2014, color_2014)
do_fine_sweep(data_2022, color_2022)
plt.xlabel('r', fontsize=14)
plt.ylabel(r'$\log_{10}$(residual)', fontsize=14)
plt.title('Fine r sweep around global minimum', fontsize=13)
plt.grid(True, ls=':')
plt.legend(fontsize=12)

plt.tight_layout()
plt.show()