"""
Conjecture II, extended analysis.

(A) Larger Farey alphabets. The original Farey-9 has 13 symbols.  We now
    repeat the chain-type count with Farey-12 (≥25 symbols) and Farey-15
    (≥35 symbols) — increasing resolution near r = 1 and r = 2.

(B) Markov-chain test on 3-planet systems.  Each system has two chain
    symbols (s₁, s₂).  Under "independent" formation/filtering the joint
    distribution of (s₁, s₂) factorises.  Test it against a chi-square
    of independence on the contingency table.  A non-trivial coupling
    is direct evidence that consecutive resonances in a chain are
    correlated — a stronger statement than just "chain types
    concentrate."

(C) Entropy comparison.  Compute Shannon entropy of the chain-type
    distribution and compare to the uniform entropy log|Σ|^L.  This is
    a single-number summary of "lattice rigidity."
"""

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


def farey_pairs(max_sum):
    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):
    best = min(pairs, key=lambda pq: abs(r - pq[0] / pq[1]))
    return best, abs(r - best[0] / best[1])


def load_systems(path):
    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)
    return {h: sorted(ps) for h, ps in systems.items() if len(ps) >= 2}


def shannon_entropy(counts):
    """Shannon entropy of a count distribution (in bits)."""
    n = sum(counts)
    if n == 0:
        return 0.0
    H = 0.0
    for c in counts:
        if c > 0:
            p = c / n
            H -= p * math.log2(p)
    return H


def coupon_expected_distinct(K, N):
    """Expected number of distinct cells when N balls fall uniformly into K bins."""
    return K * (1 - (1 - 1/K) ** N) if K > 0 else 0


