"""
Two analyses in one pass.

(A) Conjecture I, RV-only test. The χ²=21.5 deviation from Gauss-Kuzmin
    in the pooled multi-planet sample could be due to Kepler transit-survey
    selection effects (the well-known 3:2/2:1 pile-ups, Fabrycky+ 2014).
    RV-discovered systems have very different biases — they favor short
    periods and high masses but do *not* preferentially detect near-resonant
    pairs. If the χ² signal survives in the RV subset, selection bias is
    not the explanation.

(B) Conjecture II, first concrete test (resonance lattice rigidity).
    For each adjacent pair in n≥3-planet systems, find the (p,q) with
    gcd=1, p > q ≥ 1, p+q ≤ 9 that minimizes |r - p/q|.  The "chain type"
    of an n-planet system is the tuple of these labels. We ask:
      (i) How many distinct chain types appear vs combinatorially possible?
     (ii) For 3-planet systems, plot (r₁₂, r₂₃) and look for clustering.
    Lattice rigidity predicts a discrete, sparse set of populated chain
    types — not a uniform cloud.
"""

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


def continued_fraction(x, depth=8, eps=1e-12):
    qs = []
    for _ in range(depth):
        a = math.floor(x)
        qs.append(int(a))
        frac = x - a
        if frac < eps:
            break
        x = 1.0 / frac
    return qs


def partial_quotients_after_first(r, depth=8):
    return continued_fraction(r, depth=depth)[1:]


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


def gk_pmf(max_m=10):
    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()
    return pmf


def histogram(values, max_m=10):
    counts = np.zeros(max_m + 1, dtype=int)
    for v in values:
        if v < 1:
            continue
        counts[min(v, max_m)] += 1
    total = counts.sum()
    freqs = counts / total if total else counts.astype(float)
    return counts, freqs


def chi2_vs_gk(counts, gk):
    n = counts.sum()
    expected = gk * n
    mask = expected[1:] > 0
    chi2 = (((counts[1:] - expected[1:]) ** 2) / expected[1:])[mask].sum()
    return chi2, int(mask.sum() - 1)


def load(path):
    rows = defaultdict(list)  # host -> list of (period, dm)
    ages = {}
    with open(path) as f:
        for row in csv.DictReader(f):
            try:
                p = float(row["pl_orbper"])
            except (ValueError, TypeError):
                continue
            if not math.isfinite(p) or p <= 0:
                continue
            host = row["hostname"]
            dm = (row.get("discoverymethod") or "").strip().strip('"')
            rows[host].append((p, dm))
            try:
                age = float(row["st_age"])
                if math.isfinite(age) and age > 0:
                    ages[host] = age
            except (ValueError, TypeError):
                pass
    return rows, ages


def rv_only_systems(rows, min_planets=2):
    """Systems where every detected planet was RV-discovered."""
    out = {}
    for host, plist in rows.items():
        if len(plist) < min_planets:
            continue
        if all(dm == "Radial Velocity" for _, dm in plist):
            out[host] = sorted(p for p, _ in plist)
    return out


def adjacent_ratios(systems):
    """systems: {host: [periods sorted asc]} -> list of (host, r > 1, n_planets)."""
    out = []
    for host, periods in systems.items():
        n = len(periods)
        for i in range(n - 1):
            r = periods[i + 1] / periods[i]
            if r > 1 and math.isfinite(r):
                out.append((host, r, n))
    return out


# ---------------- Conjecture II machinery -------------------

def farey_pairs(max_sum):
    """All (p, q) with gcd(p,q)=1, p > q >= 1, p+q <= max_sum, sorted by p/q."""
    pairs = []
    for q in range(1, max_sum):
        for p in range(q + 1, max_sum - q + 1):
            if math.gcd(p, q) == 1:
                pairs.append((p, q))
    pairs.sort(key=lambda pq: pq[0] / pq[1])
    return pairs


def nearest_resonance(r, pairs):
    """Return ((p, q), |r - p/q|, log10|r - p/q|)."""
    best = min(pairs, key=lambda pq: abs(r - pq[0] / pq[1]))
    delta = abs(r - best[0] / best[1])
    return best, delta


# ----------------------- Main ------------------------------

