#!/usr/bin/env python3
"""
Kapitza's pendulum — does a vibrating pivot really stabilize the upside-down state?

This verifier does NOT trust the textbook formula. It integrates the EXACT equation
of motion of a rigid pendulum whose pivot is driven vertically, and asks two things
the analytic averaging theory predicts:

  1. The inverted equilibrium (bob pointing straight UP) becomes stable when
         (a * w)^2  >  2 * g * L
     i.e. the dimensionless drive  K = (a/L) * (w/w0)  exceeds  sqrt(2) ≈ 1.41421,
     where w0 = sqrt(g/L) is the natural small-swing frequency about the down state.

  2. Once stable, the bob does a slow oscillation about the top whose frequency the
     time-averaged "effective potential" predicts exactly:
         Omega_slow = sqrt( a^2 w^2 / (2 L^2)  -  g / L ).

We find the numerical threshold by brute force (sweep the drive, watch whether the
pendulum released near the top falls over or stays), then compare to the formula,
and we measure the slow oscillation period from the full simulation and compare to
the formula. Every number printed here is what the page shows. No averaging is used
in the simulation itself — only RK4 on the true nonlinear ODE.

Equation of motion (theta measured from the DOWNWARD vertical; theta = pi is "up"):

    theta'' = -( g - a*w^2 * cos(w t) ) / L * sin(theta)

Pure NumPy, fixed-step RK4, fully reproducible:  python3 verify.py
"""

import json
import math
import os
import numpy as np

# ----------------------------------------------------------------------------
# Exact dynamics
# ----------------------------------------------------------------------------

def rk4_run(theta0, omega0, *, g, L, a, w, periods, steps_per_period):
    """Integrate the exact driven-pendulum ODE with fixed-step RK4.

    Returns arrays (t, theta, omega). theta is measured from the downward vertical;
    omega = theta-dot is the true integrated angular velocity.
    """
    T = 2.0 * math.pi / w
    dt = T / steps_per_period
    n = int(round(periods * steps_per_period))

    def accel(t, th):
        # vertical pivot acceleration s'' = -a w^2 cos(w t); effective gravity g + s''
        return -((g - a * w * w * math.cos(w * t)) / L) * math.sin(th)

    t = 0.0
    th = theta0
    om = omega0
    ts = np.empty(n + 1)
    ths = np.empty(n + 1)
    oms = np.empty(n + 1)
    ts[0] = 0.0
    ths[0] = th
    oms[0] = om
    for i in range(1, n + 1):
        # state y = (th, om); y' = (om, accel)
        k1th = om
        k1om = accel(t, th)
        k2th = om + 0.5 * dt * k1om
        k2om = accel(t + 0.5 * dt, th + 0.5 * dt * k1th)
        k3th = om + 0.5 * dt * k2om
        k3om = accel(t + 0.5 * dt, th + 0.5 * dt * k2th)
        k4th = om + dt * k3om
        k4om = accel(t + dt, th + dt * k3th)
        th += (dt / 6.0) * (k1th + 2 * k2th + 2 * k3th + k4th)
        om += (dt / 6.0) * (k1om + 2 * k2om + 2 * k3om + k4om)
        t += dt
        ts[i] = t
        ths[i] = th
        oms[i] = om
    return ts, ths, oms


def stays_up(*, g, L, a, w, kick=0.10, fall=math.pi / 2, slow_periods=12.0,
             steps_per_period=240):
    """Release the pendulum near the top with a small tilt; report whether it stays.

    Returns (stable: bool, max_excursion_rad). 'Stable' means it never tilts more
    than `fall` away from straight up over the integration window. The window is
    scaled to several predicted slow-oscillation periods (or a floor when the slow
    motion is undefined below threshold), so we always watch long enough to see it
    fall if it's going to.
    """
    # predicted slow angular frequency about the top (real only above threshold)
    arg = a * a * w * w / (2 * L * L) - g / L
    if arg > 0:
        Omega = math.sqrt(arg)
        T_slow = 2 * math.pi / Omega
    else:
        T_slow = None
    T_drive = 2 * math.pi / w
    if T_slow is not None:
        total_time = max(slow_periods * T_slow, 40 * T_drive)
    else:
        total_time = 60 * T_drive  # below threshold it falls fast; this is plenty
    periods = total_time / T_drive

    _, th, _ = rk4_run(math.pi + kick, 0.0, g=g, L=L, a=a, w=w,
                       periods=periods, steps_per_period=steps_per_period)
    excursion = np.abs(th - math.pi)
    max_exc = float(np.max(excursion))
    return (max_exc < fall), max_exc


