"""
Copyright Malcolm J. Macleod (malcolm@simulationuniverse.com)

 * 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.

"""

"""
branch_animation.py  --  branch evolution over time
=====================================================
Same format and conventions as spiral_branch.py.

WHAT IS ANIMATED
------------------
One branch, opening at t_start and closing at t_start + PSI, shown over
TWO cycles so the spinor return is visible: the state goes to -xi after
one cycle and back to +xi only after two. A single-cycle animation shows
the mechanics but misses the result.

Panels:
  (a) Psi = A e^{i phi} in the complex plane  -- THE primary panel.
      Sweeps from +Omega on the real axis, through the upper half-plane,
      to -Omega^2. That single trace is the half-cycle, the collapse
      sign, and the origin of spin-1/2, all at once.
  (b) the three monopole phasors on the unit circle, and their product
      (the holonomy), which must return to +1 at closure.
  (c) the monopole triplet in 3D on its fixed cone.
  (d) amplitude and phase against branch time.
  (e) the spinor sign trace over two cycles.

THE TIMESCALE PROBLEM, AND WHAT IS DONE ABOUT IT
--------------------------------------------------
Three rates span eight orders of magnitude:
    scaffolding tau-cycle   0.0667 ticks
    spiral step (one box)   1      tick   (= 15 tau-cycles)
    branch cycle            PSI    ticks  (= ONE slow tau-loop)
At ~20,000 ticks per frame the scaffolding runs ~300,000 tau-cycles per
frame and CANNOT be drawn at true rate. It appears as a separate inset
running on its own time base, labelled as such -- not synchronised, and
not pretending to be.

Likewise the point-state lasts 1 tick, i.e. 5e-5 of a frame. It is drawn
as a flash. Any visible width overstates its duration by ~4 orders of
magnitude, so the caption says so.

WHY THERE IS NO NESTED LOOP
-----------------------------
The pseudocode structure is

    FOR t_spiral = t_start TO t_end:
        FOR t_tau = 1 TO 15:
            generate tau; generate branch

which is PSI*15 = 374,418,240 iterations -- infeasible in Python and
unnecessary. Every branch phase is LINEAR in the elapsed tick count k:

    branch u(k)      = k/PSI          (ONE tau-loop over PSI steps)
    branch phase(k)  = pi * u(k)
    monopole phi_j(k)= pi * n_j * k/PSI

so each frame is evaluated directly from its own k. This is not an
approximation of the loop; it is the loop's closed form. Phases are
computed, never accumulated -- accumulating 3.7e8 increments would drift
by ~1e-8, which is larger than several of the quantities being tested.

NOTE ON PSI AND alpha
-----------------------
PSI must be an integer but the resolution of alpha is 11-12 digits and 
PSI electron to 22 digits so we cannot know if PSI electron is an integer.
"""

import math

# ----------------------------------------------------------------------
# CONSTANTS
# ----------------------------------------------------------------------
PI = math.pi
TAU = 2.0*PI
ALPHA_INV = 137.035999084          # CODATA; see note above
OMEGA = math.sqrt(PI**math.e * math.e**(1.0 - math.e))
CONE = math.acos(1.0/3.0)          # 70.5288 deg, forced by 3 cos t = 1

SIGMA = math.sqrt(2.0*ALPHA_INV)*OMEGA**5
PSI_RAW = SIGMA**3/TAU
PSI = float(round(PSI_RAW))

CYCLES = 2                         # 2 shows the spinor return; 1 does not
NO_FRAMES = round(PSI/10000)*CYCLES
B = 15                             # tau-cycles per spiral step

t_start = 10**12
t_end = t_start + PSI

PREVIEW_FRAC = 0.30                # where a static capture is taken:
                                   # 30% into cycle 1, so the trace has
                                   # length and the phasors have separated

