"""
Lithwick-Wu pair test at 2nd-order mean-motion resonances.

Test 9 found a clear LW asymmetric pile-up at four first-order MMRs
(2:1, 3:2, 4:3, 5:4), with shared MLE ε = 0.017, σ = 0.004, γ = 5.

Does the same structure exist at *2nd-order* MMRs (p − q = 2)?

Physical expectation (Lithwick & Wu 2012; Murray–Dermott §8): second-order
resonances are weaker — libration width ∝ μ instead of μ^(2/3), so for
typical μ ~ 10^-4 the resonance width is smaller by factor ~ μ^(1/3) ~ 0.05.
So if LW is physical, we'd expect:
  - ε(2nd-order) substantially SMALLER than ε(1st-order)
  - σ(2nd-order) substantially SMALLER
  - γ(2nd-order) plausibly SMALLER (less trapping efficiency at higher order)

2nd-order MMR centers in Farey-20 (gcd=1, p > q ≥ 2, p − q = 2):
  5:3 → {r} = 2/3 ≈ 0.6667
  7:5 → {r} = 2/5 = 0.4
  9:7 → {r} = 2/7 ≈ 0.2857
  11:9 → {r} = 2/9 ≈ 0.2222

We exclude integer ratios (3:1, 5:1, …, also 2nd-order) because they
share {r} = 0 with 1st-order 2:1 and would mix the signals.
"""

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


LN2 = math.log(2.0)

# 2nd-order MMR centers (p-q = 2, q >= 2, gcd=1)
SECOND_ORDER_RES = [(5, 3), (7, 5), (9, 7), (11, 9)]
SECOND_ORDER_CENTERS = np.array([2/3, 2/5, 2/7, 2/9])
SECOND_ORDER_LABELS  = ["5:3", "7:5", "9:7", "11:9"]

# 1st-order for comparison
FIRST_ORDER_RES = [(2, 1), (3, 2), (4, 3), (5, 4)]
FIRST_ORDER_CENTERS = np.array([0.0, 0.5, 1.0/3.0, 0.25])
FIRST_ORDER_LABELS  = ["2:1", "3:2", "4:3", "5:4"]


def load_ratios(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)
    out = []
    for periods in systems.values():
        if len(periods) < 2:
            continue
        periods = sorted(periods)
        for i in range(len(periods) - 1):
            r = periods[i + 1] / periods[i]
            if r > 1 and math.isfinite(r):
                out.append(r)
    return np.array(out)


def gauss_density(x):
    return 1.0 / ((1.0 + x) * LN2)


def wrapped_gaussian_vec(x, mu, sigma, n_wraps=3):
    ks = np.arange(-n_wraps, n_wraps + 1)
    delta = x[:, None] - mu + ks[None, :]
    coef = 1.0 / (sigma * math.sqrt(2 * math.pi))
    return (coef * np.exp(-0.5 * (delta / sigma) ** 2)).sum(axis=1)


def lw_density_vec(x, centers, eps, sigma, gamma):
    K = len(centers)
    total = np.zeros_like(x)
    norm = 1.0 / (1.0 + gamma)
    for c in centers:
        total += norm * wrapped_gaussian_vec(x, c - eps, sigma)
        total += gamma * norm * wrapped_gaussian_vec(x, c + eps, sigma)
    return total / K


def lw_normalizer(centers, eps, sigma, gamma, n_grid=2001):
    xs = np.linspace(0, 1, n_grid)
    ys = lw_density_vec(xs, centers, eps, sigma, gamma)
    return float(np.trapezoid(ys, xs))


def fit_lw(fracs, centers, alpha_grid, eps_grid, sigma_grid, gamma_grid, gauss_base):
    """4-D grid search with shared γ; Z/f computed once per (eps, sigma, gamma)."""
    best = (-np.inf, None, None, None, None)
    for eps in eps_grid:
        for sigma in sigma_grid:
            for gamma in gamma_grid:
                f = lw_density_vec(fracs, centers, eps, sigma, gamma)
                Z = lw_normalizer(centers, eps, sigma, gamma)
                if Z <= 0:
                    continue
                f_norm = f / Z
                for alpha in alpha_grid:
                    p = alpha * gauss_base + (1 - alpha) * f_norm
                    if np.any(p <= 0):
                        continue
                    ll = float(np.log(p).sum())
                    if ll > best[0]:
                        best = (ll, alpha, eps, sigma, gamma)
    return best


