"""
First-pass test of Conjecture I (Diophantine Astrophysics):
The distribution of partial quotients in continued-fraction expansions of
adjacent planetary period ratios should differ from the Gauss-Kuzmin law,
and old systems should show stronger suppression of large partial quotients.

Gauss-Kuzmin: for a uniform random real in (0,1), partial quotients a_k
(k >= 1) follow P(a_k = m) = log2(1 + 1/(m*(m+2))).

If period ratios behave like generic irrationals, we expect Gauss-Kuzmin.
If KAM-type stability selects them, we expect suppression of large a_k
(large a_k means the irrational is close to a low-denominator rational,
i.e. near a strong resonance — which is dynamically destabilizing).
"""

import math
import csv
from collections import defaultdict
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt


def continued_fraction(x, depth=8, eps=1e-12):
    """Return the first `depth` partial quotients of x > 0."""
    quotients = []
    for _ in range(depth):
        a = math.floor(x)
        quotients.append(int(a))
        frac = x - a
        if frac < eps:
            break
        x = 1.0 / frac
    return quotients


def gauss_kuzmin(m):
    return math.log2(1.0 + 1.0 / (m * (m + 2)))


def load_planets(path):
    systems = defaultdict(list)
    ages = {}
    with open(path) as f:
        reader = csv.DictReader(f)
        for row in reader:
            try:
                p = float(row["pl_orbper"])
            except (ValueError, TypeError):
                continue
            if not math.isfinite(p) or p <= 0:
                continue
            host = row["hostname"]
            systems[host].append(p)
            try:
                ages[host] = float(row["st_age"])
            except (ValueError, TypeError):
                pass
    return systems, ages


def adjacent_ratios(systems):
    """For each multi-planet system, return list of (host, P_outer/P_inner)."""
    out = []
    for host, periods in systems.items():
        if len(periods) < 2:
            continue
        periods = sorted(periods)
        for i in range(len(periods) - 1):
            r = periods[i + 1] / periods[i]
            if r > 1.0 and math.isfinite(r):
                out.append((host, r))
    return out


def partial_quotients_after_first(r, depth=8):
    """Return a_1, a_2, ..., i.e. drop a_0 (the integer part)."""
    cf = continued_fraction(r, depth=depth)
    return cf[1:]


def histogram(values, max_m=10):
    """Empirical frequency of a_k values, clipped at max_m."""
    counts = np.zeros(max_m + 1, dtype=int)  # index 0 unused; indices 1..max_m, max_m = "10+"
    for v in values:
        if v < 1:
            continue
        idx = min(v, max_m)
        counts[idx] += 1
    total = counts.sum()
    freqs = counts / total if total else counts.astype(float)
    return counts, freqs


def gauss_kuzmin_pmf(max_m=10):
    """PMF clipped at max_m: lump >=max_m into the last bin."""
    pmf = np.zeros(max_m + 1)
    for m in range(1, max_m):
        pmf[m] = gauss_kuzmin(m)
    pmf[max_m] = 1.0 - pmf[1:max_m].sum()  # tail
    return pmf


def main():
    systems, ages = load_planets("ps.csv")
    n_multi = sum(1 for ps in systems.values() if len(ps) >= 2)
    print(f"systems total: {len(systems)}; multi-planet systems: {n_multi}")

    ratios = adjacent_ratios(systems)
    print(f"adjacent ratio pairs: {len(ratios)}")

    # Collect a_1, a_2, ... a_8 across all pairs
    all_aks = []  # flat list
    per_host = []  # (host, list of a_k)
    for host, r in ratios:
        aks = partial_quotients_after_first(r, depth=8)
        per_host.append((host, r, aks))
        all_aks.extend(aks)

    max_m = 10
    counts, freqs = histogram(all_aks, max_m=max_m)
    gk = gauss_kuzmin_pmf(max_m=max_m)

    print("\n  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 test by hand
    n = counts.sum()
    expected = gk * n
    mask = expected[1:] > 0
    chi2 = (((counts[1:] - expected[1:]) ** 2) / expected[1:])[mask].sum()
    dof = mask.sum() - 1
    print(f"\nchi^2 = {chi2:.1f} on {dof} dof  (large => non-Gauss-Kuzmin)")

    # Save ratios with their a_k values
    with open("ratios.csv", "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["hostname", "ratio"] + [f"a{k}" for k in range(1, 9)])
        for host, r, aks in per_host:
            row = [host, f"{r:.6f}"] + aks + [""] * (8 - len(aks))
            w.writerow(row)

    # ---- Plot 1: empirical vs Gauss-Kuzmin -----------------------------
    fig, ax = plt.subplots(figsize=(7, 4))
    xs = np.arange(1, max_m + 1)
    width = 0.4
    ax.bar(xs - width/2, freqs[1:], width, label=f"Exoplanets (n={n})", color="#c0392b")
    ax.bar(xs + width/2, gk[1:], width, label="Gauss-Kuzmin", color="#2c3e50", alpha=0.7)
    ax.set_xlabel("partial quotient $a_k$ (k >= 1)")
    ax.set_ylabel("frequency")
    ax.set_title("Partial-quotient distribution of adjacent exoplanet period ratios")
    xt = [str(m) for m in range(1, max_m)] + [f">={max_m}"]
    ax.set_xticks(xs)
    ax.set_xticklabels(xt)
    ax.legend()
    fig.tight_layout()
    fig.savefig("partial_quotients.png", dpi=140)

    # ---- Age stratification (Conjecture I) -----------------------------
    young, old = [], []
    n_young_sys = n_old_sys = 0
    for host, r, aks in per_host:
        if host not in ages:
            continue
        age = ages[host]
        if age < 2.0:
            young.extend(aks); n_young_sys += 1
        elif age > 5.0:
            old.extend(aks); n_old_sys += 1

    print(f"\nyoung (<2 Gyr) ratio-pairs contributing a_k: {n_young_sys}, total a_k samples: {len(young)}")
    print(f"old   (>5 Gyr) ratio-pairs contributing a_k: {n_old_sys}, total a_k samples: {len(old)}")

    if young and old:
        _, fy = histogram(young, max_m=max_m)
        _, fo = histogram(old, max_m=max_m)

        # Fraction of "large" a_k (>=4) -- prediction: lower in old
        frac_large_young = sum(1 for v in young if v >= 4) / len(young)
        frac_large_old = sum(1 for v in old if v >= 4) / len(old)
        print(f"frac a_k >= 4  young: {frac_large_young:.3f}   old: {frac_large_old:.3f}")
        print("Conjecture I predicts: frac(old) < frac(young)")

        fig, ax = plt.subplots(figsize=(7, 4))
        ax.bar(xs - width/2, fy[1:], width, label=f"young <2 Gyr (n={len(young)})", color="#27ae60")
        ax.bar(xs + width/2, fo[1:], width, label=f"old >5 Gyr (n={len(old)})", color="#8e44ad")
        ax.set_xlabel("partial quotient $a_k$")
        ax.set_ylabel("frequency")
        ax.set_title("Age stratification: partial-quotient distribution")
        ax.set_xticks(xs); ax.set_xticklabels(xt)
        ax.legend()
        fig.tight_layout()
        fig.savefig("age_stratified.png", dpi=140)


if __name__ == "__main__":
    main()