# TURN TRIPLE. The holonomy under the pi convention requires sum(n) EVEN
# (sum phi = pi sum n must be 0 mod 2pi). Precession = (3 n3 - S)/4, and
# a 2-cycle return needs 3 n3 - S = 2. With S even that forces n3 even.
#   (-2,0,0): S=-2 even, n3=0, precession +1/2  -> spinor          OK
#   ( 0,0,1): S=+1 ODD  -- inadmissible, though it also gives +1/2.
# An earlier draft of this file used (0,0,1); the diagnostic caught it.
TURNS = (-2, 0, 0)


# ----------------------------------------------------------------------
# BRANCH STATE  --  all closed-form in k, never accumulated
# ----------------------------------------------------------------------

def branch_u(k):
    """ln(tau) WITHIN the current cycle, 0..1.

    tau RESETS at every cycle boundary (Axiom, growth law), so this is
    the fractional part -- except at an exact boundary, where the cycle
    has just completed and u = 1 (giving A = Omega^2, the closure
    amplitude) rather than 0.
    """
    c = k/PSI
    u = c - math.floor(c)
    if u == 0.0 and k > 0:
        u = 1.0
    return u


def branch_u_total(k):
    """ln(tau) ACCUMULATED across cycles. Used only by the phase."""
    return k/PSI


def amplitude(k):
    """A = Omega^(u+1) with u RESETTING each cycle: Omega -> Omega^2,
    then back to Omega. The amplitude does not compound across cycles."""
    return OMEGA**(branch_u(k) + 1.0)


def tau_phase(k):
    """phi = pi * accumulated u. This does NOT reset -- if it did, every
    closure would give -Omega^2 and the state would never return to
    +Omega^2, so there would be no two-cycle return and no spinor.
    tau resets; the accumulated phase does not."""
    return PI*branch_u_total(k)


def psi_state(k):
    """Psi = A e^{i phi}. Psi(PSI) = -Omega^2 (Euler)."""
    a, p = amplitude(k), tau_phase(k)
    return a*math.cos(p), a*math.sin(p)


def monopole_phases(k, turns=TURNS):
    return [PI*n*k/PSI for n in turns]


def holonomy_product(k, turns=TURNS):
    s = sum(monopole_phases(k, turns))
    return math.cos(s), math.sin(s)


def branch_delta(k, turns=TURNS):
    """N-S azimuth: phi_3 - (phi_1+phi_2)/2."""
    p = monopole_phases(k, turns)
    return p[2] - 0.5*(p[0] + p[1])


def precession(turns=TURNS):
    """Turns per cycle = (3 n3 - S)/4. Read off the dynamics."""
    return (turns[2] - 0.5*(turns[0] + turns[1]))/2.0


def spinor_sign(k):
    """+1 -> -1 -> +1: one loop multiplies the state by e^{i pi} = -1."""
    return math.cos(tau_phase(k))


def admissible(turns=TURNS):
    return sum(turns) % 2 == 0


def monopole_vectors(k, turns=TURNS, axis=(0.0, 0.0, 1.0)):
    """Three unit vectors at CONE from the spin axis, azimuths 120 apart,
    carried round by the N-S azimuth. Orientations, not positions."""
    s = axis
    tmp = (1.0, 0.0, 0.0) if abs(s[0]) < 0.9 else (0.0, 1.0, 0.0)
    u = (tmp[1]*s[2]-tmp[2]*s[1], tmp[2]*s[0]-tmp[0]*s[2], tmp[0]*s[1]-tmp[1]*s[0])
    un = math.sqrt(sum(c*c for c in u)); u = tuple(c/un for c in u)
    w = (s[1]*u[2]-s[2]*u[1], s[2]*u[0]-s[0]*u[2], s[0]*u[1]-s[1]*u[0])
    d = branch_delta(k, turns)
    out = []
    for j in range(3):
        a = d + TAU*j/3
        out.append(tuple(math.cos(CONE)*s[i]
                         + math.sin(CONE)*(math.cos(a)*u[i] + math.sin(a)*w[i])
                         for i in range(3)))
    return out


