#!/usr/bin/env python3
"""Re-run the externally driven Rikitake dynamo of Kharshiladze, Tsulukidze, Kezua and Zilpimiani (2025),
"Control of Inversions of the Earth's Magnetic Field with Rikitake's Model", Journals of Georgian Geophysical
Society 28(2), 94-101 (https://ggs.openjournals.ge/index.php/GGS/article/view/10461), and test what its match
with the reversal record rests on. Fact-Checked Science, https://factchecked.science/rikitake-electromechanical/

    python3 forced_rikitake.py                 prints every number from our re-run on the page, in a few seconds
    python3 forced_rikitake.py --json out.json also writes all runs, and the series behind the two charts

Only the Python 3 standard library is needed, and no download: the reversal record is written out below.

The model (the paper's eq. 6, in its dimensionless form). Two disks, with currents x1, x2 and spin rates
y1, y2; x0 is the "external current", added to the current that makes the field threading each disk:

    dx1/dt = -mu*x1 + y1*(x2 + x0)        dy1/dt = 1 - x1*(x2 + x0)
    dx2/dt = -mu*x2 + y2*(x1 + x0)        dy2/dt = 1 - (x1 + x0)*x2

with mu = 1 and starting values x1 = 1, x2 = -1, y1 = 0, y2 = 1, as in the captions of its Figs. 1-2. The
PDF's mathematical symbols can't be extracted as text; we read them from the page layout and checked the
reading twice: with x0 = 1 held constant the model settles at x1 = x2 = (-1 + sqrt 5)/2 = 0.618, the plateau of
the paper's Fig. 3, and a square wave whose amplitude, 0.21, is set from the plateau of Fig. 2(b) flips about
12 units after each switch, as that figure does. One time unit is 10,000 years (the paper, p. 96), so the last
3.5 million years are 350 units.

The record: the polarity boundaries of Cande and Kent (1995), https://doi.org/10.1029/94JB03098, the record
the paper used, as tabulated in PmagPy's get_ts('ck95') (Tauxe et al. 2016, https://doi.org/10.1002/2016GC006307).

The runs
  A  the real record, outside current +-1 with weak noise, as in the paper's Fig. 3
  B  20 made-up records, each with 12 reversals at random times
  C  the real record with weaker outside currents, +-0.5 down to +-0.01
  D  no outside current at all
  E  Bullard's single disk with an outside current added the same way (an illustration: the 2026 abstract
     gives no equations for its Bullard model)
"Followed" means the model's current changes sign within 50,000 years after the outside current does.
Random numbers are seeded, so every run gives the same numbers.
"""
import json
import math
import random
import sys

# Cande & Kent (1995) polarity boundaries in millions of years ago, 0-3.58 Ma (PmagPy get_ts('ck95'))
CK95 = [0.0, 0.780, 0.990, 1.070, 1.770, 1.950, 2.140, 2.150, 2.581, 3.040, 3.110, 3.220, 3.330, 3.580]
SPAN_MA = 3.5
UNITS_PER_MA = 100.0  # one time unit is 10,000 years
T_END = SPAN_MA * UNITS_PER_MA
MU = 1.0
DT = 0.005


def polarity_from_boundaries(bounds_ma):
    """Return a function age (Ma) -> +1 (normal polarity) or -1 (reversed); the youngest interval is normal."""
    def pol(age):
        k = 0
        for b in bounds_ma[1:]:
            if age >= b:
                k += 1
            else:
                break
        return 1 if k % 2 == 0 else -1
    return pol


def flips_in_window(bounds_ma):
    """Model times (units from the start, 3.5 Ma) of the polarity changes inside 0-3.5 Ma."""
    return sorted((SPAN_MA - b) * UNITS_PER_MA for b in bounds_ma[1:] if 0 < b < SPAN_MA)


def rhs(s, x0):
    """The paper's eq. 6."""
    x1, x2, y1, y2 = s
    return (-MU * x1 + y1 * (x2 + x0),
            -MU * x2 + y2 * (x1 + x0),
            1.0 - x1 * (x2 + x0),
            1.0 - (x1 + x0) * x2)


def rhs_bullard(s, x0):
    """Bullard's single disk without friction, with the outside current added as in eq. 6:
    dx/dt = -mu*x + y*(x + x0), dy/dt = 1 - x*(x + x0). With x0 = 0 the current can never change sign."""
    x, y = s
    return (-MU * x + y * (x + x0), 1.0 - x * (x + x0))


