"""
Diophantine analysis of the main asteroid belt — parallel to Tests 1 and 5
on exoplanets.

Hypothesis (Conjecture I‴, population-comparison reformulation):
KAM-stability filtering dominates in long-lived, low-mass, dense populations
(asteroids). Resonance trapping during migration dominates in short-lived,
moderate-mass, sparse populations (exoplanets). The two should appear in
opposite signs of the Khinchin, Lévy, and pile-up tests.

Predictions for asteroids (KAM-filtered):
  - Khinchin K(P_J/P_a) < K_0 = 2.685     (more noble irrationals)
  - Lévy L(P_J/P_a) < L_0 = 1.187
  - {r} distribution shows GAPS at low-order rationals (the Kirkwood gaps),
    not pile-ups.

We compute period ratios r = P_Jupiter / P_asteroid for 100 k main-belt
asteroids (2.0 ≤ a ≤ 3.5 AU) from JPL SBDB.
"""

import json
import math
from collections import defaultdict
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_JUPITER_DAYS = 11.862 * 365.25  # = 4332.6 days


def load_asteroid_data(path):
    with open(path) as f:
        d = json.load(f)
    rows = []
    for full_name, a, e, per, H in d["data"]:
        try:
            a_au = float(a); e_v = float(e); per_d = float(per)
            H_v = float(H) if H is not None else float("nan")
        except (TypeError, ValueError):
            continue
        if not (math.isfinite(a_au) and math.isfinite(per_d) and per_d > 0):
            continue
        rows.append((full_name.strip(), a_au, e_v, per_d, H_v))
    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 gauss_kuzmin_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():
    print("Loading asteroid data...")
    asteroids = load_asteroid_data("asteroids_mb.json")
    print(f"  {len(asteroids)} asteroids with valid (a, per)")

    # Compute period ratio r = P_jupiter / P_asteroid
    rs = []
    a_vals = []
    for name, a, e, per, H in asteroids:
        r = P_JUPITER_DAYS / per
        if r > 1 and math.isfinite(r):
            rs.append(r)
            a_vals.append(a)
    rs = np.array(rs)
    a_vals = np.array(a_vals)
    print(f"  computed {len(rs)} period ratios in (1, ∞)")
    print(f"  r range: [{rs.min():.3f}, {rs.max():.3f}]")
    print(f"  a range: [{a_vals.min():.3f}, {a_vals.max():.3f}] AU")

    # Restrict to main-belt-relevant range
    sel = (rs > 1.5) & (rs < 5.0)
    rs = rs[sel]; a_vals = a_vals[sel]
    print(f"  after restricting to r ∈ (1.5, 5.0): {len(rs)} asteroids")

    # =============================================================
    # Test A — partial-quotient distribution vs Gauss-Kuzmin
    # =============================================================
    all_aks = []
    ratio_lnK = []
    ratio_levy = []
    for r in rs:
        aks, lnK, L = per_ratio_stats(r, 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
    counts = np.zeros(max_m + 1, dtype=int)
    for v in all_aks:
        if v >= 1:
            counts[min(v, max_m)] += 1
    total = counts.sum()
    freqs = counts / total
    gk = gauss_kuzmin_pmf(max_m=max_m)

    print(f"\nTest A — partial-quotient distribution:")
    print(f"  n_aks = {total} from {len(rs)} ratios")
    print(f"  m   empirical   Gauss-Kuzmin   ratio")
    for m in range(1, max_m + 1):
        label = f">={max_m}" if m == max_m else str(m)
        print(f"  {label:>4}  {freqs[m]:.4f}      {gk[m]:.4f}        {freqs[m]/gk[m]:.2f}")

    # Chi-square
    expected = gk * total
    chi2 = (((counts[1:] - expected[1:]) ** 2) / expected[1:]).sum()
    dof = max_m - 1
    print(f"  chi² = {chi2:.1f} on {dof} dof  (exoplanet was 21.5/9)")

    # =============================================================
    # Test B — Khinchin
    # =============================================================
    print(f"\nTest B — Khinchin constant:")
    print(f"  theoretical ln K_0 = {KHINCHIN_LN:.5f}  (K_0 = {KHINCHIN_K:.5f})")
    emp_mean = float(np.mean(ratio_lnK))
    emp_std = float(np.std(ratio_lnK, ddof=1))
    se = emp_std / math.sqrt(len(ratio_lnK))
    z = (emp_mean - KHINCHIN_LN) / se
    print(f"  asteroid empirical ⟨ln a⟩ = {emp_mean:.5f}  (K = {math.exp(emp_mean):.5f})")
    print(f"  SE = {se:.5f}  z = {z:+.2f}")
    print(f"  exoplanet Test 5 result: ⟨ln a⟩ = 1.194, z = +10.09  (toward rationals)")
    direction = "toward rationals" if z > 0 else "away from rationals (KAM-noble)"
    print(f"  asteroid direction: {direction}")

    # =============================================================
    # Test C — Lévy
    # =============================================================
    print(f"\nTest C — Lévy constant:")
    print(f"  theoretical L_0 = {LEVY_LN:.5f}")
    emp_L = float(np.mean(ratio_levy))
    emp_L_std = float(np.std(ratio_levy, ddof=1))
    se_L = emp_L_std / math.sqrt(len(ratio_levy))
    z_L = (emp_L - LEVY_LN) / se_L
    print(f"  asteroid empirical Lévy = {emp_L:.5f}")
    print(f"  SE = {se_L:.5f}  z = {z_L:+.2f}")
    print(f"  exoplanet Test 5 result: Lévy = 1.428, z = +10.95  (toward rationals)")

    # =============================================================
    # Test D — bin histogram of {r} with Kirkwood gap inspection
    # =============================================================
    fracs = rs - np.floor(rs)
    counts_bin, edges = np.histogram(fracs, bins=20, range=(0, 1))
    n = counts_bin.sum()
    gauss_exp = np.array([n * math.log2((1 + edges[i+1]) / (1 + edges[i]))
                          for i in range(20)])
    unif_exp = np.full(20, n / 20)
    chi2_g = ((counts_bin - gauss_exp) ** 2 / gauss_exp).sum()
    chi2_u = ((counts_bin - unif_exp) ** 2 / unif_exp).sum()
    print(f"\nTest D — {{r}} bin distribution:")
    print(f"  chi² vs Gauss = {chi2_g:.1f} on 19 dof  (exoplanet: 111.2)")
    print(f"  chi² vs uniform = {chi2_u:.1f} on 19 dof  (exoplanet: 50.5)")

    dev = (counts_bin - gauss_exp) / np.sqrt(gauss_exp)
    print(f"  top deficits (KAM-style gaps):")
    bot5 = np.argsort(dev)[:5]
    for k in bot5:
        lo, hi = edges[k], edges[k+1]
        print(f"    [{lo:.3f}, {hi:.3f}): obs={counts_bin[k]:6d}  exp={gauss_exp[k]:6.0f}  dev={dev[k]:+.2f}")

    # Also compare with a-distribution (where Kirkwood gaps live)
    print(f"\n  a-distribution (semi-major axis) gap inspection:")
    a_hist, a_edges = np.histogram(a_vals, bins=300, range=(2.0, 3.5))
    # Find the deepest gaps in normalized smoothed a distribution
    from scipy.ndimage import gaussian_filter1d
    a_smooth = gaussian_filter1d(a_hist.astype(float), sigma=2.5)
    bg = gaussian_filter1d(a_hist.astype(float), sigma=15)
    relative = a_smooth / np.maximum(bg, 1)
    # Find local minima of relative (deepest dips)
    gap_mask = (relative < 0.7)
    gap_idxs = np.where(gap_mask)[0]
    if len(gap_idxs):
        # Cluster contiguous gap regions
        groups = []
        current = [gap_idxs[0]]
        for idx in gap_idxs[1:]:
            if idx - current[-1] <= 2:
                current.append(idx)
            else:
                groups.append(current); current = [idx]
        groups.append(current)
        print(f"  found {len(groups)} candidate gap regions:")
        # Theoretical Kirkwood gap locations
        kirkwood = {
            "4:1": 2.064, "7:2": 2.255, "3:1": 2.502, "8:3": 2.706,
            "5:2": 2.825, "7:3": 2.957, "9:4": 3.029, "2:1": 3.279,
        }
        for grp in groups:
            a_mid = 0.5 * (a_edges[grp[0]] + a_edges[grp[-1] + 1])
            # Match to nearest Kirkwood
            nearest = min(kirkwood.items(), key=lambda kv: abs(kv[1] - a_mid))
            dist = abs(nearest[1] - a_mid)
            tag = nearest[0] if dist < 0.05 else "—"
            print(f"    a ≈ {a_mid:.3f} AU  (depth {min(relative[grp]):.2f})  → {tag}")

    # =============================================================
    # Plots
    # =============================================================
    # Asteroid semi-major-axis histogram with Kirkwood gap markers
    fig, ax = plt.subplots(figsize=(11, 4.5))
    ax.hist(a_vals, bins=300, color="#34495e", alpha=0.85, edgecolor="white", lw=0.2)
    kirkwood_full = {
        "4:1": 2.064, "7:2": 2.255, "3:1": 2.502, "8:3": 2.706,
        "5:2": 2.825, "7:3": 2.957, "9:4": 3.029, "2:1": 3.279,
    }
    for lbl, a_loc in kirkwood_full.items():
        ax.axvline(a_loc, color="#c0392b", lw=1.0, alpha=0.7)
        ax.text(a_loc, ax.get_ylim()[1] * 0.93, lbl, ha="center",
                fontsize=8, color="#c0392b", rotation=90)
    ax.set_xlabel("semi-major axis a (AU)")
    ax.set_ylabel("asteroid count")
    ax.set_xlim(2.0, 3.5)
    ax.set_title(f"Main-belt asteroid a-distribution (n={len(a_vals)}) with Kirkwood gaps")
    fig.tight_layout()
    fig.savefig("c1_asteroid_a_distribution.png", dpi=140)

    # Khinchin comparison: asteroid vs exoplanet
    fig, ax = plt.subplots(figsize=(8, 4.5))
    ax.hist(ratio_lnK, bins=60, density=True, color="#34495e", alpha=0.85,
            edgecolor="white", lw=0.4, label=f"asteroids (n={len(ratio_lnK)})")
    ax.axvline(KHINCHIN_LN, color="#c0392b", lw=2, label=f"Gauss-Kuzmin ln K₀ = {KHINCHIN_LN:.3f}")
    ax.axvline(emp_mean, color="#2c3e50", lw=2, ls="--",
               label=f"asteroid mean = {emp_mean:.3f} (z = {z:+.2f})")
    ax.axvline(1.194, color="#27ae60", lw=2, ls=":",
               label="exoplanet mean = 1.194 (z = +10.09)")
    ax.set_xlabel(r"per-ratio empirical $\langle\ln a_k\rangle$")
    ax.set_ylabel("density")
    ax.set_title("Khinchin test: asteroid belt vs exoplanets vs Gauss-Kuzmin")
    ax.legend(loc="upper left", fontsize=9)
    fig.tight_layout()
    fig.savefig("c1_asteroid_khinchin.png", dpi=140)

    # Period-ratio bin distribution
    fig, ax = plt.subplots(figsize=(11, 4.5))
    ax.hist(fracs, bins=80, density=True, color="#34495e", alpha=0.7, edgecolor="white", lw=0.3)
    # Annotate where Jupiter MMRs land
    # For asteroid r = P_J/P_a, integer r → integer MMR (a moves into Jupiter's, asteroid faster).
    # r=2.0 ({r}=0) → 2:1, r=3:1, etc. Half-integers: r=2.5 → 5:2, etc.
    annotations = {0.0: "integer (3:1, 4:1)", 0.5: "5:2, 7:2",
                   0.333: "7:3, 10:3", 0.667: "8:3, 11:3"}
    for x, lbl in annotations.items():
        ax.axvline(x, color="#c0392b", lw=0.6, alpha=0.5)
        ax.text(x, ax.get_ylim()[1] * 0.92, lbl, ha="center", fontsize=8, color="#c0392b", rotation=90)
    ax.set_xlabel(r"$\{r\} = P_J/P_a - \lfloor P_J/P_a \rfloor$")
    ax.set_ylabel("density")
    ax.set_title("Asteroid {r} distribution — pile-ups or gaps at Kirkwood resonances?")
    fig.tight_layout()
    fig.savefig("c1_asteroid_fracs.png", dpi=140)

    print(f"\nplots: c1_asteroid_a_distribution.png  c1_asteroid_khinchin.png  c1_asteroid_fracs.png")


if __name__ == "__main__":
    main()