def find_threshold(*, g, L, a, w0_lo, w0_hi, w0, tol_rel=1e-3):
    """Bisection on drive frequency w to find the stability threshold, at fixed a.

    w0_lo / w0_hi are bracket endpoints expressed as multiples of w0. Returns the
    critical w (rad/s).
    """
    lo = w0_lo * w0
    hi = w0_hi * w0
    assert not stays_up(g=g, L=L, a=a, w=lo)[0], "low bracket should be unstable"
    assert stays_up(g=g, L=L, a=a, w=hi)[0], "high bracket should be stable"
    while (hi - lo) / hi > tol_rel:
        mid = 0.5 * (lo + hi)
        if stays_up(g=g, L=L, a=a, w=mid)[0]:
            hi = mid
        else:
            lo = mid
    return 0.5 * (lo + hi)


def measure_slow_period(*, g, L, a, w, kick=0.30, steps_per_period=360):
    """Measure the slow oscillation period about the top from the full ODE.

    Released at theta = pi + kick (no velocity), the slow envelope swings back. We
    strip the fast micro-jitter by sampling once per drive period (stroboscopic),
    then find the time between successive maxima of (theta - pi).
    """
    T_drive = 2 * math.pi / w
    arg = a * a * w * w / (2 * L * L) - g / L
    Omega_pred = math.sqrt(arg)
    T_slow_pred = 2 * math.pi / Omega_pred
    periods = 6.0 * (T_slow_pred / T_drive)
    ts, th, _ = rk4_run(math.pi + kick, 0.0, g=g, L=L, a=a, w=w,
                        periods=periods, steps_per_period=steps_per_period)
    # stroboscopic sample: one point per drive period
    stride = steps_per_period
    ts_s = ts[::stride]
    dev = (th[::stride] - math.pi)
    # find interior local maxima
    peaks = []
    for i in range(1, len(dev) - 1):
        if dev[i] > dev[i - 1] and dev[i] >= dev[i + 1]:
            peaks.append(ts_s[i])
    if len(peaks) >= 2:
        gaps = np.diff(peaks)
        T_slow_meas = float(np.median(gaps))
    else:
        T_slow_meas = float('nan')
    return T_slow_pred, T_slow_meas


def energy_drift(*, g, L, a, w, periods=40, steps_per_period=240):
    """Sanity check the integrator: with NO drive (a=0) total energy is conserved.

    Returns the relative drift in mechanical energy over the run. A tiny number
    means RK4 at this resolution is faithful, so the threshold result isn't a
    numerical artifact.
    """
    th0 = math.pi / 3
    _, th, om = rk4_run(th0, 0.0, g=g, L=L, a=0.0, w=w,
                        periods=periods, steps_per_period=steps_per_period)
    # Energy per unit m L^2 with no drive:  E = 0.5*om^2 - (g/L)*cos(theta).
    # om here is the TRUE integrated angular velocity, so this measures the RK4
    # scheme's own energy error -- not a finite-difference artifact.
    E = 0.5 * om * om - (g / L) * np.cos(th)
    return float((np.max(E) - np.min(E)) / abs(np.mean(E)))


# ----------------------------------------------------------------------------
# Run the checks
# ----------------------------------------------------------------------------

