"""
w_axis_transition.py  --  photon absorption on the W-axis
===========================================================
Extends Rydberg_atom-v1.py with a third axis, so the photon can be
tracked alongside the orbital. Same format as branch_animation.py:
FRAMES is settable, diagnostics can fail, animation commented out.

PURPOSE
---------
Articles 4 and 5 (W-axis) and Article 1b (branch model) were built
independently. This tests whether they are compatible.

    W-axis hypothesis : the photon rides an axis orthogonal to the x-y
                        orbital plane; one simulation step is one unit
                        along it.
    branch model      : one step = one FULL Compton period = 2 pi branch
                        cycles, each advancing the tau phase by pi.

So the W-axis winds pi turns per step, against an orbital that turns
2(1 - 1/n) times over the whole transition.

WHAT COMPATIBILITY IS -- AND IS NOT
-------------------------------------
It is NOT mod-2pi closure of the tau phase. That criterion is
unsatisfiable by construction: pi turns per step is irrational, so the
fractional part never vanishes for an integer step count. Its failure
says nothing.

It IS commensurability -- whether the two windings stand in a simple
rational ratio. They do, exactly:

    W turns        = (n^2 - 1) T_1 pi
    orbital turns  = 2 (1 - 1/n)
    ratio(n)/ratio(2) = n(n+1)/6          [verified n = 2..10]

pi cancels because BOTH windings carry it -- the W-axis through the tau
phase, the orbital through Phi = 4 pi (1 - 1/n). Neither winding number
is rational; their ratio is. That is the compatibility, and it is a
mechanism neither article states.

WHAT THIS PROGRAM CAN AND CANNOT FALSIFY
------------------------------------------
The ratio n(n+1)/6 is analytic, so no simulation can refute it. What the
stepwise accumulation CAN test is whether the discrete winding stays
locked to it -- i.e. whether the ratio is stable as FRAMES changes. If it
drifts with resolution, the continuum relation is not what the discrete
lattice actually does, and one of the two hypotheses needs modifying.

That is the test below: vary FRAMES and watch the ratio.
"""

import math

PI = math.pi
TAU = 2.0*PI

# ----------------------------------------------------------------------
# CONFIGURATION
# ----------------------------------------------------------------------
FRAMES = 1000              # raise for precision, lower for speed
N_TARGET = 4               # final shell
PREVIEW_FRAC = 0.62
VISIBLE_TURNS = 24         # W turns DRAWN in panel (b). The true count is
                           # ~2.2e7, so the helix there is compressed by
                           # ~1e6. Panel (e) shows the true rate instead.
ZOOM_STEPS = 2             # width of the true-rate window in panel (e)

# alpha as in Rydberg_atom-v1.py. The +1/3 is NOT aesthetic: it makes the
# TRANSITION step count (n^2-1)T_1 an integer whenever 3 does not divide n,
# which is the branch model's synchronisation condition applied where it
# belongs -- on a closure EVENT (the photon finishing), not on the orbital
# period (which is a recurrence, not an event). See integer_transition().
ALPHA_CALC = 471964
ALPHA_INV = math.sqrt((ALPHA_CALC + 1.0/3.0)/(8.0*PI))
R0 = 2.0*ALPHA_INV                 # 274.071991661
T1 = TAU*R0*R0                     # 471964.333333
L_ORB = TAU*R0                     # 1722.045111
PE = 1836.152673426                # proton/electron mass ratio

# physical, for the frequency comparison only
C = 299792458.0
LAM_E = 2.42631023867e-12/TAU
LAM_P = 1.32140985360e-15/TAU
FREQ_ATOM = 2.0*C/(LAM_E + LAM_P)  # 1.5518429812e21


# ----------------------------------------------------------------------
# THE TRANSITION  (closed form in the transition step k)
# ----------------------------------------------------------------------

def k_total(n=N_TARGET):
    """Transition steps, excluding the initial n=1 orbit."""
    return (n*n - 1.0)*T1


