"""
Test 14 — Kuiper belt as a mixed-mode test of Conjecture I‴.

The asteroid belt (Test 13) was a clean KAM-stability filter: gaps at
low-order rationals with Jupiter. The Kuiper belt is mixed:

  - Plutinos:  resonance-trapped at 3:2 with Neptune (P_TNO / P_N = 1.5)
  - twotinos:  resonance-trapped at 2:1 (P_TNO / P_N = 2.0)
  - Cold Classical KBOs (a ≈ 42–47 AU): smooth, non-resonant
  - Hot Classical, scattered disk: mix

So the Kuiper belt should show *both* pile-ups (at 3:2, 2:1, 5:3, …) AND
the smooth Classical population. Conjecture I‴ predicts that the *sign*
of the {r}-bin deviation at low-order rationals identifies which sub-
population is sampling: positive (pile-up) for trapped, neutral/zero
for smooth, negative for KAM-filtered.

Predictions:
  - {r} = 0.5 (3:2 Plutinos):     z > 0 (excess), like exoplanets
  - {r} = 0.0 (2:1 twotinos):     z > 0 (excess)
  - Classical KBO bulk:            consistent with smooth, not Gauss-Kuzmin

Data: 5186 TNOs with 30 ≤ a ≤ 60 AU from JPL SBDB.
"""

import json
import math
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt


LN2 = math.log(2.0)
KHINCHIN_LN = 0.9878490568339858
KHINCHIN_K  = math.exp(KHINCHIN_LN)
LEVY_LN     = math.pi ** 2 / (12 * LN2)
P_NEPTUNE_DAYS = 164.79 * 365.25   # = 60189 days


def load(path):
    with open(path) as f:
        d = json.load(f)
    rows = []
    for full_name, a, e, i, per, H in d["data"]:
        try:
            rows.append((full_name.strip(),
                          float(a), float(e), float(i),
                          float(per), float(H) if H else float("nan")))
        except (TypeError, ValueError):
            continue
    return rows


def continued_fraction(x, depth=10, eps=1e-12):
    qs = []
    p_prev2, p_prev1 = 1, 0
    q_prev2, q_prev1 = 0, 1
    convergents = []
    for _ in range(depth):
        a = math.floor(x)
        qs.append(int(a))
        p = a * p_prev1 + p_prev2
        q = a * q_prev1 + q_prev2
        convergents.append((p, q))
        p_prev2, p_prev1 = p_prev1, p
        q_prev2, q_prev1 = q_prev1, q
        frac = x - a
        if frac < eps:
            break
        x = 1.0 / frac
    return qs, convergents


def per_ratio_stats(r, depth=10):
    aks, conv = continued_fraction(r, depth=depth)
    aks_post = aks[1:]
    if not aks_post:
        return [], None, None
    ln_a_mean = float(np.mean([math.log(a) for a in aks_post if a > 0]))
    q_list = [q for _, q in conv if q > 0]
    levy = None
    if len(q_list) >= 4:
        ns = np.arange(1, len(q_list) + 1, dtype=float)
        logs = np.log([q for q in q_list])
        levy = float(np.polyfit(ns[1:], logs[1:], 1)[0])
    return aks_post, ln_a_mean, levy


def gk_pmf(max_m=10):
    pmf = np.zeros(max_m + 1)
    for m in range(1, max_m):
        pmf[m] = math.log2(1.0 + 1.0 / (m * (m + 2)))
    pmf[max_m] = 1.0 - pmf[1:max_m].sum()
    return pmf


