"""
Mixture fit v3 (vectorized).

Same model as v2 but rewritten with numpy broadcasting so the (α, σ)
grid search completes in seconds instead of minutes. Also adds an
alphabet ablation across Farey-9, 12, 15, 20.
"""

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


LN2 = math.log(2.0)


def farey_pairs(max_sum):
    out = []
    for q in range(1, max_sum):
        for p in range(q + 1, max_sum - q + 1):
            if math.gcd(p, q) == 1:
                out.append((p, q))
    return out


def build_centers(max_sum, include_integer=True):
    counts = Counter()
    for p, q in farey_pairs(max_sum):
        if q > 1:
            counts[round((p % q) / q, 9)] += 1
    if include_integer:
        n_int = max_sum - 2
        if n_int > 0:
            counts[0.0] = n_int
    items = sorted(counts.items())
    centers = np.array([c for c, _ in items])
    weights = np.array([w for _, w in items], dtype=float)
    weights /= weights.sum()
    return centers, weights


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 farey_density_vec(x, centers, weights, sigma, n_wraps=3):
    """Vectorized over x (1-D array). Returns p(x) (1-D, length N)."""
    ks = np.arange(-n_wraps, n_wraps + 1)
    delta = x[:, None, None] - centers[None, :, None] + ks[None, None, :]
    coef = 1.0 / (sigma * math.sqrt(2 * math.pi))
    kernel = coef * np.exp(-0.5 * (delta / sigma) ** 2)
    per_center = kernel.sum(axis=2)
    return (per_center * weights[None, :]).sum(axis=1)


def farey_normalizer(centers, weights, sigma, n_grid=2001):
    xs = np.linspace(0, 1, n_grid)
    ys = farey_density_vec(xs, centers, weights, sigma)
    return np.trapezoid(ys, xs)


def fit_mixture(fracs, centers, weights, alpha_grid, sigma_grid, baseline="gauss"):
    base = gauss_density(fracs) if baseline == "gauss" else np.ones_like(fracs)
    ll_grid = np.full((len(alpha_grid), len(sigma_grid)), -np.inf)
    for j, sigma in enumerate(sigma_grid):
        f = farey_density_vec(fracs, centers, weights, sigma)
        Z = farey_normalizer(centers, weights, sigma)
        if Z <= 0:
            continue
        f_norm = f / Z
        for i, alpha in enumerate(alpha_grid):
            p = alpha * base + (1 - alpha) * f_norm
            with np.errstate(divide="ignore", invalid="ignore"):
                ll = np.log(p).sum()
            ll_grid[i, j] = ll
    best_idx = np.unravel_index(np.argmax(ll_grid), ll_grid.shape)
    return ll_grid, best_idx, ll_grid[best_idx], alpha_grid[best_idx[0]], sigma_grid[best_idx[1]]