def main():
    systems = load_systems("ps_dm.csv")
    threeplus = {h: ps for h, ps in systems.items() if len(ps) >= 3}
    print(f"multi-planet systems: {len(systems)}  3+ planet: {len(threeplus)}")

    for MAX_SUM in (9, 12, 15):
        print("\n" + "=" * 72)
        print(f"Farey-{MAX_SUM} alphabet")
        print("=" * 72)
        alphabet = farey_pairs(MAX_SUM)
        K = len(alphabet)
        print(f"|Σ| = {K} symbols")

        # Encode each adjacent ratio and assemble chains
        chain_for_host = {}
        all_symbols = []
        for host, periods in threeplus.items():
            chain = []
            for i in range(len(periods) - 1):
                r = periods[i + 1] / periods[i]
                (p, q), _ = nearest_resonance(r, alphabet)
                chain.append((p, q))
            chain_for_host[host] = tuple(chain)
            all_symbols.extend(chain)

        # ---------- by chain length ----------
        print(f"\n{'L':>3}  {'systems':>9}  {'distinct':>9}  {'possible':>15}  {'uniform_exp':>13}  {'entropy(bits)':>14}  {'uniform_entropy':>15}")
        for L in (2, 3, 4, 5):
            chains_at_L = [c for c in chain_for_host.values() if len(c) == L]
            if not chains_at_L:
                continue
            counter = Counter(chains_at_L)
            n_sys = len(chains_at_L)
            n_distinct = len(counter)
            possible = K ** L
            uniform_exp = coupon_expected_distinct(possible, n_sys)
            H = shannon_entropy(list(counter.values()))
            H_uniform = math.log2(min(possible, n_sys))  # cap at log2(n) since you can't have more distinct than n
            print(f"  {L}  {n_sys:9d}  {n_distinct:9d}  {possible:15,}  {uniform_exp:13.1f}  {H:14.3f}  {H_uniform:15.3f}")

        # ---------- Top chain types at L=2 ----------
        chains2 = [c for c in chain_for_host.values() if len(c) == 2]
        counter2 = Counter(chains2)
        print(f"\nTop 12 L=2 chain types (Farey-{MAX_SUM}, n={len(chains2)} 3-planet systems):")
        for (a, b), c in counter2.most_common(12):
            print(f"  {c:4d}  {a[0]:2d}:{a[1]} -- {b[0]:2d}:{b[1]}")

        # ---------- Markov-chain test on L=2 ----------
        # Build K x K contingency table on (s1, s2)
        idx = {s: i for i, s in enumerate(alphabet)}
        ct = np.zeros((K, K), dtype=int)
        for c in chains2:
            ct[idx[c[0]], idx[c[1]]] += 1
        # Aggregate independence chi-square:
        # combine rows/cols with <5 expected counts into "other" bins so chi^2 is valid
        N = ct.sum()
        rs = ct.sum(axis=1)
        cs = ct.sum(axis=0)
        expected = np.outer(rs, cs) / N if N else None

        # Quick top-cell deviation report
        print(f"\nMarkov-chain independence test (L=2, {N} systems, full {K}x{K} table)")
        chi2_full = (((ct - expected) ** 2) / np.where(expected > 0, expected, 1)).sum()
        dof_full = (K - 1) ** 2
        print(f"  chi^2 = {chi2_full:.1f} on {dof_full} dof (raw, no cell merging)")

        # Use a robust version: merge rare symbols into "other"
        used = (rs >= 5) & (cs >= 5)
        keep = [i for i, u in enumerate(used) if u]
        if len(keep) >= 3:
            small = [i for i in range(K) if i not in keep]
            ct_red = np.zeros((len(keep) + 1, len(keep) + 1), dtype=int)
            keep_idx = {i: k for k, i in enumerate(keep)}
            for i in range(K):
                for j in range(K):
                    ii = keep_idx.get(i, len(keep))
                    jj = keep_idx.get(j, len(keep))
                    ct_red[ii, jj] += ct[i, j]
            rs_r = ct_red.sum(axis=1)
            cs_r = ct_red.sum(axis=0)
            N_r = ct_red.sum()
            exp_r = np.outer(rs_r, cs_r) / N_r
            chi2_r = (((ct_red - exp_r) ** 2) / np.where(exp_r > 0, exp_r, 1)).sum()
            dof_r = (ct_red.shape[0] - 1) ** 2
            print(f"  After merging rare symbols ({len(keep)} kept + 1 'other'):")
            print(f"  chi^2 = {chi2_r:.1f} on {dof_r} dof")
            print(f"  critical at p=0.01 ≈ chi^2_dof_inv(0.99, dof) — compare to {dof_r * 2:.0f}-{dof_r * 3:.0f}-ish")

        # ---------- Plot: contingency heatmap for largest alphabet ----------
        if MAX_SUM == 12:
            # Use the most-common 12 symbols only for legibility
            top_k = [s for s, _ in Counter(all_symbols).most_common(12)]
            tk_idx = {s: i for i, s in enumerate(top_k)}
            ct_top = np.zeros((12, 12), dtype=int)
            for c in chains2:
                if c[0] in tk_idx and c[1] in tk_idx:
                    ct_top[tk_idx[c[0]], tk_idx[c[1]]] += 1
            rs_t = ct_top.sum(axis=1)
            cs_t = ct_top.sum(axis=0)
            N_t = ct_top.sum()
            if N_t > 0:
                exp_t = np.outer(rs_t, cs_t) / N_t
                with np.errstate(divide="ignore", invalid="ignore"):
                    dev = np.where(exp_t > 0, (ct_top - exp_t) / np.sqrt(exp_t), 0.0)
                vmax = max(abs(dev.min()), abs(dev.max()), 1)
                fig, ax = plt.subplots(figsize=(8.5, 7.5))
                im = ax.imshow(dev, cmap="RdBu_r", vmin=-vmax, vmax=vmax)
                labels = [f"{p}:{q}" for p, q in top_k]
                ax.set_xticks(range(12)); ax.set_yticks(range(12))
                ax.set_xticklabels(labels, rotation=45, ha="right")
                ax.set_yticklabels(labels)
                ax.set_xlabel(r"second resonance $s_2$")
                ax.set_ylabel(r"first resonance $s_1$")
                ax.set_title(f"Conjecture II: deviation from independence in 3-planet chain symbols\n(Farey-12, top-12 symbols, n={N_t} systems)")
                for i in range(12):
                    for j in range(12):
                        if abs(dev[i, j]) > vmax * 0.35:
                            ax.text(j, i, f"{dev[i,j]:+.1f}",
                                    ha="center", va="center",
                                    color="white" if abs(dev[i,j]) > vmax*0.6 else "black",
                                    fontsize=8)
                plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04, label="(obs − exp) / √exp")
                fig.tight_layout()
                fig.savefig("c2_markov_heatmap.png", dpi=140)


if __name__ == "__main__":
    main()