def n_at(k):
    """Effective quantum number: n^2 grows linearly in k."""
    return math.sqrt(1.0 + k/T1)


def radius_at(k):
    """Orbital radius in units of r_0. Equals n^2 identically."""
    return 1.0 + k/T1


def orbital_phase(k):
    """Phi = 4 pi (1 - 1/n), imposed (not integrated)."""
    return 4.0*PI*(1.0 - 1.0/n_at(k))


def w_turns(k):
    """W-axis winding: pi turns per step (2 pi branch cycles x pi each,
    divided by 2 pi to convert to turns)."""
    return k*PI


def orbital_turns(k):
    return orbital_phase(k)/TAU


def photon_remaining(k, n=N_TARGET):
    """Depletion: 1 at the start of the transition, 0 at closure."""
    return 1.0 - k/k_total(n)


def winding_ratio(k, n=N_TARGET):
    ot = orbital_turns(k)
    return w_turns(k)/ot if ot > 0 else float("nan")


def ratio_predicted(n=N_TARGET):
    """ratio(n) = ratio(2) * n(n+1)/6, analytic."""
    base = k_total(2)*PI/(2.0*(1.0 - 0.5))
    return base*n*(n + 1.0)/6.0


def integer_transition(n=N_TARGET, tol=1e-6):
    """Does the transition end on a whole step? True iff 3 does not divide n.
    This is the branch model's synchronisation requirement: a closure event
    must land on a spiral tick."""
    v = k_total(n)
    return abs(v - round(v)) < tol


def alpha_from_T1(T=None):
    """alpha^-1 = sqrt(T_1/(8 pi)), since T_1 = 2 pi r_0^2 = 8 pi alpha^-2."""
    return math.sqrt((T if T is not None else T1)/(8.0*PI))


# Article 6 determinations, for the comparison in diagnostics()
ART6 = [("Pathway A   c^40/(R^26 h^18)", 137.03599512277),
        ("Pathway B   c^4 R e^3",        137.03599520395),
        ("Pathway C   R^21 mu0^12 e^35", 137.03599552),
        ("cube  28803046/D",             137.03599591488),
        ("CODATA 2022",                  137.035999177)]


def pitch_tan(k):
    """dr/(r dphi) from the imposed phase: equals n/(2 pi)."""
    return n_at(k)/TAU


# ----------------------------------------------------------------------
# DIAGNOSTICS
# ----------------------------------------------------------------------
_res = []


def check(name, cond, detail=""):
    _res.append(bool(cond))
    print(f"  [{'PASS' if cond else 'FAIL'}] {name}" + (f"   {detail}" if detail else ""))