def main():
    rows, ages = load("ps_dm.csv")
    print(f"hosts with any planet: {len(rows)}")

    # ----- All multi (reference) -----
    all_multi = {h: sorted(p for p, _ in pl) for h, pl in rows.items() if len(pl) >= 2}
    print(f"all multi-planet systems: {len(all_multi)}")
    all_pairs = adjacent_ratios(all_multi)
    aks_all = []
    for h, r, n in all_pairs:
        aks_all.extend(partial_quotients_after_first(r))
    cnt, frq = histogram(aks_all)
    chi, dof = chi2_vs_gk(cnt, gk_pmf())
    print(f"\n[ALL]    pairs={len(all_pairs)}  a_k samples={cnt.sum()}  "
          f"frac(a=1)={frq[1]:.3f}  chi^2={chi:.1f}/{dof}dof")

    # ----- RV-only -----
    rv_systems = rv_only_systems(rows, min_planets=2)
    rv_pairs = adjacent_ratios(rv_systems)
    print(f"\nRV-only multi-planet systems: {len(rv_systems)}  "
          f"adjacent pairs: {len(rv_pairs)}")
    aks_rv = []
    for h, r, n in rv_pairs:
        aks_rv.extend(partial_quotients_after_first(r))
    cnt_rv, frq_rv = histogram(aks_rv)
    chi_rv, dof_rv = chi2_vs_gk(cnt_rv, gk_pmf())
    print(f"[RV]     a_k samples={cnt_rv.sum()}  frac(a=1)={frq_rv[1]:.3f}  "
          f"chi^2={chi_rv:.1f}/{dof_rv}dof")

    # Per-m table
    gk = gk_pmf()
    print("\n  m   ALL-emp  RV-emp   GK     ratio_RV  ratio_ALL")
    for m in range(1, 11):
        label = f">={10}" if m == 10 else str(m)
        print(f"  {label:>3}  {frq[m]:.4f}  {frq_rv[m]:.4f}  {gk[m]:.4f}  "
              f"{frq_rv[m]/gk[m]:.2f}      {frq[m]/gk[m]:.2f}")

    # ----- Transit-only for contrast -----
    transit_systems = {}
    for host, plist in rows.items():
        if len(plist) < 2:
            continue
        if all(dm == "Transit" for _, dm in plist):
            transit_systems[host] = sorted(p for p, _ in plist)
    tr_pairs = adjacent_ratios(transit_systems)
    aks_tr = []
    for h, r, n in tr_pairs:
        aks_tr.extend(partial_quotients_after_first(r))
    cnt_tr, frq_tr = histogram(aks_tr)
    chi_tr, dof_tr = chi2_vs_gk(cnt_tr, gk_pmf())
    print(f"\n[TRANSIT-only] systems={len(transit_systems)}  pairs={len(tr_pairs)}  "
          f"a_k samples={cnt_tr.sum()}  frac(a=1)={frq_tr[1]:.3f}  "
          f"chi^2={chi_tr:.1f}/{dof_tr}dof")

    # Plot: ALL vs RV vs Transit vs GK
    fig, ax = plt.subplots(figsize=(9, 4.5))
    xs = np.arange(1, 11)
    w = 0.2
    ax.bar(xs - 1.5*w, frq[1:],   w, label=f"all multi (n={cnt.sum()})",     color="#7f8c8d")
    ax.bar(xs - 0.5*w, frq_rv[1:], w, label=f"RV-only (n={cnt_rv.sum()})",   color="#27ae60")
    ax.bar(xs + 0.5*w, frq_tr[1:], w, label=f"transit-only (n={cnt_tr.sum()})", color="#c0392b")
    ax.bar(xs + 1.5*w, gk[1:],     w, label="Gauss-Kuzmin", color="#2c3e50", alpha=0.7)
    xt = [str(m) for m in range(1, 10)] + [">=10"]
    ax.set_xticks(xs); ax.set_xticklabels(xt)
    ax.set_xlabel("partial quotient $a_k$")
    ax.set_ylabel("frequency")
    ax.set_title("Discovery-method comparison: partial-quotient distribution vs Gauss-Kuzmin")
    ax.legend()
    fig.tight_layout()
    fig.savefig("rv_vs_transit_vs_gk.png", dpi=140)

    # ============== Conjecture II ==============
    print("\n" + "=" * 60)
    print("CONJECTURE II — resonance chain combinatorics")
    print("=" * 60)

    pairs_table = farey_pairs(max_sum=9)
    print(f"Farey 'resonance alphabet' (p,q): gcd=1, p>q>=1, p+q<=9: {len(pairs_table)} symbols")

    # Three+-planet systems
    n3 = {h: ps for h, ps in all_multi.items() if len(ps) >= 3}
    print(f"3+ planet systems: {len(n3)}")

    chain_types = Counter()
    deltas = []
    chain_to_hosts = defaultdict(list)
    for host, periods in n3.items():
        chain = []
        ok = True
        for i in range(len(periods) - 1):
            r = periods[i + 1] / periods[i]
            (p, q), d = nearest_resonance(r, pairs_table)
            chain.append((p, q))
            deltas.append((r, p, q, d, host))
            if d / (p / q) > 0.15:  # >15% off — too far to be a "resonance"
                ok = False
        chain_types[tuple(chain)] += 1
        chain_to_hosts[tuple(chain)].append(host)

    print(f"distinct chain types observed: {len(chain_types)}")
    print(f"most common chain types:")
    for ct, c in chain_types.most_common(10):
        pretty = " - ".join(f"{p}:{q}" for p, q in ct)
        print(f"  {c:4d}  {pretty}")

    # Distribution of chain lengths
    by_len = Counter(len(ct) for ct in chain_types)
    print(f"\ndistribution of chain lengths (= n_planets-1):")
    for L, c in sorted(by_len.items()):
        possible = len(pairs_table) ** L
        observed = sum(1 for ct in chain_types if len(ct) == L)
        n_systems = sum(chain_types[ct] for ct in chain_types if len(ct) == L)
        print(f"  L={L}: {observed} distinct types observed in {n_systems} systems  "
              f"(combinatorially possible w/ |alphabet|={len(pairs_table)}: {possible:,})")

    # Distance-to-nearest-resonance distribution
    deltas_arr = np.array([d[3] for d in deltas])
    rs_arr = np.array([d[0] for d in deltas])
    print(f"\ndistance-to-nearest-resonance (n={len(deltas)}):")
    print(f"  median |r - p/q|: {np.median(deltas_arr):.4f}")
    print(f"  fraction within 1% of a Farey-9 ratio: {(deltas_arr < 0.01).mean():.3f}")
    print(f"  fraction within 5% of a Farey-9 ratio: {(deltas_arr < 0.05).mean():.3f}")

    # ---- Plot: r vs nearest resonance distance ----
    fig, ax = plt.subplots(figsize=(9, 4.5))
    ax.scatter(rs_arr, deltas_arr, s=8, alpha=0.4, color="#2c3e50")
    for p, q in pairs_table:
        ax.axvline(p / q, color="#c0392b", lw=0.4, alpha=0.3)
    ax.set_xlim(1.0, 4.0)
    ax.set_yscale("log")
    ax.set_xlabel("period ratio $r = P_{i+1}/P_i$")
    ax.set_ylabel(r"$|r - p/q|$  (nearest Farey-9 commensurability)")
    ax.set_title("Conjecture II: proximity to nearest mean-motion resonance")
    fig.tight_layout()
    fig.savefig("resonance_proximity.png", dpi=140)

    # ---- Plot: 3-planet systems in the (r12, r23) plane ----
    threes = [(host, ps) for host, ps in all_multi.items() if len(ps) == 3]
    fours_plus = [(host, ps) for host, ps in all_multi.items() if len(ps) >= 4]
    pts3 = np.array([[ps[1]/ps[0], ps[2]/ps[1]] for _, ps in threes])
    print(f"\n3-planet systems exactly: {len(threes)}")
    print(f"4+ planet systems:         {len(fours_plus)}")

    fig, ax = plt.subplots(figsize=(7, 7))
    ax.scatter(pts3[:, 0], pts3[:, 1], s=14, alpha=0.55, color="#16a085",
               label=f"3-planet systems (n={len(threes)})")
    # Overlay 4+ systems' first two ratios
    if fours_plus:
        pts4 = np.array([[ps[1]/ps[0], ps[2]/ps[1]] for _, ps in fours_plus])
        ax.scatter(pts4[:, 0], pts4[:, 1], s=12, alpha=0.55, color="#8e44ad",
                   marker="x", label=f"4+ planet systems, first two ratios (n={len(fours_plus)})")
    # Famous benchmarks
    benchmarks = {
        "Trappist-1 b/c/d": (1.5117/1.5108 * 1.5108/1.0, None),  # placeholder; will fix
    }
    # Use real Trappist-1 period ratios
    trappist_periods = [1.51087, 2.42180, 4.04961, 6.09961, 9.20669, 12.35294, 18.76673]
    tp = trappist_periods
    for i in range(len(tp) - 2):
        ax.scatter([tp[i+1]/tp[i]], [tp[i+2]/tp[i+1]], s=80, marker="*",
                   color="#e74c3c", edgecolor="black", linewidth=0.5,
                   label="Trappist-1" if i == 0 else None, zorder=5)
    # Grid lines for low-order resonances
    for p, q in [(2,1),(3,2),(4,3),(5,3),(5,4),(7,5),(3,1)]:
        v = p / q
        ax.axvline(v, color="#888", lw=0.4, alpha=0.5)
        ax.axhline(v, color="#888", lw=0.4, alpha=0.5)
    ax.set_xlim(1, 4)
    ax.set_ylim(1, 4)
    ax.set_xlabel(r"$r_{12} = P_2 / P_1$")
    ax.set_ylabel(r"$r_{23} = P_3 / P_2$")
    ax.set_title("Conjecture II: 3-planet systems in the period-ratio plane")
    ax.legend(loc="upper right", fontsize=8)
    ax.set_aspect("equal")
    fig.tight_layout()
    fig.savefig("c2_three_planet_plane.png", dpi=140)


if __name__ == "__main__":
    main()