def main():
    g = 9.81           # m/s^2
    L = 0.25           # m  (a 25 cm pendulum)
    w0 = math.sqrt(g / L)
    K_crit_theory = math.sqrt(2.0)          # dimensionless: (a/L)(w/w0) at threshold
    aw_crit_theory = math.sqrt(2.0 * g * L)  # m/s: peak pivot speed a*w at threshold

    print("=" * 72)
    print("Kapitza's pendulum — exact-ODE check of the inverted-state threshold")
    print("=" * 72)
    print(f"g = {g} m/s^2   L = {L} m   w0 = sqrt(g/L) = {w0:.5f} rad/s")
    print(f"Analytic threshold:  (a*w)^2 > 2 g L   <=>   a*w > {aw_crit_theory:.5f} m/s")
    print(f"Dimensionless:       K = (a/L)(w/w0) > sqrt(2) = {K_crit_theory:.5f}")
    print()

    # 0) integrator faithfulness
    drift = energy_drift(g=g, L=L, a=0.0, w=w0)
    print(f"[0] integrator energy drift over 40 swings (a=0): {drift:.2e} (relative)")
    print()

    # 1) numerical threshold at several amplitudes -> compare to formula
    print("[1] Numerical stability threshold of the UP state (released at pi+0.10 rad):")
    print(f"    {'a (m)':>8} {'a/L':>7} {'w_crit (rad/s)':>15} {'a*w_crit (m/s)':>15} "
          f"{'K_num':>9} {'K_num/sqrt2':>12}")
    rows = []
    for a in (0.010, 0.015, 0.020, 0.030, 0.040):
        # bracket: low multiple unstable, high multiple stable
        w_crit = find_threshold(g=g, L=L, a=a, w0_lo=2.0, w0_hi=60.0, w0=w0)
        aw = a * w_crit
        K = (a / L) * (w_crit / w0)
        rows.append({
            "a_m": a, "a_over_L": a / L, "w_crit_rad_s": w_crit,
            "aw_crit_m_s": aw, "K_num": K, "K_ratio": K / K_crit_theory,
            "freq_crit_hz": w_crit / (2 * math.pi),
        })
        print(f"    {a:8.3f} {a/L:7.3f} {w_crit:15.3f} {aw:15.4f} {K:9.4f} {K/K_crit_theory:12.4f}")
    K_nums = [r["K_num"] for r in rows]
    print(f"\n    mean numerical K at threshold = {np.mean(K_nums):.4f}  "
          f"(theory sqrt(2) = {K_crit_theory:.4f}); the slight excess is the known "
          f"O((w0/w)^2) correction to leading-order averaging.")
    print()

    # 2) slow-oscillation frequency: full ODE vs effective-potential formula
    print("[2] Slow oscillation about the top — full ODE vs effective potential:")
    print(f"    {'a (m)':>8} {'w (rad/s)':>10} {'K':>6} {'T_slow pred (s)':>16} "
          f"{'T_slow meas (s)':>16} {'ratio':>8}")
    slow_rows = []
    # pick comfortably-above-threshold settings
    for a, w in ((0.020, 180.0), (0.030, 140.0), (0.040, 120.0)):
        K = (a / L) * (w / w0)
        Tp, Tm = measure_slow_period(g=g, L=L, a=a, w=w)
        slow_rows.append({"a_m": a, "w_rad_s": w, "K": K,
                          "T_slow_pred_s": Tp, "T_slow_meas_s": Tm,
                          "ratio": (Tm / Tp) if Tp else float('nan')})
        print(f"    {a:8.3f} {w:10.1f} {K:6.2f} {Tp:16.4f} {Tm:16.4f} {Tm/Tp:8.4f}")
    print()

    # 3) the headline demo numbers used on the page
    a_demo, L_demo = 0.020, 0.25
    w_thr = math.sqrt(2 * g * L_demo) / a_demo   # w at threshold for a=2cm
    print("[3] Page demo (a = 2 cm, L = 25 cm):")
    print(f"    threshold drive frequency = {w_thr:.2f} rad/s = {w_thr/(2*math.pi):.2f} Hz")
    print(f"    (below this the bob falls; above it, it stands up and springs back)")

    # the side equilibria that appear above threshold (the new barriers):
    # cos(theta*) = -2 g L / (a^2 w^2); shown for a strongly-driven case
    a_s, w_s = 0.030, 140.0
    c = -2 * g * L / (a_s * a_s * w_s * w_s)
    theta_star = math.degrees(math.acos(c)) if -1 <= c <= 1 else None
    print(f"    above threshold (a=3cm, w=140): the unstable barrier sits at "
          f"theta* = {theta_star:.1f} deg from straight DOWN "
          f"(cos theta* = {c:.4f}).")

    out = {
        "params": {"g": g, "L": L, "w0": w0},
        "theory": {
            "threshold_aw_m_s": aw_crit_theory,
            "threshold_K_dimensionless": K_crit_theory,
            "effective_potential": "U/(mL^2) = -(g/L)cos(theta) + (a^2 w^2/(4 L^2)) sin^2(theta)",
        },
        "integrator_energy_drift_rel": drift,
        "threshold_sweep": rows,
        "threshold_K_mean": float(np.mean(K_nums)),
        "slow_oscillation": slow_rows,
        "page_demo": {
            "a_m": a_demo, "L_m": L_demo,
            "threshold_w_rad_s": w_thr,
            "threshold_freq_hz": w_thr / (2 * math.pi),
            "barrier_theta_deg_from_down_example": theta_star,
            "barrier_example_a_m": a_s, "barrier_example_w_rad_s": w_s,
        },
    }
    here = os.path.dirname(os.path.abspath(__file__))
    with open(os.path.join(here, "numbers.json"), "w") as f:
        json.dump(out, f, indent=2)
    print(f"\nwrote {os.path.join(here, 'numbers.json')}")


if __name__ == "__main__":
    main()