def run(forcing, noise_sd=0.0, seed=1, t_end=T_END, s0=(1.0, -1.0, 0.0, 1.0), f=rhs):
    """Fourth-order Runge-Kutta. forcing(t) gives x0; noise is added to x0, held for 0.1 units at a time.
    Returns the first current, x1 (or x), every 0.1 units (1,000 years)."""
    rng = random.Random(seed)
    s = list(s0)
    m = len(s)
    t = 0.0
    n = int(round(t_end / DT))
    hold = int(round(0.1 / DT))
    eps = 0.0
    xs = []
    for i in range(n):
        if noise_sd and i % hold == 0:
            eps = rng.gauss(0.0, noise_sd)
        x0 = forcing(t) + eps
        k1 = f(s, x0)
        k2 = f([s[j] + 0.5 * DT * k1[j] for j in range(m)], x0)
        k3 = f([s[j] + 0.5 * DT * k2[j] for j in range(m)], x0)
        k4 = f([s[j] + DT * k3[j] for j in range(m)], x0)
        s = [s[j] + DT / 6.0 * (k1[j] + 2 * k2[j] + 2 * k3[j] + k4[j]) for j in range(m)]
        t += DT
        if not all(math.isfinite(v) for v in s) or abs(s[0]) > 1e6:
            return xs, "diverged at t=%.2f" % t
        if i % 20 == 0:
            xs.append((round(t, 4), s[0]))
    return xs, "ok"


def polarity_changes(xs, thr=0.5):
    """Lasting polarity changes, with hysteresis: counted when the current passes -thr after having been above
    +thr, or the reverse, so brief dips across zero inside one polarity don't count."""
    state = None
    times = []
    for t, x in xs:
        if x > thr:
            if state == -1:
                times.append(t)
            state = 1
        elif x < -thr:
            if state == 1:
                times.append(t)
            state = -1
    return times


def zero_crossings(xs):
    out = []
    for (ta, a), (tb, b) in zip(xs, xs[1:]):
        if a == 0:
            continue
        if (a > 0) != (b > 0):
            out.append(ta + (tb - ta) * (abs(a) / (abs(a) + abs(b))))
    return out


def sign_agreement(xs, forcing):
    """Share of the samples in which the model's current has the same sign as the forcing."""
    agree = sum(1 for t, x in xs if (x > 0) == (forcing(t) > 0))
    return agree / len(xs)


def match_flips(flips, crossings, window=5.0):
    """For each flip of the outside current, the first sign change of the model within 5 units (50,000 years)."""
    lags = []
    for f in flips:
        after = [c - f for c in crossings if 0 <= c - f <= window]
        lags.append(min(after) if after else None)
    return lags


def lags_until_next(flips, crossings, t_end=T_END):
    """For each flip of the outside current, the delay to the model's first sign change before the next flip
    (None if it doesn't change sign in that interval). Units of the model (one unit = 10,000 years)."""
    out = []
    for i, f in enumerate(flips):
        nxt = flips[i + 1] if i + 1 < len(flips) else t_end
        after = [c - f for c in crossings if f <= c < nxt]
        out.append(min(after) if after else None)
    return out


def forcing_from_bounds(bounds_ma, amp):
    pol = polarity_from_boundaries(bounds_ma)
    return lambda t: amp * pol(SPAN_MA - t / UNITS_PER_MA)


def random_bounds(n_flips, rng, min_gap_ma=0.01):
    """A made-up record: n_flips boundaries at random ages within 0-3.5 Ma (youngest interval normal)."""
    while True:
        b = sorted(rng.uniform(0.0, SPAN_MA) for _ in range(n_flips))
        if all(b2 - b1 >= min_gap_ma for b1, b2 in zip([0.0] + b, b + [SPAN_MA])):
            return [0.0] + b + [SPAN_MA + 1.0]