def main():
    rs = load_ratios("ps_dm.csv")
    fracs = rs - np.floor(rs)
    fracs = fracs[(fracs >= 0) & (fracs < 1)]
    print(f"n fractional parts: {len(fracs)}")

    gauss_base = gauss_density(fracs)
    ll_pure_gauss = float(np.log(gauss_base).sum())
    print(f"pure Gauss LL: {ll_pure_gauss:.2f}")

    # ----------- count of pairs in narrow windows around each 2nd-order center
    print("\n--- narrow-window counts at 2nd-order centers (Test 8 analog) ---")
    print(f"  {'res':>5}  {'center':>7}  {'below':>6}  {'near':>5}  {'above':>6}  "
          f"{'b/a':>6}  {'z(b−a)':>9}")
    for (p, q), c in zip(SECOND_ORDER_RES, SECOND_ORDER_CENTERS):
        # ±5% window, split at ±0.02
        d = ((fracs - c + 0.5) % 1) - 0.5
        below = (d < -0.02) & (d > -0.05)
        near = np.abs(d) <= 0.02
        above = (d > 0.02) & (d < 0.05)
        nb = int(below.sum()); nn = int(near.sum()); na = int(above.sum())
        ba = nb / na if na > 0 else float("nan")
        ntot = nb + na
        z = (nb - ntot * 0.5) / math.sqrt(max(ntot * 0.25, 1)) if ntot > 0 else 0
        print(f"  {p}:{q:>3}  {c:7.4f}  {nb:6d}  {nn:5d}  {na:6d}  {ba:6.2f}  {z:+9.2f}")

    # ----------- shared (ε, σ, γ) fit at 2nd-order centers
    print("\n--- LW pair fit at 2nd-order centers (shared ε, σ, γ) ---")
    alpha_grid = np.linspace(0.80, 0.99, 12)
    # Narrower grid for 2nd-order: smaller ε, smaller σ
    eps_grid = np.linspace(0.002, 0.04, 22)
    sigma_grid = np.exp(np.linspace(math.log(0.002), math.log(0.020), 16))
    gamma_grid = np.exp(np.linspace(math.log(0.2), math.log(5.0), 21))

    ll_2nd, a_2nd, eps_2nd, sigma_2nd, gamma_2nd = fit_lw(
        fracs, SECOND_ORDER_CENTERS, alpha_grid, eps_grid, sigma_grid, gamma_grid, gauss_base
    )
    print(f"2nd-order MLE: α={a_2nd:.3f}  ε={eps_2nd:.4f}  σ={sigma_2nd:.4f}  γ={gamma_2nd:.3f}")
    print(f"LL = {ll_2nd:.2f}  ΔLL vs pure Gauss = {ll_2nd - ll_pure_gauss:+.2f}")

    # Symmetric (γ=1) at same (α, ε, σ)
    f_sym = lw_density_vec(fracs, SECOND_ORDER_CENTERS, eps_2nd, sigma_2nd, 1.0)
    Z_sym = lw_normalizer(SECOND_ORDER_CENTERS, eps_2nd, sigma_2nd, 1.0)
    f_sym_norm = f_sym / Z_sym
    p_sym = a_2nd * gauss_base + (1 - a_2nd) * f_sym_norm
    ll_sym = float(np.log(p_sym).sum()) if np.all(p_sym > 0) else -np.inf
    print(f"\nSymmetric (γ=1) at same (α, ε, σ):  LL = {ll_sym:.2f}")
    print(f"ΔLL_asym = {ll_2nd - ll_sym:+.2f}  (chi² on 1 dof, critical@0.05 = 3.84)")
    print(f"z_eff for asymmetry: {math.sqrt(2 * max(0, ll_2nd - ll_sym)):.2f}")

    # ----------- comparison to 1st-order (run same fit, for reference)
    print("\n--- LW pair fit at 1st-order centers (for direct comparison) ---")
    # Use the same grid scale for fair comparison
    ll_1st, a_1st, eps_1st, sigma_1st, gamma_1st = fit_lw(
        fracs, FIRST_ORDER_CENTERS, alpha_grid, eps_grid, sigma_grid, gamma_grid, gauss_base
    )
    print(f"1st-order MLE: α={a_1st:.3f}  ε={eps_1st:.4f}  σ={sigma_1st:.4f}  γ={gamma_1st:.3f}")
    print(f"LL = {ll_1st:.2f}  ΔLL vs pure Gauss = {ll_1st - ll_pure_gauss:+.2f}")

    # ----------- side-by-side
    print("\n=== summary: 1st-order vs 2nd-order LW pair fit ===")
    print(f"  {'param':>8} {'1st-order':>12} {'2nd-order':>12}  {'ratio':>7}")
    print(f"  {'α':>8} {a_1st:12.3f} {a_2nd:12.3f}  {'':>7}")
    print(f"  {'ε':>8} {eps_1st:12.4f} {eps_2nd:12.4f}  {eps_2nd/eps_1st if eps_1st>0 else float('nan'):7.2f}")
    print(f"  {'σ':>8} {sigma_1st:12.4f} {sigma_2nd:12.4f}  {sigma_2nd/sigma_1st if sigma_1st>0 else float('nan'):7.2f}")
    print(f"  {'γ':>8} {gamma_1st:12.3f} {gamma_2nd:12.3f}  {gamma_2nd/gamma_1st if gamma_1st>0 else float('nan'):7.2f}")
    print(f"  {'ΔLL':>8} {ll_1st-ll_pure_gauss:+12.2f} {ll_2nd-ll_pure_gauss:+12.2f}  {'':>7}")

    # ----------- profile likelihood for γ at 2nd-order
    print("\n--- profile likelihood for γ at 2nd-order (other params at MLE) ---")
    gamma_profile = np.exp(np.linspace(math.log(0.2), math.log(5.0), 81))
    lls_gamma = []
    for g in gamma_profile:
        f = lw_density_vec(fracs, SECOND_ORDER_CENTERS, eps_2nd, sigma_2nd, g)
        Z = lw_normalizer(SECOND_ORDER_CENTERS, eps_2nd, sigma_2nd, g)
        if Z <= 0:
            lls_gamma.append(-np.inf); continue
        f_norm = f / Z
        p = a_2nd * gauss_base + (1 - a_2nd) * f_norm
        if np.any(p <= 0):
            lls_gamma.append(-np.inf); continue
        lls_gamma.append(float(np.log(p).sum()))
    lls_gamma = np.array(lls_gamma)
    ll_max = lls_gamma.max()
    g_mle = gamma_profile[np.argmax(lls_gamma)]
    in_ci68 = gamma_profile[lls_gamma > ll_max - 0.5]
    in_ci95 = gamma_profile[lls_gamma > ll_max - 1.92]
    print(f"γ MLE: {g_mle:.3f}")
    print(f"68% CI: [{in_ci68.min():.3f}, {in_ci68.max():.3f}]")
    print(f"95% CI: [{in_ci95.min():.3f}, {in_ci95.max():.3f}]")
    # Test against γ = 1
    f_at_1 = lw_density_vec(fracs, SECOND_ORDER_CENTERS, eps_2nd, sigma_2nd, 1.0)
    Z_at_1 = lw_normalizer(SECOND_ORDER_CENTERS, eps_2nd, sigma_2nd, 1.0)
    p_at_1 = a_2nd * gauss_base + (1 - a_2nd) * (f_at_1 / Z_at_1)
    ll_at_1 = float(np.log(p_at_1).sum()) if np.all(p_at_1 > 0) else -np.inf
    print(f"LL at γ=1: {ll_at_1:.2f}, ΔLL(MLE − γ=1) = {ll_max - ll_at_1:+.2f}")
    print(f"z_eff vs γ=1: {math.sqrt(2 * max(0, ll_max - ll_at_1)):.2f}")

    # ----------- per-resonance γ scan at 2nd-order (fixed ε, σ)
    print("\n--- per-resonance γ at 2nd-order (others = 1) ---")
    g_scan = np.exp(np.linspace(math.log(0.2), math.log(5.0), 25))
    for i, (res, c) in enumerate(zip(SECOND_ORDER_RES, SECOND_ORDER_CENTERS)):
        best_g = None
        best_ll = -np.inf
        for g in g_scan:
            gammas = np.ones(len(SECOND_ORDER_CENTERS))
            gammas[i] = g
            # Reuse lw_density_vec with shared γ — but we have per-center γ now.
            # Inline: compute density manually with per-center γ:
            total = np.zeros_like(fracs)
            for j, (cj, gj) in enumerate(zip(SECOND_ORDER_CENTERS, gammas)):
                norm = 1.0 / (1.0 + gj)
                total += norm * wrapped_gaussian_vec(fracs, cj - eps_2nd, sigma_2nd)
                total += gj * norm * wrapped_gaussian_vec(fracs, cj + eps_2nd, sigma_2nd)
            f = total / len(SECOND_ORDER_CENTERS)
            # Normalize numerically
            xs_n = np.linspace(0, 1, 2001)
            ys_n = np.zeros_like(xs_n)
            for j, (cj, gj) in enumerate(zip(SECOND_ORDER_CENTERS, gammas)):
                norm = 1.0 / (1.0 + gj)
                ys_n += norm * wrapped_gaussian_vec(xs_n, cj - eps_2nd, sigma_2nd)
                ys_n += gj * norm * wrapped_gaussian_vec(xs_n, cj + eps_2nd, sigma_2nd)
            ys_n /= len(SECOND_ORDER_CENTERS)
            Z = float(np.trapezoid(ys_n, xs_n))
            if Z <= 0:
                continue
            f_norm = f / Z
            p = a_2nd * gauss_base + (1 - a_2nd) * f_norm
            if np.any(p <= 0):
                continue
            ll = float(np.log(p).sum())
            if ll > best_ll:
                best_ll = ll
                best_g = g
        bg_str = f"{best_g:.3f}" if best_g is not None else "nan"
        print(f"  {res[0]}:{res[1]:<3} best γ = {bg_str}  LL = {best_ll:.2f}")

    # ----------- plots
    xs = np.linspace(0.0001, 0.9999, 2000)
    ys_gauss = gauss_density(xs)
    ys_lw_2nd = lw_density_vec(xs, SECOND_ORDER_CENTERS, eps_2nd, sigma_2nd, gamma_2nd)
    Z2 = lw_normalizer(SECOND_ORDER_CENTERS, eps_2nd, sigma_2nd, gamma_2nd)
    ys_lw_2nd /= Z2
    ys_lw_2nd_sym = lw_density_vec(xs, SECOND_ORDER_CENTERS, eps_2nd, sigma_2nd, 1.0)
    Z2sym = lw_normalizer(SECOND_ORDER_CENTERS, eps_2nd, sigma_2nd, 1.0)
    ys_lw_2nd_sym /= Z2sym
    ys_mix_2nd = a_2nd * ys_gauss + (1 - a_2nd) * ys_lw_2nd

    fig, ax = plt.subplots(figsize=(10, 5.5))
    ax.hist(fracs, bins=80, density=True, color="#7f8c8d", alpha=0.4,
            edgecolor="white", linewidth=0.5, label=f"data (n={len(fracs)})")
    ax.plot(xs, ys_gauss, color="#2c3e50", lw=1.2, alpha=0.5, label="pure Gauss")
    ax.plot(xs, ys_lw_2nd_sym, color="#8e44ad", lw=1.2, ls="--", alpha=0.6,
            label=fr"sym 2nd-order LW (γ=1, ε={eps_2nd:.3f}, σ={sigma_2nd:.3f})")
    ax.plot(xs, ys_lw_2nd, color="#c0392b", lw=1.5, alpha=0.85,
            label=fr"2nd-order LW (γ={gamma_2nd:.2f})")
    ax.plot(xs, ys_mix_2nd, color="#16a085", lw=2.5,
            label=fr"best fit:  α={a_2nd:.3f}, ΔLL = {ll_2nd - ll_pure_gauss:.1f}")
    for c, lbl in zip(SECOND_ORDER_CENTERS, SECOND_ORDER_LABELS):
        ax.axvline(c, color="#c0392b", lw=0.7, alpha=0.5)
        ax.text(c + 0.005, ax.get_ylim()[1] * 0.93, lbl, fontsize=8, color="#c0392b")
    ax.set_xlabel(r"$\{r\}$")
    ax.set_ylabel("density")
    ax.set_title(rf"LW pair fit at 2nd-order MMRs:  "
                 rf"$\alpha={a_2nd:.3f}$, $\varepsilon={eps_2nd:.4f}$, "
                 rf"$\sigma={sigma_2nd:.4f}$, $\gamma={gamma_2nd:.2f}$")
    ax.legend(loc="upper left", fontsize=8)
    ax.set_xlim(0, 1)
    fig.tight_layout()
    fig.savefig("c1_lw_2nd_fit.png", dpi=140)

    # Comparison plot: ε(1st) vs ε(2nd), and other params
    fig, axes = plt.subplots(1, 4, figsize=(13, 4), sharey=False)
    params_1st = [a_1st, eps_1st, sigma_1st, gamma_1st]
    params_2nd = [a_2nd, eps_2nd, sigma_2nd, gamma_2nd]
    names = [r"$\alpha$", r"$\varepsilon$", r"$\sigma$", r"$\gamma$"]
    for ax, n, p1, p2 in zip(axes, names, params_1st, params_2nd):
        ax.bar([0, 1], [p1, p2], color=["#c0392b", "#9b59b6"])
        ax.set_xticks([0, 1])
        ax.set_xticklabels(["1st", "2nd"])
        ax.set_title(n, fontsize=14)
        ax.annotate(f"{p1:.3g}", (0, p1), ha="center", va="bottom", fontsize=10)
        ax.annotate(f"{p2:.3g}", (1, p2), ha="center", va="bottom", fontsize=10)
    fig.suptitle("LW pair parameters: 1st-order vs 2nd-order MMRs")
    fig.tight_layout()
    fig.savefig("c1_lw_order_compare.png", dpi=140)

    print("\nplots: c1_lw_2nd_fit.png  c1_lw_order_compare.png")


if __name__ == "__main__":
    main()
