"""
Conjecture I, second test: stratify by *dynamical* age, not stellar age.

Dynamical age N_orbits = (stellar_age_Gyr * 1e9 yr) / (P_inner / 365.25 yr)
                      = stellar_age_Gyr * 365.25e9 / P_inner_days

This is "how many times the innermost planet has gone around" — the right
clock for KAM-style stability filtering, since the timescale on which
non-survivors are pruned is orbital, not stellar-evolutionary.

Prediction (Conjecture I):
  systems with high N_orbits should show a *stronger* concentration on
  period ratios with bounded continued-fraction partial quotients
  (more a_k = 1, fewer large a_k) than systems with low N_orbits.

We also separately test ≥3-planet systems, where resonance chains are
unambiguous and selection effects on a single ratio are weaker.
"""

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):
    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 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 load_systems(path):
    """Return {host: {periods: [P, ...], age: Gyr or None}}."""
    systems = defaultdict(lambda: {"periods": [], "age": None})
    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"]
            systems[host]["periods"].append(p)
            try:
                age = float(row["st_age"])
                if math.isfinite(age) and age > 0:
                    systems[host]["age"] = age
            except (ValueError, TypeError):
                pass
    return systems


def collect_pairs(systems, min_planets=2):
    """Yield (host, ratio, log10_Norbits or None, n_planets)."""
    out = []
    for host, info in systems.items():
        periods = sorted(info["periods"])
        n = len(periods)
        if n < min_planets:
            continue
        age = info["age"]
        p_inner = periods[0]
        log_N = None
        if age is not None and p_inner > 0:
            N = age * 365.25e9 / p_inner
            if N > 0 and math.isfinite(N):
                log_N = math.log10(N)
        for i in range(n - 1):
            r = periods[i + 1] / periods[i]
            if r > 1 and math.isfinite(r):
                out.append((host, r, log_N, n))
    return out


def stratify_by_log_N(pairs, depth=8):
    """Return per-pair (log_N, a_k_list) for pairs with measured log_N."""
    data = []
    for host, r, log_N, n in pairs:
        if log_N is None:
            continue
        aks = partial_quotients_after_first(r, depth=depth)
        data.append((log_N, aks))
    return data


def summary_for_group(aks_pooled, label, gk=None, max_m=10):
    counts, freqs = histogram(aks_pooled, max_m=max_m)
    n = counts.sum()
    frac_a1 = sum(1 for v in aks_pooled if v == 1) / n
    frac_large = sum(1 for v in aks_pooled if v >= 4) / n
    print(f"  {label:>30}  n={n:5d}  frac(a=1)={frac_a1:.3f}  frac(a>=4)={frac_large:.3f}")
    if gk is not None:
        expected = gk * n
        mask = expected[1:] > 0
        chi2 = (((counts[1:] - expected[1:]) ** 2) / expected[1:])[mask].sum()
        print(f"  {'':>30}  chi^2 vs GK = {chi2:.1f} (dof={mask.sum()-1})")
    return freqs, counts


def two_proportion_z(k1, n1, k2, n2):
    """Two-proportion z-test. Returns (diff, z, two-sided p)."""
    p1, p2 = k1 / n1, k2 / n2
    p = (k1 + k2) / (n1 + n2)
    se = math.sqrt(p * (1 - p) * (1 / n1 + 1 / n2))
    if se == 0:
        return p1 - p2, float("nan"), float("nan")
    z = (p1 - p2) / se
    # two-sided p via erfc
    p_two = math.erfc(abs(z) / math.sqrt(2))
    return p1 - p2, z, p_two