def unit_sum(k, turns=TURNS):
    ms = monopole_vectors(k, turns)
    tot = tuple(sum(m[i] for m in ms) for i in range(3))
    return math.sqrt(sum(c*c for c in tot))


# ----------------------------------------------------------------------
# SCAFFOLDING TAU  --  own time base; NOT synchronised to the frames
# ----------------------------------------------------------------------

def scaffold_state(phase01):
    """One scaffolding tau-cycle, parameterised 0..1. Shown in an inset
    at an arbitrary rate: at true rate it would run ~300,000 cycles per
    frame and be invisible."""
    a = OMEGA**(phase01 + 1.0)
    p = PI*phase01
    return a*math.cos(p), a*math.sin(p)


# ----------------------------------------------------------------------
# DIAGNOSTICS  (fast; runs without matplotlib)
# ----------------------------------------------------------------------

def diagnostics():
    print("="*70)
    print("BRANCH ANIMATION -- parameters")
    print("="*70)
    print(f"  alpha^-1      = {ALPHA_INV}")
    print(f"  Omega         = {OMEGA:.12f}")
    print(f"  Omega^5       = {OMEGA**5:.10f}")
    print(f"  SIGMA         = {SIGMA:.8f}")
    print(f"  PSI raw       = {PSI_RAW:.6f}")
    print(f"  PSI           = {PSI:,.0f}   (rounded; frac discarded "
          f"{abs(PSI_RAW-PSI):.4f})")
    print(f"  cycles shown  = {CYCLES}")
    print(f"  NO_FRAMES     = {NO_FRAMES}   -> {CYCLES*PSI/NO_FRAMES:,.0f} ticks/frame")
    print(f"  turn triple   = {TURNS}, sum = {sum(TURNS)} "
          f"({'admissible' if admissible() else 'NOT admissible -- sum must be even'})")
    print(f"  precession    = {precession():+.3f} turns/cycle "
          f"-> returns after {abs(1/precession()):.0f} cycles")

    print(f"\n  spiral background")
    print(f"    t_start = {t_start:.4e}  radius {math.sqrt(t_start):,.2f}")
    print(f"    t_end   = {t_end:.4e}  radius {math.sqrt(t_end):,.2f}")
    print(f"    change over one cycle: x{math.sqrt(t_end)/math.sqrt(t_start):.7f}"
          f"  (effectively static -- a late branch)")

    print(f"\n  timescales")
    for nm, v in [("scaffolding tau-cycle", 1/B), ("spiral step", 1.0),
                  ("branch cycle", PSI), ("spinor period", CYCLES*PSI)]:
        print(f"    {nm:<24}{v:>16,.4f} ticks")
    tpf = CYCLES*PSI/NO_FRAMES
    print(f"    scaffolding cycles per frame: {tpf*B:,.0f}  -> inset only")
    print(f"    point-state = 1 tick = {1/tpf:.2e} frames -> drawn as a flash")

    print(f"\n  key states")
    for c in range(CYCLES + 1):
        k = c*PSI
        x, y = psi_state(k)
        print(f"    after {c} cycle(s): phase = {tau_phase(k)/PI:.3f} pi, "
              f"Psi = ({x:+.4f}, {y:+.4f}), sign = {spinor_sign(k):+.0f}")
    print(f"\n  invariants (constant -- not animated)")
    print(f"    cone half-angle      {math.degrees(CONE):.4f} deg")
    print(f"    |m_j|                {math.sqrt(sum(c*c for c in monopole_vectors(0)[0])):.12f}")
    print(f"    |sum m_j|            {unit_sum(0):.12f}")


# ----------------------------------------------------------------------
# ANIMATION  --  commented out; uncomment the call at the bottom to run
# ----------------------------------------------------------------------