def summarise(label, xs, status, forcing, flips):
    cr = zero_crossings(xs)
    lags = match_flips(flips, cr)
    matched = [l for l in lags if l is not None]
    pc = polarity_changes(xs, thr=0.3)
    return {
        "run": label,
        "status": status,
        "n_forcing_flips": len(flips),
        "n_x1_zero_crossings": len(cr),
        "n_x1_lasting_polarity_changes_hyst0.3": len(pc),
        "flips_followed_within_50kyr": len(matched),
        "median_lag_kyr": round(sorted(matched)[len(matched) // 2] * 10.0, 1) if matched else None,
        "max_lag_kyr": round(max(matched) * 10.0, 1) if matched else None,
        "share_of_time_same_sign_as_forcing": round(sign_agreement(xs, forcing), 3),
    }


def steady_x1(amp):
    """The model's steady current under a constant outside current amp: x1 = x2 = (-amp + sqrt(amp^2 + 4))/2."""
    return (-amp + (amp * amp + 4) ** 0.5) / 2


def every(xs, k):
    return [round(x, 3) for _, x in xs[::k]]


def compute():
    results = {"model": "Kharshiladze et al. (2025), eq. 6, as reconstructed", "mu": MU, "dt": DT,
               "time_unit_years": 1e4, "span_Ma": SPAN_MA, "ck95_boundaries_Ma": CK95, "runs": []}
    real_flips = flips_in_window(CK95)
    series = {}

    # Checks of the reconstruction
    xs, st = run(lambda t: 1.0, t_end=200.0)
    results["check_constant_x0_1"] = {"x1_at_end": round(xs[-1][1], 4), "expected": round(steady_x1(1.0), 4), "status": st}
    amp2b = (1 - 0.9 ** 2) / 0.9   # the amplitude that gives Fig. 2(b)'s plateau of 0.9
    sq = lambda t: amp2b * (1 if int(t // 100) % 2 == 0 else -1)
    xs, st = run(sq, t_end=800.0)
    switches = [100.0 * k for k in range(1, 8)]
    results["check_fig2b_square_wave"] = {
        "amplitude": round(amp2b, 3), "half_period_units": 100, "status": st,
        "x1_lasting_polarity_changes_at_units": [round(t, 1) for t in polarity_changes(xs, thr=0.3)],
        "zero_crossing_delays_after_switch_units": [None if l is None else round(l, 1)
                                                    for l in lags_until_next(switches, zero_crossings(xs), 800.0)],
        "plateau_x1_at_t_95": round([x for t, x in xs if t >= 95][0], 3)}

    # A. the real record, outside current +-1 with weak noise (the paper's Fig. 3)
    f_real = forcing_from_bounds(CK95, 1.0)
    xs, st = run(f_real, noise_sd=0.05, seed=11)
    resA = summarise("A: real record, outside current 1", xs, st, f_real, real_flips)
    resA["steady_x1"] = round(steady_x1(1.0), 3)
    resA["dynamo_to_outside_current_ratio"] = resA["steady_x1"]
    resA["x1_crossing_times_Ma_ago"] = [round(SPAN_MA - c / UNITS_PER_MA, 3) for c in zero_crossings(xs)]
    resA["ck95_flip_times_Ma_ago"] = [round(SPAN_MA - f / UNITS_PER_MA, 3) for f in real_flips]
    results["runs"].append(resA)
    series["t_Ma_ago"] = [round(SPAN_MA - t / UNITS_PER_MA, 3) for t, _ in xs[::5]]
    series["forcing_real"] = [f_real(t) for t, _ in xs[::5]]
    series["x1_real_amp1"] = every(xs, 5)

    # B. 20 made-up records with the same number of reversals
    rng = random.Random(2026)
    fake_rows = []
    for k in range(20):
        fb = random_bounds(len(real_flips), rng)
        f_fake = forcing_from_bounds(fb, 1.0)
        xs, st = run(f_fake, noise_sd=0.05, seed=100 + k)
        row = summarise("B%02d: made-up record, outside current 1" % k, xs, st, f_fake, flips_in_window(fb))
        row["share_of_time_same_sign_as_REAL_record"] = round(sign_agreement(xs, f_real), 3)
        row["made_up_flip_times_Ma_ago"] = [round(b, 3) for b in fb[1:-1]]
        fake_rows.append(row)
        if k == 0:
            series["forcing_made_up_B00"] = [f_fake(t) for t, _ in xs[::5]]
            series["x1_made_up_B00"] = every(xs, 5)
    results["runs"].extend(fake_rows)
    results["B_summary"] = {
        "runs": len(fake_rows),
        "flips_followed_total": sum(r["flips_followed_within_50kyr"] for r in fake_rows),
        "flips_total": sum(r["n_forcing_flips"] for r in fake_rows),
        "max_lag_kyr": max(r["max_lag_kyr"] for r in fake_rows),
        "mean_share_same_sign_as_own_forcing": round(sum(r["share_of_time_same_sign_as_forcing"] for r in fake_rows) / len(fake_rows), 3),
        "min_share_same_sign_as_own_forcing": min(r["share_of_time_same_sign_as_forcing"] for r in fake_rows),
        "mean_share_same_sign_as_real_record": round(sum(r["share_of_time_same_sign_as_REAL_record"] for r in fake_rows) / len(fake_rows), 3),
        "range_share_same_sign_as_real_record": [min(r["share_of_time_same_sign_as_REAL_record"] for r in fake_rows),
                                                 max(r["share_of_time_same_sign_as_REAL_record"] for r in fake_rows)],
    }

    # C. the real record with weaker outside currents
    for amp in (0.5, 0.3, 0.2, 0.1, 0.05, 0.03, 0.01):
        f = forcing_from_bounds(CK95, amp)
        xs, st = run(f, noise_sd=0.05 * amp, seed=21)
        row = summarise("C: real record, outside current %g" % amp, xs, st, f, real_flips)
        row["steady_x1"] = round(steady_x1(amp), 3)
        row["dynamo_to_outside_current_ratio"] = round(row["steady_x1"] / amp, 2)
        row["lags_until_next_flip_kyr"] = [None if l is None else round(l * 10.0, 1)
                                           for l in lags_until_next(real_flips, zero_crossings(xs))]
        results["runs"].append(row)
        if amp == 0.1:
            series["x1_real_amp0.1"] = every(xs, 5)
    results["series_every_5kyr"] = series

    # D. no outside current
    xs, st = run(lambda t: 0.0)
    win = 50  # 50 samples of 1,000 years
    means = [sum(x for _, x in xs[i:i + win]) / win for i in range(0, len(xs) - win + 1, win)]
    results["runs"].append({"run": "D: no outside current", "status": st,
                            "n_x1_zero_crossings": len(zero_crossings(xs)),
                            "n_x1_lasting_polarity_changes_hyst0.3": len(polarity_changes(xs, thr=0.3)),
                            "n_x1_lasting_polarity_changes_hyst0.5": len(polarity_changes(xs, thr=0.5)),
                            "n_sign_changes_of_50kyr_means": sum(1 for a, b in zip(means, means[1:]) if (a > 0) != (b > 0)),
                            "share_of_time_positive": round(sum(1 for t, x in xs if x > 0) / len(xs), 3)})

    # E. Bullard's disk with an outside current (illustration)
    xs, st = run(lambda t: 0.0, s0=(1.0, 0.0), f=rhs_bullard)
    results["runs"].append({"run": "E: Bullard, no outside current", "status": st,
                            "n_x1_zero_crossings": len(zero_crossings(xs)),
                            "min_x": round(min(x for _, x in xs), 3), "max_x": round(max(x for _, x in xs), 3)})
    for amp in (1.0, 0.5, 0.3, 0.1, 0.03):
        fr = forcing_from_bounds(CK95, amp)
        xs, st = run(fr, noise_sd=0.05 * amp, seed=31, s0=(1.0, 0.0), f=rhs_bullard)
        results["runs"].append(summarise("E: Bullard, real record, outside current %g" % amp, xs, st, fr, real_flips))
    fb = random_bounds(len(real_flips), random.Random(77))
    ff = forcing_from_bounds(fb, 1.0)
    xs, st = run(ff, noise_sd=0.05, seed=32, s0=(1.0, 0.0), f=rhs_bullard)
    row = summarise("E: Bullard, made-up record, outside current 1", xs, st, ff, flips_in_window(fb))
    row["share_of_time_same_sign_as_REAL_record"] = round(sign_agreement(xs, f_real), 3)
    results["runs"].append(row)
    return results


def pct(v):
    return "%.1f%%" % (100 * v)


def report(r):
    out = []
    c1, c2 = r["check_constant_x0_1"], r["check_fig2b_square_wave"]
    out.append("Checks of the reconstruction")
    out.append("  constant outside current 1: the model settles at %.3f (expected %.3f, the plateau of Fig. 3)" % (c1["x1_at_end"], c1["expected"]))
    out.append("  square wave with the amplitude set from Fig. 2b's plateau (%.2f), half-period 100: plateau %.2f (Fig. 2b: about 0.9)"
               % (c2["amplitude"], c2["plateau_x1_at_t_95"]))
    out.append("    the model's current crosses zero %s units after each switch (Fig. 2b, by eye: about 11-14)"
               % ", ".join("-" if l is None else "%g" % l for l in c2["zero_crossing_delays_after_switch_units"]))
    out.append("    lasting flips (past +-0.3) at %s" % ", ".join("%g" % t for t in c2["x1_lasting_polarity_changes_at_units"]))
    out.append("")
    out.append("Rikitake's model, real or made-up record (followed = sign change within 50,000 years)")
    out.append("  %-38s %8s %8s %9s %13s %10s %7s" % ("run", "ratio*", "inverse", "followed", "slowest (yr)", "same sign", "sign changes"))
    for row in r["runs"]:
        if row["run"].startswith(("A", "C")):
            out.append("  %-38s %8s %8s %9s %13s %10s %7d" % (
                row["run"], "%.2f" % row["dynamo_to_outside_current_ratio"], "%.2f" % (1 / row["dynamo_to_outside_current_ratio"]),
                "%d of %d" % (row["flips_followed_within_50kyr"], row["n_forcing_flips"]),
                "-" if row["max_lag_kyr"] is None else "%d" % round(row["max_lag_kyr"] * 1000),
                pct(row["share_of_time_same_sign_as_forcing"]), row["n_x1_zero_crossings"]))
    out.append("  * the dynamo's own steady current divided by the outside current (the paper estimates 1 to 100);"
               " inverse: the outside current as a share of the dynamo's own")
    for row in r["runs"]:
        if row["run"] == "C: real record, outside current 0.2":
            later = [l for l in row["lags_until_next_flip_kyr"] if l is not None and l > 50]
            out.append("  at +-0.2 the model follows %d more reversals later, before the next one comes, after %s years"
                       % (len(later), ", ".join("{:,}".format(int(round(l * 1000))) for l in sorted(later))))
    b = r["B_summary"]
    slowest = "{:,}".format(int(round(b["max_lag_kyr"] * 1000)))
    out.append("  B, 20 made-up records: followed %d of %d made-up reversals, the slowest after %s years; same sign as"
               " their own record %s on average (lowest %s), as the real record %s (range %s-%s)"
               % (b["flips_followed_total"], b["flips_total"], slowest,
                  pct(b["mean_share_same_sign_as_own_forcing"]), pct(b["min_share_same_sign_as_own_forcing"]),
                  pct(b["mean_share_same_sign_as_real_record"]), pct(b["range_share_same_sign_as_real_record"][0]),
                  pct(b["range_share_same_sign_as_real_record"][1])))
    b00 = [row for row in r["runs"] if row["run"].startswith("B00")][0]
    out.append("  B00, the made-up run charted: followed %d of %d; same sign as its own record %s, as the real record %s"
               % (b00["flips_followed_within_50kyr"], b00["n_forcing_flips"], pct(b00["share_of_time_same_sign_as_forcing"]),
                  pct(b00["share_of_time_same_sign_as_REAL_record"])))
    d = [row for row in r["runs"] if row["run"].startswith("D")][0]
    out.append("  D, no outside current: %d lasting polarity changes (past +-0.5) in 3.5 model-Myr (%d past +-0.3, %d sign"
               " changes of 50,000-year means, %d zero crossings); the real record has 12"
               % (d["n_x1_lasting_polarity_changes_hyst0.5"], d["n_x1_lasting_polarity_changes_hyst0.3"],
                  d["n_sign_changes_of_50kyr_means"], d["n_x1_zero_crossings"]))
    out.append("")
    out.append("Bullard's disk with an outside current (illustration)")
    for row in r["runs"]:
        if row["run"].startswith("E"):
            if "min_x" in row:
                out.append("  %-44s sign changes %d; current stays between %.2f and %.2f" % (row["run"], row["n_x1_zero_crossings"], row["min_x"], row["max_x"]))
            else:
                extra = ""
                if "share_of_time_same_sign_as_REAL_record" in row:
                    extra = "; same sign as the real record %s" % pct(row["share_of_time_same_sign_as_REAL_record"])
                slow = ("" if row["max_lag_kyr"] is None else
                        ", the slowest after {:,} years".format(int(round(row["max_lag_kyr"] * 1000))))
                out.append("  %-44s followed %d of %d%s; same sign as its record %s%s" % (
                    row["run"], row["flips_followed_within_50kyr"], row["n_forcing_flips"], slow,
                    pct(row["share_of_time_same_sign_as_forcing"]), extra))
    return "\n".join(out)


def main():
    args = sys.argv[1:]
    out_json = None
    if args[:1] == ["--json"] and len(args) > 1:
        out_json = args[1]
    elif args:
        sys.exit(__doc__)
    r = compute()
    print(report(r))
    if out_json:
        with open(out_json, "w", encoding="utf-8") as fh:
            json.dump(r, fh, indent=1)
        print("\nwrote", out_json)


if __name__ == "__main__":
    main()