"""
Conjecture I, deeper analysis.

Three independent measures of how "Gauss-Kuzmin-like" the partial-quotient
expansion of an observed period ratio is:

1. KHINCHIN CONSTANT. For almost every real x (Lebesgue), the geometric
   mean of partial quotients converges to K_0 = 2.6854520010...  More
   precisely:  E[ln a_k] under Gauss-Kuzmin = ln K_0 ≈ 0.98785.
   If period ratios are "more rational" / "more noble" than generic,
   their empirical ln-a-mean should be < 0.98785.

2. LEVY CONSTANT. For almost every real, log q_n / n -> L_0 = π²/(12 ln 2)
   ≈ 1.18657, where q_n is the denominator of the n-th convergent.
   Noble-ish irrationals have smaller L (golden ratio: L = ln φ ≈ 0.481).
   Period-ratio shift should be downward.

3. CONDITIONAL DEPENDENCE. Under Gauss-Kuzmin, (a_k) are NOT iid, but
   the asymptotic joint distribution is the iid product. For finite k
   the dependence is weak. We test P(a_{k+1} | a_k) vs P(a_{k+1}):
   strong departure from independence is a structural signal.

We also localize: stratify by r ranges (1<r<2, 2<r<3, etc.) to see
where the bulk of the deviation lives.
"""

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


KHINCHIN_LN = 0.9878490568339858    # E[ln a] under Gauss-Kuzmin
KHINCHIN_K  = math.exp(KHINCHIN_LN)  # ~ 2.6854520
LEVY_LN     = math.pi ** 2 / (12 * math.log(2))  # ~ 1.18657


def continued_fraction_and_convergents(x, depth=10, eps=1e-12):
    """Return (a_list, convergents) where convergents are (p_n, q_n) pairs."""
    a_list = []
    p_prev2, p_prev1 = 1, 0   # p_{-1}=1, p_{-2}=0 -- standard recurrence
    q_prev2, q_prev1 = 0, 1
    convergents = []
    for _ in range(depth):
        a = math.floor(x)
        a_list.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 a_list, convergents


def per_ratio_stats(r, depth=10):
    """Return (a_list_after_first, log_q_list, ln_a_mean_post1, levy_estimate)."""
    a_list, conv = continued_fraction_and_convergents(r, depth=depth)
    a_post = a_list[1:]  # drop a_0 (integer part of r > 1)
    if not a_post:
        return [], [], None, None
    ln_a_mean = float(np.mean([math.log(a) for a in a_post if a > 0]))
    # Levy estimate: slope of log q_n vs n for n >= 2 (n=1 trivial: q_1 = 1)
    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])
        # use indices >= 2 (drop trivial first point)
        slope = np.polyfit(ns[1:], logs[1:], 1)[0]
        levy = float(slope)
    return a_post, q_list, ln_a_mean, levy


def load_ratios(path):
    """Return list of (host, r) for adjacent pairs in multi-planet systems."""
    systems = defaultdict(list)
    with open(path) as f:
        for row in csv.DictReader(f):
            try:
                p = float(row["pl_orbper"])
            except (ValueError, TypeError):
                continue
            if math.isfinite(p) and p > 0:
                systems[row["hostname"]].append(p)
    out = []
    for host, periods in systems.items():
        if len(periods) < 2:
            continue
        for a, b in zip(sorted(periods)[:-1], sorted(periods)[1:]):
            r = b / a
            if r > 1 and math.isfinite(r):
                out.append((host, r))
    return out


def gauss_kuzmin_pmf(max_m=30):
    return np.array([0.0] + [math.log2(1 + 1/(m*(m+2))) for m in range(1, max_m + 1)])