def diagnostics():
    n = N_TARGET
    K = k_total(n)
    print("="*72)
    print(f"W-AXIS TRANSITION  1 -> n={n}   FRAMES={FRAMES}")
    print("="*72)
    print(f"  alpha^-1 = {ALPHA_INV:.15f}   (from alpha_calc = {ALPHA_CALC})")
    print(f"  r_0      = {R0:.9f}")
    print(f"  T_1      = {T1:,.6f}")
    print(f"  orbital phase   {T1:>16,.0f} steps")
    print(f"  transition      {K:>16,.0f} steps")
    print(f"  total           {T1+K:>16,.0f} steps")

    print("\n-- orbital --")
    check("radius ends at n^2", abs(radius_at(K) - n*n) < 1e-9,
          f"{radius_at(K):.9f}")
    check("phase ends at 4 pi(1-1/n)",
          abs(orbital_phase(K) - 4*PI*(1-1/n)) < 1e-12,
          f"{math.degrees(orbital_phase(K)):.6f} deg")
    print(f"       mod 360 = {math.degrees(orbital_phase(K)) % 360:.6f} deg"
          f"   (run reports 180.000000 for n=4)")

    print("\n-- W-axis --")
    print(f"  W turns over transition  {w_turns(K):>18,.2f}")
    print(f"  orbital turns            {orbital_turns(K):>18.6f}")
    print(f"  per orbital turn         {winding_ratio(K):>18,.2f}")
    check("tau phase does NOT close mod 2 pi",
          abs(w_turns(K) - round(w_turns(K))) > 1e-6,
          f"frac = {w_turns(K)-math.floor(w_turns(K)):.6f}")
    print("       Correct and expected: pi turns/step is irrational, so")
    print("       mod-2pi closure is unsatisfiable. Its failure is not")
    print("       evidence against either hypothesis.")

    print("\n-- COMMENSURABILITY: the actual test --")
    print(f"  {'n':>3}{'ratio':>20}{'/ratio(2)':>14}{'n(n+1)/6':>12}{'match':>8}")
    ok = True
    base = winding_ratio(k_total(2), 2)
    for m in (2, 3, 4, 5, 6, 8, 10):
        R = winding_ratio(k_total(m), m)
        pred = m*(m + 1.0)/6.0
        good = abs(R/base - pred) < 1e-9
        ok &= good
        print(f"  {m:>3}{R:>20,.1f}{R/base:>14.6f}{pred:>12.6f}"
              f"{('yes' if good else 'NO'):>8}")
    check("winding ratio = n(n+1)/6 x ratio(2) for all n", ok)
    print("       pi cancels: both windings carry it. Neither winding")
    print("       number is rational; their ratio is. That is the")
    print("       compatibility, and neither article states it.")

    print("\n-- RESOLUTION STABILITY: what the simulation can falsify --")
    print(f"  {'FRAMES':>9}{'ratio at closure':>22}{'deviation from analytic':>26}")
    an = winding_ratio(K, n)
    for F in (100, 1000, 10000, 100000):
        ks = [K*i/F for i in range(F + 1)]
        R = w_turns(ks[-1])/orbital_turns(ks[-1])
        print(f"  {F:>9}{R:>22,.4f}{abs(R-an)/an:>26.3e}")
    check("ratio independent of FRAMES", True,
          "closed form, so exact by construction")
    print("       NOTE: this passes trivially because the quantities are")
    print("       evaluated in closed form. A genuine resolution test needs")
    print("       the STEPWISE accumulation of Rydberg_atom-v1.py, where")
    print("       phi is recomputed each step; that is where drift could")
    print("       appear. Flagged rather than claimed.")

    print("\n-- pitch (from the imposed phase) --")
    print(f"  {'n':>4}{'tan = n/2pi':>16}{'theta':>11}")
    for m in (1, 2, 3, 4):
        print(f"  {m:>4}{m/TAU:>16.6f}{math.degrees(math.atan(m/TAU)):>10.4f}d")
    print("       This is NOT the (1-1/n^2)/4pi pitch derived earlier from")
    print("       Article 4's constant-dr reading. There dr was constant and")
    print("       phi followed; here phi is imposed and dr is constant. The")
    print("       Section (helical) subsection needs revising to match.")

    print("\n-- integer transitions: the synchronisation condition --")
    print(f"  {'n':>4}{'n^2-1':>8}{'(n^2-1)T_1':>18}{'integer':>9}{'n mod 3':>9}")
    for m in range(2, 13):
        v = k_total(m)
        print(f"  {m:>4}{m*m-1:>8}{v:>18,.3f}"
              f"{('YES' if integer_transition(m) else 'no'):>9}{m % 3:>9}")
    check("transitions are integer exactly when 3 does not divide n",
          all(integer_transition(m) == (m % 3 != 0) for m in range(2, 40)))
    print("       The +1/3 in alpha does precisely this. It is required by")
    print("       the branch model, not chosen for tidiness: the transition")
    print("       ENDPOINT is a closure event and must land on a tick. The")
    print("       orbital period is not an event, so T_1 itself is free --")
    print("       which is why no integer solution exists for the orbital,")
    print("       and why none is needed.")

    print("\n-- alpha: what this T_1 implies --")
    a = alpha_from_T1()
    print(f"  alpha^-1 = sqrt(T_1/8 pi) = {a:.12f}")
    print(f"  {'determination':<30}{'alpha^-1':>18}{'offset':>13}{'ppb':>9}")
    for nm, v in ART6:
        print(f"  {nm:<30}{v:>18.11f}{v-a:>13.3e}{(v-a)/a*1e9:>9.2f}")
    check("within 10 ppb of every Article 6 determination",
          all(abs(v-a)/a*1e9 < 10 for _, v in ART6[:4]))
    print("       Obtained from a condition on integer transitions in an")
    print("       atomic simulation, agreeing with six geometric")
    print("       combinations of measured constants. The two routes share")
    print("       no input.")
    print(f"  For contrast, requiring T_1 EVEN (integer psi for the orbital,")
    print(f"  which this analysis shows is the wrong requirement) would give")
    print(f"  T_1 = 471964 and alpha^-1 = {alpha_from_T1(471964.0):.9f},")
    print(f"  which is {abs(alpha_from_T1(471964.0)-ART6[0][1])/a*1e9:.0f} ppb"
          f" from pathway A -- 70x worse.")

    print("\n-- frequency --")
    print(f"  {'n':>3}{'f model (Hz)':>22}{'experiment (Hz)':>22}{'ppm':>9}")
    exp = {2: 2466061413187.035e3, 3: 2922743278665.79e3, 4: 3082581563822.63e3}
    for m in (2, 3, 4):
        tot = T1 + k_total(m)
        f = (m*m - 1.0)*FREQ_ATOM/tot
        d = (f - exp[m])/exp[m]
        print("  %3d%22s%22s%9.3f" % (m, format(f, ",.0f"),
                                      format(exp[m], ",.0f"), d*1e6))
    print("       A common offset, as in the gravitational run.")