def main():
    systems = load_systems("ps.csv")
    n_total = len(systems)
    n_multi = sum(1 for s in systems.values() if len(s["periods"]) >= 2)
    n_three = sum(1 for s in systems.values() if len(s["periods"]) >= 3)
    print(f"systems: {n_total}  multi: {n_multi}  3+ planets: {n_three}")

    gk = gk_pmf(max_m=10)

    # ---- all multi-planet pairs ----
    pairs_all = collect_pairs(systems, min_planets=2)
    print(f"\nall adjacent pairs: {len(pairs_all)}")
    data = stratify_by_log_N(pairs_all)
    print(f"pairs with measured stellar age: {len(data)}")

    log_Ns = np.array([d[0] for d in data])
    if len(log_Ns) > 0:
        print(f"log10(N_orbits) range: [{log_Ns.min():.2f}, {log_Ns.max():.2f}]  median: {np.median(log_Ns):.2f}")

    # ---- median split on dynamical age ----
    med = np.median(log_Ns)
    low_aks, high_aks = [], []
    for log_N, aks in data:
        (low_aks if log_N < med else high_aks).extend(aks)

    print("\n--- Median split on log10(N_orbits) [all multi-planet] ---")
    summary_for_group(low_aks, f"low  (log N < {med:.2f})", gk=gk)
    summary_for_group(high_aks, f"high (log N >= {med:.2f})", gk=gk)

    # Two-proportion test on frac(a >= 4)
    k_low_large = sum(1 for v in low_aks if v >= 4)
    k_high_large = sum(1 for v in high_aks if v >= 4)
    d, z, p = two_proportion_z(k_high_large, len(high_aks),
                                k_low_large, len(low_aks))
    print(f"  frac(a>=4) high - low = {d:+.4f}  z={z:+.2f}  p={p:.3f}")
    print(f"  (Conjecture I predicts NEGATIVE: high N_orbits => fewer large a_k)")

    # Same test on frac(a = 1)
    k_low_a1 = sum(1 for v in low_aks if v == 1)
    k_high_a1 = sum(1 for v in high_aks if v == 1)
    d, z, p = two_proportion_z(k_high_a1, len(high_aks), k_low_a1, len(low_aks))
    print(f"  frac(a=1)  high - low = {d:+.4f}  z={z:+.2f}  p={p:.3f}")
    print(f"  (Conjecture I predicts POSITIVE: high N_orbits => more a_k=1)")

    # ---- 3+ planet systems only ----
    pairs_three = [p for p in pairs_all if p[3] >= 3]
    data3 = stratify_by_log_N(pairs_three)
    print(f"\n3+ planet adjacent pairs with measured age: {len(data3)}")
    if data3:
        log_Ns3 = np.array([d[0] for d in data3])
        med3 = np.median(log_Ns3)
        low3, high3 = [], []
        for log_N, aks in data3:
            (low3 if log_N < med3 else high3).extend(aks)
        print(f"\n--- Median split on log10(N_orbits) [3+ planets] ---")
        summary_for_group(low3, f"low  (log N < {med3:.2f})", gk=gk)
        summary_for_group(high3, f"high (log N >= {med3:.2f})", gk=gk)
        k_low_large = sum(1 for v in low3 if v >= 4)
        k_high_large = sum(1 for v in high3 if v >= 4)
        d, z, p = two_proportion_z(k_high_large, len(high3),
                                    k_low_large, len(low3))
        print(f"  frac(a>=4) high - low = {d:+.4f}  z={z:+.2f}  p={p:.3f}")
        k_low_a1 = sum(1 for v in low3 if v == 1)
        k_high_a1 = sum(1 for v in high3 if v == 1)
        d, z, p = two_proportion_z(k_high_a1, len(high3), k_low_a1, len(low3))
        print(f"  frac(a=1)  high - low = {d:+.4f}  z={z:+.2f}  p={p:.3f}")

    # ---- Quartile breakdown for a visual trend ----
    print("\n--- Quartile breakdown (all multi-planet, by log N_orbits) ---")
    if len(data) >= 8:
        sorted_data = sorted(data, key=lambda d: d[0])
        n_per = len(sorted_data) // 4
        chunks = [sorted_data[i*n_per:(i+1)*n_per] for i in range(3)]
        chunks.append(sorted_data[3*n_per:])
        frac_a1_q = []
        frac_large_q = []
        log_N_centers = []
        labels = []
        for i, chunk in enumerate(chunks):
            aks = []
            for log_N, aks_i in chunk:
                aks.extend(aks_i)
            n = len(aks)
            frac_a1 = sum(1 for v in aks if v == 1) / n
            frac_large = sum(1 for v in aks if v >= 4) / n
            log_N_med = np.median([d[0] for d in chunk])
            print(f"  Q{i+1} (median log N={log_N_med:5.2f}): n_aks={n}  frac(a=1)={frac_a1:.3f}  frac(a>=4)={frac_large:.3f}")
            frac_a1_q.append(frac_a1)
            frac_large_q.append(frac_large)
            log_N_centers.append(log_N_med)
            labels.append(f"Q{i+1}\n{log_N_med:.1f}")

        # Plot quartile trend
        fig, ax = plt.subplots(figsize=(7, 4.5))
        ax2 = ax.twinx()
        ax.plot(range(4), frac_a1_q, "o-", color="#27ae60", label="frac($a_k$ = 1)")
        ax2.plot(range(4), frac_large_q, "s-", color="#c0392b", label="frac($a_k \\geq$ 4)")
        ax.axhline(gauss_kuzmin(1), color="#27ae60", ls=":", alpha=0.5,
                   label="GK: $P(a_k$=1)")
        ax2.axhline(sum(gauss_kuzmin(m) for m in range(4, 10)) + (1 - sum(gauss_kuzmin(m) for m in range(1, 10))),
                    color="#c0392b", ls=":", alpha=0.5, label="GK: $P(a_k \\geq$ 4)")
        ax.set_xticks(range(4))
        ax.set_xticklabels(labels)
        ax.set_xlabel("dynamical age quartile (median log$_{10}$ N_orbits)")
        ax.set_ylabel("frac($a_k$ = 1)", color="#27ae60")
        ax2.set_ylabel("frac($a_k \\geq$ 4)", color="#c0392b")
        ax.set_title("Partial-quotient stats vs dynamical age")
        lines1, labels1 = ax.get_legend_handles_labels()
        lines2, labels2 = ax2.get_legend_handles_labels()
        ax.legend(lines1 + lines2, labels1 + labels2, loc="center right", fontsize=8)
        fig.tight_layout()
        fig.savefig("dynamical_age_quartiles.png", dpi=140)

    # ---- All-systems pooled vs GK (sanity, same as analyze.py) ----
    print("\n--- All multi-planet pairs pooled (sanity) ---")
    all_aks = []
    for host, r, _, _ in pairs_all:
        all_aks.extend(partial_quotients_after_first(r))
    summary_for_group(all_aks, "all pairs", gk=gk)


if __name__ == "__main__":
    main()