def main():
    tnos = load("tnos.json")
    print(f"loaded {len(tnos)} TNOs with valid records")

    a_vals = np.array([t[1] for t in tnos])
    per_d = np.array([t[4] for t in tnos])

    # r = P_TNO / P_Neptune > 1 (TNOs are slower than Neptune)
    r = per_d / P_NEPTUNE_DAYS
    sel = (r > 1.0) & np.isfinite(r) & (r < 5.0)
    r = r[sel]; a_vals = a_vals[sel]
    print(f"after r > 1 cut: {len(r)} TNOs in r ∈ ({r.min():.3f}, {r.max():.3f})")
    print(f"a range: [{a_vals.min():.2f}, {a_vals.max():.2f}] AU")

    # ----- canonical resonance locations -----
    # P_TNO/P_N → resonance label (TNO does q orbits per p Neptune; p > q for r > 1)
    KNOWN_RES = {
        "4:3":   4.0/3.0,
        "5:4":   5.0/4.0,
        "3:2":   3.0/2.0,   # Plutinos
        "5:3":   5.0/3.0,
        "7:4":   7.0/4.0,
        "2:1":   2.0/1.0,   # twotinos
        "5:2":   5.0/2.0,
        "7:3":   7.0/3.0,
        "8:3":   8.0/3.0,
        "3:1":   3.0/1.0,
    }

    # ----- bin {r} -----
    fracs = r - np.floor(r)
    bins = 40
    counts, edges = np.histogram(fracs, bins=bins, range=(0, 1))
    n_total = counts.sum()

    # Compute Gauss-measure expected counts in each bin
    gauss_exp = np.array([
        n_total * math.log2((1 + edges[i+1]) / (1 + edges[i]))
        for i in range(bins)
    ])
    chi2_g = ((counts - gauss_exp) ** 2 / gauss_exp).sum()
    dof = bins - 1
    print(f"\nχ² vs Gauss on {bins}-bin {{r}}: {chi2_g:.1f} on {dof} dof")

    dev = (counts - gauss_exp) / np.sqrt(gauss_exp)
    print(f"\ntop bin excesses (pile-ups → trapping signature):")
    top = np.argsort(-dev)[:8]
    for k in top:
        lo, hi = edges[k], edges[k+1]
        # Identify which resonance(s) this bin contains
        in_bin = [name for name, ratio in KNOWN_RES.items()
                  if lo <= (ratio - math.floor(ratio)) < hi]
        in_bin_str = ", ".join(in_bin) if in_bin else "—"
        print(f"  [{lo:.3f}, {hi:.3f}): obs={counts[k]:5d}  exp={gauss_exp[k]:5.0f}  "
              f"dev={dev[k]:+.2f}  contains: {in_bin_str}")
    print(f"\ntop bin deficits (gaps → filtering signature):")
    bot = np.argsort(dev)[:8]
    for k in bot:
        lo, hi = edges[k], edges[k+1]
        in_bin = [name for name, ratio in KNOWN_RES.items()
                  if lo <= (ratio - math.floor(ratio)) < hi]
        in_bin_str = ", ".join(in_bin) if in_bin else "—"
        print(f"  [{lo:.3f}, {hi:.3f}): obs={counts[k]:5d}  exp={gauss_exp[k]:5.0f}  "
              f"dev={dev[k]:+.2f}  contains: {in_bin_str}")

    # ----- narrow-window pile-up test at known resonances -----
    print(f"\nnarrow-window pile-up test at each known TNO resonance (±2% in r):")
    print(f"  {'res':>5}  {'r_res':>7}  {'observed':>9}  {'expected (Gauss in window)':>27}  z")
    WIN = 0.02
    for name, r_res in KNOWN_RES.items():
        # Count TNOs with r in (r_res - WIN, r_res + WIN)
        in_win = np.sum(np.abs(r - r_res) < WIN)
        # Gauss-measure expected within window:
        # convert window to fractional-part window. r_res in (1, 3) so just take fractional part window.
        # Actually for absolute r, we want the same WIN width.
        # Gauss expected over a window of width 2*WIN centered at r_res in fractional space:
        frac_res = r_res - math.floor(r_res)
        lo_f, hi_f = max(0, frac_res - WIN), min(1, frac_res + WIN)
        gauss_frac = math.log2((1 + hi_f) / (1 + lo_f))
        # Account for the fact that bins outside this floor might have density.
        # Approximation: just compare to gauss expectation in this {r} window scaled by total
        expected = gauss_frac * n_total
        # But the "in_win" counts ALL stars in this absolute r window (could come from different floor(r)).
        # For r_res > 1 and not crossing integer, this is one fractional bin.
        # Better: count empirical density per unit r, compare to gauss density at r_res
        z = (in_win - expected) / math.sqrt(expected) if expected > 0 else float("nan")
        print(f"  {name:>5}  {r_res:7.4f}  {in_win:9d}  {expected:27.1f}  {z:+.2f}")

    # ----- partial-quotient distribution -----
    all_aks = []
    ratio_lnK = []
    ratio_levy = []
    for ri in r:
        aks, lnK, L = per_ratio_stats(ri, depth=10)
        all_aks.extend(aks)
        if lnK is not None:
            ratio_lnK.append(lnK)
        if L is not None:
            ratio_levy.append(L)
    ratio_lnK = np.array(ratio_lnK)
    ratio_levy = np.array(ratio_levy)

    max_m = 10
    pq_counts = np.zeros(max_m + 1, dtype=int)
    for v in all_aks:
        if v >= 1:
            pq_counts[min(v, max_m)] += 1
    total_pq = pq_counts.sum()
    freqs = pq_counts / total_pq
    gk = gk_pmf(max_m=max_m)
    expected_pq = gk * total_pq
    chi2_pq = (((pq_counts[1:] - expected_pq[1:]) ** 2) / expected_pq[1:]).sum()
    print(f"\npartial-quotient χ² vs Gauss-Kuzmin: {chi2_pq:.1f} on 9 dof")
    print(f"  (exoplanet 21.5 / 9; asteroid 8170 / 9)")

    # Khinchin
    emp_mean = float(np.mean(ratio_lnK))
    se = float(np.std(ratio_lnK, ddof=1) / math.sqrt(len(ratio_lnK)))
    z_K = (emp_mean - KHINCHIN_LN) / se
    print(f"\nKhinchin: TNO ⟨ln a⟩ = {emp_mean:.5f}  (K = {math.exp(emp_mean):.3f})")
    print(f"  Gauss-Kuzmin: ln K_0 = {KHINCHIN_LN:.5f}  (K_0 = {KHINCHIN_K:.3f})")
    print(f"  z = {z_K:+.2f}")
    print(f"  (exoplanet z = +10.09; asteroid z = +8.92)")

    # Lévy
    emp_L = float(np.mean(ratio_levy))
    se_L = float(np.std(ratio_levy, ddof=1) / math.sqrt(len(ratio_levy)))
    z_L = (emp_L - LEVY_LN) / se_L
    print(f"\nLévy: TNO empirical = {emp_L:.5f}")
    print(f"  Gauss-Kuzmin: L_0 = {LEVY_LN:.5f}")
    print(f"  z = {z_L:+.2f}")
    print(f"  (exoplanet z = +10.95; asteroid z = +11.91)")

    # ----- plot: a-distribution with resonance markers -----
    fig, ax = plt.subplots(figsize=(11, 4.5))
    ax.hist(a_vals, bins=120, color="#16a085", alpha=0.85, edgecolor="white", lw=0.4)
    # Mark known resonance a-values
    # P_TNO/P_N = ratio → a_TNO = (ratio)^(2/3) * 30.07 AU
    A_NEPTUNE = 30.07
    for name, ratio in KNOWN_RES.items():
        a_loc = (ratio) ** (2/3) * A_NEPTUNE
        if a_vals.min() < a_loc < a_vals.max():
            ax.axvline(a_loc, color="#c0392b", lw=1.0, alpha=0.7)
            ax.text(a_loc, ax.get_ylim()[1] * 0.93, name, ha="center", fontsize=8,
                    color="#c0392b", rotation=90)
    ax.set_xlabel("semi-major axis a (AU)")
    ax.set_ylabel("TNO count")
    ax.set_title(f"Kuiper-belt a-distribution (n={len(a_vals)}) with Neptune MMR locations")
    fig.tight_layout()
    fig.savefig("c1_kuiper_a_distribution.png", dpi=140)

    # ----- plot: {r} distribution with sign coloring -----
    fig, ax = plt.subplots(figsize=(11, 4.5))
    colors = ["#16a085" if dev[k] >= 0 else "#c0392b" for k in range(bins)]
    centers = 0.5 * (edges[:-1] + edges[1:])
    width = (edges[1] - edges[0]) * 0.95
    ax.bar(centers, counts, width=width, color=colors, alpha=0.85, edgecolor="white", lw=0.3)
    # Overlay Gauss expected as line
    ax.plot(centers, gauss_exp, color="#2c3e50", lw=2, label="Gauss-measure expected")
    # Mark known resonance fractional parts
    for name, ratio in KNOWN_RES.items():
        f = ratio - math.floor(ratio)
        ax.axvline(f, color="#2c3e50", lw=0.5, alpha=0.5)
        ax.text(f, ax.get_ylim()[1] * 0.92, name, ha="center", fontsize=7,
                color="#2c3e50", rotation=90)
    ax.set_xlabel(r"$\{r\} = P_{\rm TNO}/P_N - \lfloor P_{\rm TNO}/P_N\rfloor$")
    ax.set_ylabel("TNO count")
    ax.set_title(f"Kuiper-belt {{r}} distribution: pile-ups (green, trapping) vs gaps (red, filtering)")
    ax.legend(loc="upper right")
    fig.tight_layout()
    fig.savefig("c1_kuiper_fracs.png", dpi=140)

    # ----- comparison bar plot: Khinchin and Lévy across three populations -----
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.5))
    pops = ["null\n(Gauss-Kuzmin)", "exoplanets\n(trapping)", "asteroids\n(filtering)", "TNOs\n(mixed)"]
    K_vals = [KHINCHIN_LN, 1.194, 1.002, emp_mean]
    L_vals = [LEVY_LN, 1.428, 1.203, emp_L]
    colors_p = ["#7f8c8d", "#27ae60", "#34495e", "#16a085"]
    axes[0].bar(pops, K_vals, color=colors_p, alpha=0.85)
    axes[0].axhline(KHINCHIN_LN, color="#c0392b", lw=1, ls="--")
    axes[0].set_ylabel(r"empirical $\langle \ln a \rangle$")
    axes[0].set_title("Khinchin")
    axes[0].set_ylim(0.97, 1.22)
    for i, v in enumerate(K_vals):
        axes[0].annotate(f"{v:.3f}", (i, v), ha="center", va="bottom", fontsize=9)
    axes[1].bar(pops, L_vals, color=colors_p, alpha=0.85)
    axes[1].axhline(LEVY_LN, color="#c0392b", lw=1, ls="--")
    axes[1].set_ylabel(r"empirical Lévy slope")
    axes[1].set_title("Lévy")
    axes[1].set_ylim(1.15, 1.46)
    for i, v in enumerate(L_vals):
        axes[1].annotate(f"{v:.3f}", (i, v), ha="center", va="bottom", fontsize=9)
    fig.suptitle("Three populations vs Gauss-Kuzmin null: same direction, different magnitudes")
    fig.tight_layout()
    fig.savefig("c1_three_population_compare.png", dpi=140)

    print(f"\nplots: c1_kuiper_a_distribution.png  c1_kuiper_fracs.png  c1_three_population_compare.png")


if __name__ == "__main__":
    main()