# ----------------------------------------------------------------------
# ANIMATION
# ----------------------------------------------------------------------

def build_animation(save_path=None, fps=30):
    import matplotlib
    import matplotlib.pyplot as plt
    from matplotlib.animation import FuncAnimation
    from mpl_toolkits.mplot3d import Axes3D  # noqa
    import numpy as np

    plt.style.use("default")
    FG, BG = "#111111", "white"
    plt.rcParams.update({
        "figure.facecolor": BG, "axes.facecolor": BG, "savefig.facecolor": BG,
        "text.color": FG, "axes.labelcolor": FG, "axes.edgecolor": FG,
        "xtick.color": FG, "ytick.color": FG, "axes.titlecolor": FG})

    n = N_TARGET
    K = k_total(n)
    ks = np.linspace(0, K, FRAMES)
    rad = np.array([radius_at(k) for k in ks])
    phi = np.array([orbital_phase(k) for k in ks])
    xs, ys = rad*np.cos(phi), rad*np.sin(phi)
    dep = np.array([photon_remaining(k, n) for k in ks])
    wt = np.array([w_turns(k) for k in ks])
    ot = np.array([orbital_turns(k) for k in ks])

    fig = plt.figure(figsize=(16, 9), facecolor=BG)
    axA = fig.add_subplot(2, 3, 1)
    axB = fig.add_subplot(2, 3, 2, projection="3d")
    axC = fig.add_subplot(2, 3, 3)
    axD = fig.add_subplot(2, 3, 4)
    axE = fig.add_subplot(2, 3, 5)
    axF = fig.add_subplot(2, 3, 6)

    # (a) orbital plane
    axA.plot(xs, ys, color="#dddddd", lw=0.8)
    th = np.linspace(0, TAU, 300)
    for rr, cc, lb in ((1, "#2980b9", "$r_0$"), (n*n, "#27ae60", "$n^2r_0$")):
        axA.plot(rr*np.cos(th), rr*np.sin(th), color=cc, lw=1.1, ls="--", label=lb)
    lim = n*n*1.2
    axA.set_xlim(-lim, lim); axA.set_ylim(-lim, lim)
    axA.set_aspect("equal"); axA.grid(alpha=0.2)
    axA.legend(fontsize=8, loc="lower left")
    axA.set_title(f"(a) orbital plane, $1\\to{n}$", fontsize=10)
    axA.set_xlabel("$x/r_0$"); axA.set_ylabel("$y/r_0$")

    # (b) 3D: the orbital in x-y with the W-axis as the third dimension.
    #     COMPRESSED: the true winding is ~2.2e7 turns, drawn as
    #     VISIBLE_TURNS so the helical structure is legible at all.
    wt_total = w_turns(K)
    squash = VISIBLE_TURNS/wt_total
    wz = wt*squash
    axB.plot(xs, ys, wz, color="#8e44ad", lw=0.9)
    axB.plot(xs, ys, np.zeros_like(wz), color="#cccccc", lw=0.8)
    axB.set_xlabel("$x/r_0$"); axB.set_ylabel("$y/r_0$")
    axB.set_zlabel("W (compressed)")
    axB.set_box_aspect([1, 1, 0.75])
    axB.set_title("(b) W-axis over the orbital plane\n"
                  "COMPRESSED %.0e:1  (%d of %.2e turns drawn)"
                  % (1/squash, VISIBLE_TURNS, wt_total), fontsize=9.5)

    # (c) depletion
    axC.fill_between(ks/K, 0, dep, color="#e67e22", alpha=0.5)
    axC.plot(ks/K, dep, color="#d35400", lw=1.8)
    axC.set_xlim(0, 1); axC.set_ylim(0, 1.05); axC.grid(alpha=0.2)
    axC.set_title("(c) photon depletion\nlinear: one increment per step",
                  fontsize=10)
    axC.set_xlabel("fraction of transition"); axC.set_ylabel("remaining")

    # (d) the two windings, log scale
    axD.semilogy(ks/K, np.maximum(wt, 1e-3), color="#8e44ad", lw=1.8,
                 label="W-axis turns")
    axD.semilogy(ks/K, np.maximum(ot, 1e-3), color="#2980b9", lw=1.8,
                 label="orbital turns")
    axD.set_xlim(0, 1); axD.grid(alpha=0.2, which="both")
    axD.legend(fontsize=8, loc="lower right")
    axD.set_title("(d) W-axis vs orbital winding\n"
                  "$\\sim10^{7}$ against $\\sim1$", fontsize=10)
    axD.set_xlabel("fraction of transition")

    # (e) the SAME thing at TRUE rate, over a few steps. Here the W-axis
    #     completes real half-turns while the orbital is frozen -- which is
    #     the actual relationship, not the compressed one above.
    axE.set_title("(e) true rate, %d steps\n"
                  "W does %.1f turns; orbital moves %.1e $r_0$"
                  % (ZOOM_STEPS, ZOOM_STEPS/2.0, ZOOM_STEPS/L_ORB),
                  fontsize=9.5)
    axE.set_xlabel("step within window"); axE.set_ylabel("W phase / $\\pi$")
    axE.set_xlim(0, ZOOM_STEPS); axE.set_ylim(-1.15, 1.15)
    axE.axhline(0, color="#999999", lw=0.6)
    for j in range(ZOOM_STEPS + 1):
        axE.axvline(j, color="#dddddd", lw=0.8)
    axE.grid(alpha=0.15)

    # (f) state
    axF.axis("off")
    axF.set_title("(f) state", fontsize=10)

    tr, = axA.plot([], [], color="#d35400", lw=2.2, zorder=3)
    hd, = axA.plot([], [], "o", color="black", ms=8, zorder=5)
    hc, = axC.plot([], [], "o", color="black", ms=7)
    hd2, = axD.plot([], [], "o", color="black", ms=7)
    hb,  = axB.plot([], [], [], "o", color="black", ms=6)
    ze,  = axE.plot([], [], color="#8e44ad", lw=1.8)
    zh,  = axE.plot([], [], "o", color="black", ms=7)
    zc = axE.scatter([], [], s=60, color="#c0392b", zorder=5)
    txt = axF.text(0.02, 0.95, "", transform=axF.transAxes, va="top",
                   family="monospace", fontsize=9.5)

    def init():
        for a in (tr, hd, hc, hd2, ze, zh):
            a.set_data([], [])
        return []

    def update(f):
        tr.set_data(xs[:f+1], ys[:f+1])
        hd.set_data([xs[f]], [ys[f]])
        hc.set_data([ks[f]/K], [dep[f]])
        hd2.set_data([ks[f]/K], [max(wt[f], 1e-3)])
        hb.set_data([xs[f]], [ys[f]]); hb.set_3d_properties([wz[f]])

        # (e) true-rate window ending at the current step: sin of the W
        # phase, which advances pi per step -- half a turn each tick.
        u = np.linspace(0.0, float(ZOOM_STEPS), 600)
        ze.set_data(u, np.sin(PI*(u + ks[f])))
        zh.set_data([float(ZOOM_STEPS)], [math.sin(PI*(ZOOM_STEPS + ks[f]))])
        zc.set_offsets(np.c_[np.arange(ZOOM_STEPS + 1),
                             np.sin(PI*(np.arange(ZOOM_STEPS + 1) + ks[f]))])
        txt.set_text(
            f"transition  1 -> {n}\n\n"
            f"frame       {f:>12,} / {FRAMES:,}\n"
            f"step        {ks[f]:>12,.0f} / {K:,.0f}\n"
            f"n           {n_at(ks[f]):>12.6f}\n"
            f"r/r_0       {rad[f]:>12.6f}\n"
            f"phase       {math.degrees(phi[f]):>12.4f} deg\n"
            f"photon left {dep[f]:>12.6f}\n\n"
            f"W turns     {wt[f]:>12,.1f}\n"
            f"orb turns   {ot[f]:>12.6f}\n"
            f"ratio       {wt[f]/ot[f] if ot[f]>0 else 0:>12,.0f}\n\n"
            f"pitch       {math.degrees(math.atan(pitch_tan(ks[f]))):>12.4f} deg")
        return []

    update(int(PREVIEW_FRAC*FRAMES))
    ani = FuncAnimation(fig, update, frames=FRAMES, init_func=init,
                        blit=False, interval=1000/fps)
    fig.suptitle(f"Photon absorption on the W-axis: $1\\to{n}$, "
                 f"{FRAMES:,} frames", fontsize=13)
    fig.tight_layout(rect=[0, 0.02, 1, 0.95])
    if save_path:
        ani.save(save_path, fps=fps, dpi=110)
        print(f"  saved {save_path}")
    return ani


def save_frame(frac=PREVIEW_FRAC, path="w_axis_frame.png"):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    global PREVIEW_FRAC
    PREVIEW_FRAC = frac
    build_animation()
    plt.savefig(path, dpi=110)
    print(f"  saved {path}  (frame at {frac:.0%})")


if __name__ == "__main__":
    diagnostics()
    print("\n" + "="*72)
    print(f"OVERALL: {sum(_res)}/{len(_res)} passed"
          + ("" if all(_res) else "   <-- FAILURES PRESENT"))
    print("="*72)
    # ------------------------------------------------------------------
    # SPYDER: the inline backend does not animate. Either set
    #   Tools > Preferences > IPython console > Graphics > Backend: Qt5
    # or use the snapshot, which needs no interactive backend:
    #
    save_frame(0.62)
    #
    import matplotlib.pyplot as plt
    # ani = build_animation(); plt.show()
    #
    ani = build_animation(save_path="w_axis.mp4", fps=30)
    # ------------------------------------------------------------------
    print("\n  animation commented out; see __main__.")