def chi2_binned(fracs, n_bins=20):
    counts, edges = np.histogram(fracs, bins=n_bins, range=(0, 1))
    n = counts.sum()
    gauss_exp = np.array([
        n * (math.log2((1 + edges[i+1]) / (1 + edges[i])))
        for i in range(n_bins)
    ])
    unif_exp = np.full(n_bins, n / n_bins)
    chi2_g = ((counts - gauss_exp) ** 2 / gauss_exp).sum()
    chi2_u = ((counts - unif_exp) ** 2 / unif_exp).sum()
    return counts, edges, gauss_exp, unif_exp, chi2_g, chi2_u


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

    # ---------- non-parametric chi-square ----------
    counts, edges, gexp, uexp, chi2_g, chi2_u = chi2_binned(fracs, n_bins=20)
    print(f"\n20-bin chi-square:")
    print(f"  chi² vs Gauss measure: {chi2_g:.1f} on 19 dof")
    print(f"  chi² vs uniform:       {chi2_u:.1f} on 19 dof")
    print(f"  critical at p=0.001 ≈ 43.8")

    dev = (counts - gexp) / np.sqrt(gexp)
    print(f"\nlargest deviations from Gauss:")
    top = np.argsort(-np.abs(dev))[:5]
    for k in top:
        lo, hi = edges[k], edges[k+1]
        print(f"  [{lo:.3f}, {hi:.3f}): obs={counts[k]}  exp={gexp[k]:.1f}  dev={dev[k]:+.2f}")

    # ---------- alphabet ablation ----------
    print("\n" + "=" * 60)
    print("Alphabet ablation")
    print("=" * 60)

    alpha_grid = np.linspace(0.0, 1.0, 51)
    sigma_grid = np.exp(np.linspace(math.log(0.001), math.log(0.5), 41))

    base_vals = gauss_density(fracs)
    ll_pure_gauss = np.log(base_vals).sum()
    print(f"\npure Gauss LL: {ll_pure_gauss:.2f}")
    print(f"pure uniform LL: 0.00")
    print(f"Δ(pure Gauss − pure uniform) = {ll_pure_gauss:.2f}")

    ablation = {}
    for max_sum in (9, 12, 15, 20):
        centers, weights = build_centers(max_sum=max_sum, include_integer=True)
        ll_grid, idx, ll_best, a_best, s_best = fit_mixture(
            fracs, centers, weights, alpha_grid, sigma_grid, baseline="gauss"
        )
        dll = ll_best - ll_pure_gauss
        ablation[max_sum] = (centers, weights, ll_grid, a_best, s_best, ll_best, dll)
        print(f"  Farey-{max_sum:>2}: |Σ|={len(centers):2d}  α={a_best:.3f}  σ={s_best:.4f}  "
              f"LL={ll_best:.2f}  ΔLL = {dll:+.2f}  (2ΔLL = {2*dll:.1f})")

    centers, weights, ll_grid, a_g, s_g, ll_best_g, dll_g = ablation[12]

    # ---------- uniform baseline ----------
    ll_grid_u, _, ll_best_u, a_u, s_u = fit_mixture(
        fracs, centers, weights, alpha_grid, sigma_grid, baseline="uniform"
    )
    print(f"\nFarey-12 vs uniform baseline: α={a_u:.3f}  σ={s_u:.4f}  LL={ll_best_u:.2f}  "
          f"ΔLL vs pure unif = {ll_best_u:+.2f}")
    print(f"Δ(Gauss+Farey best − unif+Farey best) = {ll_best_g - ll_best_u:+.2f}")

    # ---------- profile likelihood on α (Farey-12, Gauss baseline) ----------
    a_profile = np.linspace(0.0, 1.0, 101)
    f_at_sg = farey_density_vec(fracs, centers, weights, s_g)
    Z_at_sg = farey_normalizer(centers, weights, s_g)
    fn = f_at_sg / Z_at_sg
    lls_a = np.array([np.log(a * base_vals + (1 - a) * fn).sum() for a in a_profile])

    in_ci = a_profile[lls_a > ll_best_g - 0.5]
    in95 = a_profile[lls_a > ll_best_g - 1.92]
    a_ci = (in_ci.min(), in_ci.max()) if len(in_ci) else (a_g, a_g)
    a_95 = (in95.min(), in95.max()) if len(in95) else (a_g, a_g)
    print(f"\nProfile CI on α (Farey-12):")
    print(f"  68%: [{a_ci[0]:.3f}, {a_ci[1]:.3f}]   Farey weight: [{1-a_ci[1]:.3f}, {1-a_ci[0]:.3f}]")
    print(f"  95%: [{a_95[0]:.3f}, {a_95[1]:.3f}]   Farey weight: [{1-a_95[1]:.3f}, {1-a_95[0]:.3f}]")

    # ---------- plots ----------
    xs = np.linspace(0.0001, 0.9999, 2000)
    ys_gauss = gauss_density(xs)
    ys_farey = farey_density_vec(xs, centers, weights, s_g) / Z_at_sg
    ys_mix = a_g * ys_gauss + (1 - a_g) * ys_farey

    fig, ax = plt.subplots(figsize=(9, 5))
    ax.hist(fracs, bins=80, density=True, color="#7f8c8d", alpha=0.45,
            edgecolor="white", linewidth=0.5, label=f"data (n={len(fracs)})")
    ax.plot(xs, ys_gauss, color="#2c3e50", lw=2, label="pure Gauss measure")
    ax.plot(xs, ys_farey, color="#c0392b", lw=1.5, alpha=0.7,
            label=f"Farey-12 component (σ={s_g:.4f})")
    ax.plot(xs, ys_mix, color="#16a085", lw=2.5,
            label=fr"best mixture:  α = {a_g:.3f}")
    for c, w in zip(centers, weights):
        ax.axvline(c, color="#c0392b", lw=0.5 + w * 4, alpha=0.3)
    ax.set_xlabel(r"$\{r\} = r - \lfloor r \rfloor$")
    ax.set_ylabel("density")
    ax.set_title(rf"Period-ratio mixture fit (Farey-12):  Farey weight $1-\alpha = {1-a_g:.3f}$,  $\Delta$LL = {dll_g:.1f}")
    ax.legend(loc="upper right", fontsize=9)
    ax.set_xlim(0, 1)
    fig.tight_layout()
    fig.savefig("c1_mixture_fit.png", dpi=140)

    # Profile plot
    fig, ax = plt.subplots(figsize=(8, 4.5))
    ax.plot(a_profile, lls_a, color="#2c3e50", lw=2)
    ax.axvline(a_g, color="#c0392b", lw=1.5, ls="--", label=f"MLE α = {a_g:.3f}")
    ax.axhline(ll_best_g - 0.5, color="#888", lw=1, ls=":", label="68% level")
    ax.axhline(ll_best_g - 1.92, color="#bbb", lw=1, ls=":", label="95% level")
    ax.axvspan(a_ci[0], a_ci[1], color="#16a085", alpha=0.20,
               label=f"68%: [{a_ci[0]:.3f}, {a_ci[1]:.3f}]")
    ax.axvspan(a_95[0], a_95[1], color="#16a085", alpha=0.08,
               label=f"95%: [{a_95[0]:.3f}, {a_95[1]:.3f}]")
    ax.set_xlabel(r"$\alpha$  (Gauss component weight)")
    ax.set_ylabel("log-likelihood")
    ax.set_title(rf"Profile log-likelihood for $\alpha$  (σ = {s_g:.4f})")
    ax.legend(fontsize=8)
    fig.tight_layout()
    fig.savefig("c1_mixture_profile.png", dpi=140)

    # 2D contour
    AA, SS = np.meshgrid(alpha_grid, sigma_grid, indexing="ij")
    dll_arr = ll_grid - ll_best_g
    fig, ax = plt.subplots(figsize=(8, 5))
    levels = [-30, -9, -4.6, -1.92, -0.5, 0]
    cs = ax.contourf(AA, SS, dll_arr, levels=levels, cmap="viridis")
    ax.set_yscale("log")
    ax.set_xlabel(r"$\alpha$  (Gauss component weight)")
    ax.set_ylabel(r"$\sigma$  (Farey bandwidth)")
    ax.set_title(r"2-D likelihood contours,  $\Delta$LL relative to best")
    ax.plot([a_g], [s_g], "r*", markersize=14,
            label=f"MLE  α={a_g:.3f}  σ={s_g:.4f}")
    plt.colorbar(cs, ax=ax, label=r"$\Delta$LL")
    ax.legend()
    fig.tight_layout()
    fig.savefig("c1_mixture_2d.png", dpi=140)

    # Alphabet ablation
    fig, ax = plt.subplots(figsize=(8, 4.5))
    ax2 = ax.twinx()
    ks = [9, 12, 15, 20]
    sizes = [len(ablation[k][0]) for k in ks]
    a_vals = [ablation[k][3] for k in ks]
    dlls = [ablation[k][6] for k in ks]
    ax.plot(ks, [1 - a for a in a_vals], "o-", color="#c0392b", lw=2,
            label="Farey fraction (1 − α)")
    ax2.plot(ks, dlls, "s-", color="#2c3e50", lw=2,
             label="ΔLL vs pure Gauss")
    for k, n in zip(ks, sizes):
        ax.annotate(f"|Σ|={n}", (k, 1 - dict(zip(ks, a_vals))[k]),
                    textcoords="offset points", xytext=(8, 5), fontsize=9, color="#c0392b")
    ax.set_xlabel("Farey alphabet bound  (p + q ≤ N)")
    ax.set_ylabel("Farey fraction (1 − α)", color="#c0392b")
    ax2.set_ylabel(r"$\Delta$LL", color="#2c3e50")
    ax.set_xticks(ks)
    ax.set_title("Alphabet ablation: how the mixture fit depends on Farey alphabet")
    fig.tight_layout()
    fig.savefig("c1_mixture_ablation.png", dpi=140)

    print("\nplots: c1_mixture_fit.png  c1_mixture_profile.png  c1_mixture_2d.png  c1_mixture_ablation.png")

    # ============================================================
    # NARROW-WINDOW PILE-UP TEST
    # ============================================================
    # The mixture-model σ hits the upper boundary at every alphabet size,
    # which means a simple "Gauss + sharp Farey peaks" model isn't the
    # right object. What IS robust is: in a narrow window around each
    # low-order Farey fraction, how much observed mass is there vs the
    # Gauss-measure expectation?
    print("\n" + "=" * 60)
    print("Narrow-window pile-up at low-order Farey centers (±2% window)")
    print("=" * 60)
    WIN = 0.02
    centers_low, _ = build_centers(max_sum=9, include_integer=True)
    centers_to_pq = {}
    for p, q in farey_pairs(9):
        if q > 1:
            fr = round((p % q) / q, 9)
            for i, c in enumerate(centers_low):
                if abs(c - fr) < 1e-9:
                    if i not in centers_to_pq or (p + q) < (centers_to_pq[i][0] + centers_to_pq[i][1]):
                        centers_to_pq[i] = (p, q)
                    break
    centers_to_pq[0] = ("integer", 1)
    print(f"\n  {'center':>8}  {'label':>10}  {'obs in window':>14}  {'Gauss-exp':>10}  "
          f"{'ratio':>6}  {'z':>5}")
    pile_up_table = []
    for idx, c in enumerate(centers_low):
        # circular distance
        in_win = np.sum(np.abs(((fracs - c + 0.5) % 1) - 0.5) < WIN)
        # Gauss-measure mass within ±WIN (handle wrap at 0)
        lo, hi = c - WIN, c + WIN
        gauss_mass = 0.0
        if lo < 0:
            gauss_mass += math.log2((1 + 1) / (1 + 1 + lo))  # wraps to [1+lo, 1)
            lo = 0
        if hi > 1:
            gauss_mass += math.log2((1 + hi - 1) / (1 + 0))   # wraps to [0, hi-1)
            hi = 1
        gauss_mass += math.log2((1 + hi) / (1 + lo))
        expected = gauss_mass * len(fracs)
        ratio = in_win / expected if expected > 0 else float("nan")
        z = (in_win - expected) / math.sqrt(expected) if expected > 0 else 0
        pq = centers_to_pq.get(idx, ("?", "?"))
        label = "integer (2:1, 3:1, …)" if pq[0] == "integer" else f"{pq[0]}:{pq[1]}"
        print(f"  {c:8.4f}  {label:>10}  {in_win:14d}  {expected:10.1f}  {ratio:6.2f}  {z:+5.2f}")
        pile_up_table.append((c, label, in_win, expected, ratio, z))

    # Plot the pile-up bar chart
    fig, ax = plt.subplots(figsize=(9, 4.5))
    labels = [row[1] for row in pile_up_table]
    obs = [row[2] for row in pile_up_table]
    exp = [row[3] for row in pile_up_table]
    x_pos = np.arange(len(labels))
    w = 0.4
    ax.bar(x_pos - w/2, obs, w, color="#c0392b", label="observed (within ±2% of center)")
    ax.bar(x_pos + w/2, exp, w, color="#2c3e50", alpha=0.7, label="Gauss-measure expected")
    for i, (_, _, o, e, _, z) in enumerate(pile_up_table):
        if abs(z) > 1:
            ax.annotate(f"z={z:+.1f}", (i, max(o, e) + 2), ha="center", fontsize=9,
                        color="#c0392b" if z > 0 else "#1b5e20")
    ax.set_xticks(x_pos)
    ax.set_xticklabels(labels, rotation=30, ha="right")
    ax.set_ylabel("count of period ratios with {r} in ±2% of center")
    ax.set_title("Narrow-window pile-up at Farey-9 fractions (n = 1564)")
    ax.legend()
    fig.tight_layout()
    fig.savefig("c1_pileup.png", dpi=140)
    print("\nplot: c1_pileup.png")


if __name__ == "__main__":
    main()