def build_animation(save_path=None, fps=30):
    """
    Five panels. Returns the FuncAnimation object.

    save_path=None  -> plt.show() (interactive, fine in Spyder)
    save_path="x.mp4" -> writes a file (needs ffmpeg)
    """
    import matplotlib
    import matplotlib.pyplot as plt
    from matplotlib.animation import FuncAnimation
    from mpl_toolkits.mplot3d import Axes3D  # noqa: F401
    import numpy as np

    # Style hardening. A dark stylesheet (Spyder's, or a global
    # plt.style) leaves titles and labels in a colour that vanishes
    # against the panel, which is why an earlier render showed no titles
    # at all. Nothing here is inherited: colours are set explicitly.
    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,
    })

    fig = plt.figure(figsize=(16, 9), facecolor=BG)
    axA = fig.add_subplot(2, 3, 1)                       # Psi complex plane
    axB = fig.add_subplot(2, 3, 2)                       # monopole phasors
    axC = fig.add_subplot(2, 3, 3, projection="3d")      # triplet in 3D
    axD = fig.add_subplot(2, 3, 4)                       # amplitude & phase
    axE = fig.add_subplot(2, 3, 5)                       # spinor sign
    axF = fig.add_subplot(2, 3, 6)                       # scaffolding inset

    kmax = CYCLES*PSI
    ks = np.linspace(0, kmax, NO_FRAMES)

    # ---- static backgrounds ----
    th = np.linspace(0, TAU, 300)
    full = [psi_state(k) for k in np.linspace(0, kmax, 1500)]
    axA.plot([p[0] for p in full], [p[1] for p in full],
             color="#dddddd", lw=1, zorder=1)
    axA.axhline(0, color="#999999", lw=0.6); axA.axvline(0, color="#999999", lw=0.6)
    axA.scatter([OMEGA], [0], s=70, color="green", zorder=4,
                label=f"$\\Psi(0)=\\Omega$")
    axA.scatter([-OMEGA**2], [0], s=140, marker="*", color="red", zorder=4,
                label=f"$\\Psi(\\psi)=-\\Omega^2$")
    # NOTE: every legend below sets loc= explicitly. matplotlib's default
    # is loc="best", which re-solves the placement on EVERY draw to avoid
    # overlapping the data. With blit=False each frame is a full redraw,
    # so the legend box hops between corners as the trace grows -- it
    # looks like flickering markers but the markers never move. Fixed
    # locations are chosen to sit in corners the data cannot reach
    # (|Psi| <= Omega^2 while the axis limit is 1.25 Omega^2, so the
    # corners are outside the trace radius).
    axA.set_aspect("equal"); axA.grid(alpha=0.2); axA.legend(fontsize=8, loc="lower left")
    axA.set_title("(a) $\\Psi = A\\,e^{i\\phi}$: half-cycle to $-\\Omega^2$",
                  fontsize=10)
    axA.set_xlabel("Re"); axA.set_ylabel("Im")

    axB.plot(np.cos(th), np.sin(th), color="#cccccc", lw=1)
    axB.axhline(0, color="#999999", lw=0.6); axB.axvline(0, color="#999999", lw=0.6)
    axB.scatter([1], [0], s=120, marker="*", color="darkgreen", zorder=4,
                label="holonomy target $+1$")
    axB.set_aspect("equal"); axB.grid(alpha=0.2); axB.legend(fontsize=8, loc="upper left")
    axB.set_title("(b) monopole phasors and holonomy product", fontsize=10)
    # NOT A BUG: with TURNS = (-2,0,0) only phi_1 advances; phi_2 and
    # phi_3 have n_j = 0 and sit permanently at phase 0, so two of the
    # three phasors coincide at (1,0) for the whole run. And since
    # phi_2 = phi_3 = 0, the holonomy sum EQUALS phi_1, so the crimson
    # holonomy marker sits on top of the phi_1 dot too. Three visible
    # markers, four artists. Choose a triple like (-1,-1,0) if you want
    # all three phasors to move.

    u_s = np.linspace(0, 2*np.pi, 30); v_s = np.linspace(0, np.pi, 15)
    axC.plot_wireframe(np.outer(np.cos(u_s), np.sin(v_s)),
                       np.outer(np.sin(u_s), np.sin(v_s)),
                       np.outer(np.ones_like(u_s), np.cos(v_s)),
                       color="#e5e5e5", linewidth=0.4)
    axC.set_box_aspect([1, 1, 1])
    axC.set_title(f"(c) triplet on its fixed cone\n"
                  f"half-angle {math.degrees(CONE):.2f}$^\\circ$ (constant)",
                  fontsize=10)

    cyc = ks/PSI
    axD.plot(cyc, [amplitude(k) for k in ks], color="#2980b9", lw=1.2,
             label="$A=\\Omega^{\\ln\\tau+1}$")
    axD.plot(cyc, [tau_phase(k)/PI for k in ks], color="#8e44ad", lw=1.2,
             label="$\\phi/\\pi$")
    axD.axhline(1, ls=":", color="gray"); axD.axhline(2, ls=":", color="gray")
    axD.grid(alpha=0.2); axD.legend(fontsize=8, loc="upper left")
    axD.set_title("(d) amplitude and phase", fontsize=10)
    axD.set_xlabel("branch cycles")

    axE.plot(cyc, [spinor_sign(k) for k in ks], color="#c0392b", lw=1.5)
    axE.axhline(1, ls=":", color="gray"); axE.axhline(-1, ls=":", color="gray")
    axE.scatter([1], [-1], s=110, color="red", zorder=4, label="$-\\xi$")
    axE.scatter([2], [1], s=130, marker="*", color="green", zorder=4,
                label="$+\\xi$")
    axE.set_ylim(-1.3, 1.3); axE.grid(alpha=0.2); axE.legend(fontsize=8, loc="lower right")
    axE.set_title("(e) spinor sign: two cycles to return", fontsize=10)
    axE.set_xlabel("branch cycles")

    # Panel (f) holds only dynamic content, so without explicit limits
    # matplotlib autoscales to an empty default and draws a blank box.
    lim = OMEGA**2*1.15
    axF.set_xlim(-lim, lim); axF.set_ylim(-lim*0.75, lim*0.75)
    axF.axhline(0, color="#999999", lw=0.6); axF.axvline(0, color="#999999", lw=0.6)
    axF.set_aspect("equal"); axF.grid(alpha=0.2)
    axF.set_title("(f) scaffolding $\\tau$-loop (INSET)\n"
                  "own time base -- ~300,000 cycles/frame at true rate",
                  fontsize=9)

    # ---- dynamic artists ----
    trace,  = axA.plot([], [], color="#d35400", lw=2.5, zorder=3)
    head,   = axA.plot([], [], "o", color="black", ms=8, zorder=5)
    ph_dots = [axB.plot([], [], "o", ms=9)[0] for _ in range(3)]
    ph_arrs = [axB.plot([], [], "-", lw=1.5)[0] for _ in range(3)]
    holo,   = axB.plot([], [], "o", color="crimson", ms=11, zorder=5)
    mono3d  = [axC.plot([], [], [], "-", lw=2.5)[0] for _ in range(3)]
    mark_d, = axD.plot([], [], "o", color="black", ms=7)
    mark_e, = axE.plot([], [], "o", color="black", ms=7)
    scaf,   = axF.plot([], [], color="#16a085", lw=2)
    scaf_h, = axF.plot([], [], "o", color="black", ms=7)
    status  = fig.text(0.015, 0.015, "", family="monospace", fontsize=9)
    cols = ["#c0392b", "#27ae60", "#2980b9"]
    for j in range(3):
        ph_dots[j].set_color(cols[j]); ph_arrs[j].set_color(cols[j])
        mono3d[j].set_color(cols[j])

    def init():
        for a in (trace, head, holo, mark_d, mark_e, scaf, scaf_h):
            a.set_data([], [])
        return []

    def update(f):
        k = ks[f]
        pts = [psi_state(kk) for kk in np.linspace(0, k, 400)]
        trace.set_data([p[0] for p in pts], [p[1] for p in pts])
        x, y = psi_state(k); head.set_data([x], [y])

        phs = monopole_phases(k)
        for j, p in enumerate(phs):
            ph_dots[j].set_data([math.cos(p)], [math.sin(p)])
            ph_arrs[j].set_data([0, math.cos(p)], [0, math.sin(p)])
        hx, hy = holonomy_product(k); holo.set_data([hx], [hy])

        for j, m in enumerate(monopole_vectors(k)):
            mono3d[j].set_data([0, m[0]], [0, m[1]])
            mono3d[j].set_3d_properties([0, m[2]])

        c = k/PSI
        mark_d.set_data([c], [amplitude(k)])
        mark_e.set_data([c], [spinor_sign(k)])

        # scaffolding inset: own time base, deliberately unsynchronised
        sp = (f*0.05) % 1.0
        spts = [scaffold_state(q) for q in np.linspace(0, sp, 120)]
        scaf.set_data([p[0] for p in spts], [p[1] for p in spts])
        sx, sy = scaffold_state(sp); scaf_h.set_data([sx], [sy])

        t_now = t_start + (k % PSI)
        status.set_text(
            f"cycle {c:6.3f} / {CYCLES}    k = {k:12,.0f} of {kmax:,.0f}\n"
            f"t_spiral = {t_now:,.0f}   radius = {math.sqrt(t_now):,.2f}\n"
            f"phase = {tau_phase(k)/PI:6.3f} pi   A = {amplitude(k):.4f}   "
            f"sign = {spinor_sign(k):+.3f}\n"
            f"|sum m| = {unit_sum(k):.6f}   (constant: symmetric triplet)")
        return []

    # Draw a real frame immediately. If the animation never steps -- as
    # with Spyder's inline backend, which captures one static image --
    # this is what gets captured, and it should be informative. Frame 0
    # is the worst choice: the trace has zero length and all three
    # phasors coincide at phase 0.
    update(int(PREVIEW_FRAC*NO_FRAMES))

    ani = FuncAnimation(fig, update, frames=NO_FRAMES, init_func=init,
                        blit=False, interval=1000/fps)
    fig.suptitle(f"Branch evolution   $\\sigma=\\sqrt{{2/\\alpha}}\\,\\Omega^5$, "
                 f"$\\psi$={PSI:,.0f},  {CYCLES} cycles", fontsize=13)
    fig.tight_layout(rect=[0, 0.06, 1, 0.96])

    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="branch_frame.png"):
    """Render ONE frame to PNG. Use this to check the panels without
    animating -- it needs no interactive backend and takes a second."""
    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%} of {CYCLES} cycles)")


if __name__ == "__main__":
    diagnostics()

    # ------------------------------------------------------------------
    # SPYDER BACKEND -- read this first.
    #
    # Spyder's default "inline" backend does NOT animate: it captures a
    # single static image, which is why an earlier run showed only the
    # static backgrounds. Either
    #
    #   Tools > Preferences > IPython console > Graphics > Backend: Qt5
    #   (then restart the kernel)
    #
    # or use the snapshot path below, which needs no interactive backend:
    #
    save_frame(0.30)          # one PNG, ~1 s, good for checking panels
    #
    # ANIMATION -- commented out. Uncomment ONE of these in Spyder.
    #
    #   interactive (fastest to check):
    import matplotlib.pyplot as plt
    # ani = build_animation()
    # plt.show()
    #
    #   write a file (needs ffmpeg):
    ani = build_animation(save_path="branch_evolution.mp4", fps=30)
    #
    # At NO_FRAMES = 2496 and ~400 curve points per frame this is roughly
    # 1e6 evaluations -- seconds of compute, but matplotlib rendering
    # dominates and an mp4 will take a few minutes.
    # ------------------------------------------------------------------
    print("\n  animation section is commented out; see __main__ to enable.")