def main():
    pairs = load_ratios("ps_dm.csv")
    print(f"loaded {len(pairs)} adjacent ratio pairs")

    ratio_levy = []
    ratio_lnK = []
    ratio_rs = []
    all_aks = []
    pair_aks = []  # list of (a_k, a_{k+1}) consecutive

    for host, r in pairs:
        aks, qs, lnK, L = per_ratio_stats(r, depth=10)
        if lnK is not None:
            ratio_lnK.append(lnK)
            ratio_rs.append(r)
        if L is not None:
            ratio_levy.append(L)
        all_aks.extend(aks)
        for k in range(len(aks) - 1):
            pair_aks.append((aks[k], aks[k + 1]))

    ratio_lnK = np.array(ratio_lnK)
    ratio_levy = np.array(ratio_levy)
    ratio_rs = np.array(ratio_rs)

    # ==================================================================
    # 1. KHINCHIN
    # ==================================================================
    print("\n========== Khinchin constant test ==========")
    print(f"theoretical E[ln a] under Gauss-Kuzmin: {KHINCHIN_LN:.5f}")
    print(f"theoretical K_0 = exp(.): {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"empirical mean of ln-Khinchin (per-ratio, depth-10 truncated): {emp_mean:.5f}")
    print(f"empirical std: {emp_std:.5f}  SE: {se:.5f}  z = {z:+.2f}")
    print(f"empirical K (= exp mean): {math.exp(emp_mean):.5f}")
    # Negative z = period ratios are MORE 'noble' than random — supports Conj I

    # ==================================================================
    # 2. LEVY
    # ==================================================================
    print("\n========== Levy constant test ==========")
    print(f"theoretical L_0 = π²/(12 ln 2): {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"empirical mean Levy estimate: {emp_L:.5f}")
    print(f"std: {emp_L_std:.5f}  SE: {se_L:.5f}  z = {z_L:+.2f}")

    # ==================================================================
    # 3. CONDITIONAL DEPENDENCE
    # ==================================================================
    print("\n========== Conditional dependence test ==========")
    # Build 5x5 contingency table on (a_k, a_{k+1}) with a in {1,2,3,4,>=5}
    def bucket(v):
        return min(v, 5) if v >= 1 else 0
    ct = np.zeros((5, 5), dtype=int)
    for a, b in pair_aks:
        i = bucket(a) - 1
        j = bucket(b) - 1
        if 0 <= i < 5 and 0 <= j < 5:
            ct[i, j] += 1
    print(f"5x5 contingency table on (a_k, a_{{k+1}}) with buckets 1,2,3,4,>=5:")
    print("           a_{k+1}=1  =2     =3     =4     >=5")
    for i in range(5):
        row_label = f"a_k={'>=5' if i==4 else i+1}"
        print(f"  {row_label:10s}  {ct[i,0]:6d} {ct[i,1]:6d} {ct[i,2]:6d} {ct[i,3]:6d} {ct[i,4]:6d}")
    # Chi-square test of independence
    n = ct.sum()
    row_sum = ct.sum(axis=1)
    col_sum = ct.sum(axis=0)
    expected = np.outer(row_sum, col_sum) / n
    chi2_indep = ((ct - expected) ** 2 / expected).sum()
    dof_indep = (5 - 1) * (5 - 1)
    print(f"chi^2 vs independence: {chi2_indep:.1f} on {dof_indep} dof")
    # Naive p-value from chi^2 cdf (Wilson-Hilferty approximation)
    # For our purposes just report the test statistic
    # Critical value at p=0.05 for dof=16 ≈ 26.3; at p=0.01 ≈ 32.0

    # ==================================================================
    # 4. LOCALIZATION IN r-SPACE
    # ==================================================================
    print("\n========== Localization in r-space ==========")
    bins = [(1.0, 1.5), (1.5, 2.0), (2.0, 2.5), (2.5, 3.0), (3.0, 5.0), (5.0, 100.0)]
    gk = gauss_kuzmin_pmf(max_m=10)
    print(f"  {'r-range':<12s}  {'n':>5s}  {'frac(a=1)':>9s}  {'mean ln a':>10s}  {'~ K(r)':>7s}")
    for lo, hi in bins:
        sel = (ratio_rs >= lo) & (ratio_rs < hi)
        if sel.sum() == 0:
            continue
        sub_aks = []
        for (host, r), lnK in zip(pairs, ratio_lnK):
            if lo <= r < hi:
                sub_aks.extend(per_ratio_stats(r, depth=10)[0])
        if sub_aks:
            frac_a1 = sum(1 for a in sub_aks if a == 1) / len(sub_aks)
            mean_ln = float(np.mean([math.log(a) for a in sub_aks if a > 0]))
            print(f"  [{lo:5.2f},{hi:5.2f})  {sel.sum():5d}  {frac_a1:9.4f}  {mean_ln:10.4f}  {math.exp(mean_ln):7.4f}")

    # ==================================================================
    # 5. PLOTS
    # ==================================================================
    fig, ax = plt.subplots(figsize=(8, 4.5))
    ax.hist(ratio_lnK, bins=40, color="#16a085", alpha=0.85, edgecolor="white", linewidth=0.5)
    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"empirical mean = {emp_mean:.3f} (z = {z:+.2f})")
    ax.set_xlabel(r"per-ratio empirical mean $\langle\ln a_k\rangle$  (depth-10 truncation)")
    ax.set_ylabel("count")
    ax.set_title("Khinchin-constant test on observed period ratios")
    ax.legend()
    fig.tight_layout()
    fig.savefig("c1_khinchin.png", dpi=140)

    # Levy histogram
    fig, ax = plt.subplots(figsize=(8, 4.5))
    ax.hist(ratio_levy, bins=40, color="#9b59b6", alpha=0.85, edgecolor="white", linewidth=0.5)
    ax.axvline(LEVY_LN, color="#c0392b", lw=2, label=f"Gauss-Kuzmin: L₀ = {LEVY_LN:.3f}")
    ax.axvline(math.log((1 + math.sqrt(5)) / 2), color="#f39c12", lw=2, ls=":",
               label=f"golden ratio: ln φ = {math.log((1+math.sqrt(5))/2):.3f}")
    ax.axvline(emp_L, color="#2c3e50", lw=2, ls="--",
               label=f"empirical mean = {emp_L:.3f} (z = {z_L:+.2f})")
    ax.set_xlabel(r"per-ratio Levy slope  $d \log q_n / dn$  (depth-10)")
    ax.set_ylabel("count")
    ax.set_xlim(0, 3)
    ax.set_title("Lévy-constant test on observed period ratios")
    ax.legend()
    fig.tight_layout()
    fig.savefig("c1_levy.png", dpi=140)

    # Conditional dependence heatmap (deviation from expected)
    deviation = (ct - expected) / np.sqrt(expected)
    fig, ax = plt.subplots(figsize=(6.5, 5))
    vmax = max(abs(deviation.min()), abs(deviation.max()))
    im = ax.imshow(deviation, cmap="RdBu_r", vmin=-vmax, vmax=vmax)
    ax.set_xticks(range(5)); ax.set_yticks(range(5))
    ax.set_xticklabels(["1", "2", "3", "4", "≥5"])
    ax.set_yticklabels(["1", "2", "3", "4", "≥5"])
    ax.set_xlabel(r"$a_{k+1}$")
    ax.set_ylabel(r"$a_k$")
    ax.set_title(f"Deviation of (a_k, a_{{k+1}}) from independence\n(χ²={chi2_indep:.1f} on {dof_indep} dof; critical@5%≈26.3)")
    for i in range(5):
        for j in range(5):
            ax.text(j, i, f"{deviation[i,j]:+.1f}",
                    ha="center", va="center", color="white" if abs(deviation[i,j]) > vmax*0.5 else "black", fontsize=9)
    plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04, label="(obs − exp) / √exp")
    fig.tight_layout()
    fig.savefig("c1_conditional.png", dpi=140)


if __name__ == "__main__":
    main()