"""
Conjecture III, first concrete test.

Conjecture III (Arithmetic Hair): Galactic disk substructure (moving
groups, phase-space spirals) is the image of rational points on a
lower-dimensional resonant submanifold of action space under the
galactic potential's flow.

This script does the simplest possible version of that test:

(1) Pull 115k Gaia DR3 stars within ~100 pc with full 6D phase space.
(2) Convert ICRS astrometry to Local Standard of Rest (LSR) galactic
    velocities (U, V, W).
(3) Identify the dominant moving-group peaks in the 2D (U, V) histogram.
(4) Compare the V positions of peaks to simple rational ratios of the
    LSR circular velocity v_c ≈ 232 km/s, and to known bar-resonance
    positions (Hercules at OLR, Hyades / Sirius at spiral-arm features).

A clean detection of "moving-group V values cluster at Farey-low p/q
fractions of v_c" would support Conjecture III. A null result would
indicate either (a) the conjecture is wrong, (b) the conjecture is
right but the substructure is in a different action-space slice, or
(c) the local Gaia sample is too dust-extinguished / kinematically
shallow to see the structure.

Bibliographic note: the canonical decomposition of nearby Gaia (U, V)
into moving groups was done by Antoja et al. 2008 and Famaey et al. 2005;
the bar / spiral resonance interpretation is in Dehnen 2000 and many
follow-ups.
"""

import math
import csv
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from astropy.coordinates import SkyCoord, Galactic
from astropy import units as u


# Standard LSR Solar peculiar velocity (Schoenrich, Binney & Dehnen 2010)
U_SUN_LSR = 11.10   # km/s, toward Galactic center
V_SUN_LSR = 12.24   # km/s, in direction of Galactic rotation
W_SUN_LSR =  7.25   # km/s, toward NGP

# Local Standard of Rest circular velocity (GRAVITY+ / Reid 2014 era)
V_CIRC_LSR = 232.0  # km/s

# Convention for U: many papers use U positive toward Galactic anti-center.
# Astropy returns galactic Cartesian with U positive toward Galactic center
# in the IAU 1991 convention. We follow that here.


def load(path):
    """Load Gaia CSV into a structured numpy array."""
    rows = []
    with open(path) as f:
        reader = csv.DictReader(f)
        for r in reader:
            try:
                rows.append((
                    float(r["ra"]), float(r["dec"]),
                    float(r["parallax"]), float(r["pmra"]), float(r["pmdec"]),
                    float(r["radial_velocity"])
                ))
            except (ValueError, TypeError):
                continue
    arr = np.array(rows, dtype=float)
    print(f"loaded {len(arr)} stars")
    return arr


def compute_uvw(ra_deg, dec_deg, parallax_mas, pmra, pmdec, rv):
    """Heliocentric galactic Cartesian velocities (U, V, W) in km/s.

    Sign convention: U toward Galactic center, V in rotation direction,
    W toward North Galactic Pole.
    """
    distance_pc = 1000.0 / parallax_mas
    c = SkyCoord(
        ra=ra_deg * u.deg,
        dec=dec_deg * u.deg,
        distance=distance_pc * u.pc,
        pm_ra_cosdec=pmra * u.mas / u.yr,
        pm_dec=pmdec * u.mas / u.yr,
        radial_velocity=rv * u.km / u.s,
        frame="icrs",
    )
    cg = c.transform_to(Galactic())
    # Galactic Cartesian velocities (relative to Sun)
    U = cg.velocity.d_x.to(u.km / u.s).value
    V = cg.velocity.d_y.to(u.km / u.s).value
    W = cg.velocity.d_z.to(u.km / u.s).value
    return U, V, W


def main():
    arr = load("gaia_100pc.csv")
    ra, dec, plx, pmra, pmdec, rv = arr.T

    # Filter unphysical
    ok = (plx > 0) & np.isfinite(rv)
    arr = arr[ok]
    ra, dec, plx, pmra, pmdec, rv = arr.T
    print(f"after filtering: {len(arr)} stars")

    # Process in batches to avoid memory blowup
    batch = 20000
    Us, Vs, Ws = [], [], []
    for start in range(0, len(arr), batch):
        end = min(start + batch, len(arr))
        print(f"  computing UVW for batch {start}–{end}")
        U, V, W = compute_uvw(ra[start:end], dec[start:end], plx[start:end],
                              pmra[start:end], pmdec[start:end], rv[start:end])
        Us.append(U); Vs.append(V); Ws.append(W)
    U = np.concatenate(Us); V = np.concatenate(Vs); W = np.concatenate(Ws)

    # Convert to LSR frame (Sun moves at +U_SUN_LSR, +V_SUN_LSR, +W_SUN_LSR)
    U_lsr = U + U_SUN_LSR
    V_lsr = V + V_SUN_LSR
    W_lsr = W + W_SUN_LSR

    # Reasonable kinematic cuts to focus on disk stars
    sane = (np.abs(U_lsr) < 200) & (np.abs(V_lsr) < 200) & (np.abs(W_lsr) < 100)
    U_lsr = U_lsr[sane]; V_lsr = V_lsr[sane]; W_lsr = W_lsr[sane]
    print(f"after sane velocity cut: {len(U_lsr)} stars")

    # ------------- 2-D (U, V) histogram and peak detection ------------
    # Bin (U, V) in km/s
    extent = [-150, 150, -150, 100]
    H, xedges, yedges = np.histogram2d(
        U_lsr, V_lsr, bins=[120, 100], range=[[-150, 150], [-150, 100]]
    )

    # Smooth slightly and find local maxima
    from scipy.ndimage import gaussian_filter, maximum_filter, label
    Hs = gaussian_filter(H, sigma=1.2)
    # Local maxima: H == maximum filter
    local_max = (Hs == maximum_filter(Hs, size=5)) & (Hs > Hs.max() * 0.06)
    # Cluster labels
    labels, n_labels = label(local_max)

    peaks = []
    for i in range(1, n_labels + 1):
        ys, xs_ = np.where(labels == i)
        if len(xs_) == 0:
            continue
        # Use centroid of the contiguous local-max cluster
        u_idx = int(round(np.mean(ys)))
        v_idx = int(round(np.mean(xs_)))
        u_val = 0.5 * (xedges[u_idx] + xedges[u_idx + 1])
        v_val = 0.5 * (yedges[v_idx] + yedges[v_idx + 1])
        height = Hs[u_idx, v_idx]
        peaks.append((u_val, v_val, height))
    peaks.sort(key=lambda p: -p[2])

    print(f"\ndetected {len(peaks)} candidate moving-group peaks in (U, V)_LSR:")
    print(f"  {'#':>3}  {'U_LSR':>7}  {'V_LSR':>7}  {'height':>8}  {'V/v_circ':>9}  identity?")
    # Known moving groups (approximate U, V LSR positions from literature)
    KNOWN = {
        "Hyades":    (-40, -20),
        "Sirius":    (  9,   3),
        "Hercules":  (-30, -50),
        "Pleiades":  (-12, -22),
        "Wolf 630":  (-25, -33),
        "Coma Ber":  ( -7,  -7),
    }
    for i, (u_v, v_v, h) in enumerate(peaks[:20]):
        # Match to nearest known
        nearest = min(KNOWN.items(), key=lambda kv: (kv[1][0] - u_v) ** 2 + (kv[1][1] - v_v) ** 2)
        dist = math.sqrt((nearest[1][0] - u_v) ** 2 + (nearest[1][1] - v_v) ** 2)
        match = nearest[0] if dist < 14 else "—"
        v_over_vc = (v_v + V_CIRC_LSR) / V_CIRC_LSR  # absolute V_phi / V_circ
        print(f"  {i+1:>3}  {u_v:7.1f}  {v_v:7.1f}  {h:8.1f}  {v_over_vc:9.4f}  {match}")

    # ------------- Conjecture-III test: Farey ratios on V_phi/V_circ ------
    # The conjecture predicts moving-group V_phi values cluster at rational
    # p/q × V_circ. Test the top peaks.
    print(f"\n--- Farey-ratio test on V_phi / V_circ (V_circ = {V_CIRC_LSR} km/s) ---")
    # All Farey-15 fractions in (0.5, 1.5) — peaks in this range are nearby disk
    def farey_in_range(max_sum, lo, hi):
        out = []
        for q in range(1, max_sum):
            for p in range(1, max_sum - q + 1):
                if math.gcd(p, q) == 1 and lo <= p/q <= hi:
                    out.append((p, q, p / q))
        out.sort(key=lambda t: t[2])
        return out
    farey_ratios = farey_in_range(15, 0.5, 1.5)
    print(f"Farey-15 fractions in (0.5, 1.5): {len(farey_ratios)} candidate ratios")

    print(f"\n  {'peak':>4}  {'V_LSR':>7}  {'V_phi/Vc':>9}  {'nearest p/q':>12}  {'|delta|':>9}")
    for i, (u_v, v_v, h) in enumerate(peaks[:20]):
        vphi_over_vc = (v_v + V_CIRC_LSR) / V_CIRC_LSR
        # Find nearest Farey-15 ratio
        best = min(farey_ratios, key=lambda pq: abs(pq[2] - vphi_over_vc))
        d = abs(best[2] - vphi_over_vc)
        print(f"  {i+1:>4}  {v_v:7.1f}  {vphi_over_vc:9.4f}  {best[0]}:{best[1]:<5d}  {d:9.4f}")

    # Also test the KNOWN moving-group V_phi values directly (literature positions)
    print(f"\n--- Farey-ratio test on canonical moving-group V positions (literature) ---")
    print(f"  {'group':>10}  {'V_LSR':>7}  {'V_phi/Vc':>9}  {'nearest p/q':>12}  {'|delta|':>9}")
    for name, (uv, vv) in KNOWN.items():
        vphi_over_vc = (vv + V_CIRC_LSR) / V_CIRC_LSR
        best = min(farey_ratios, key=lambda pq: abs(pq[2] - vphi_over_vc))
        d = abs(best[2] - vphi_over_vc)
        print(f"  {name:>10}  {vv:7.1f}  {vphi_over_vc:9.4f}  {best[0]}:{best[1]:<5d}  {d:9.4f}")
    # Add Hercules and the second Hercules clump (Antoja 2008 says clumped near V=-50)
    extras = {"Hercules (OLR)": -50, "Arcturus": -100}
    print(f"\n--- additional resonance-interpreted features ---")
    for name, vv in extras.items():
        vphi_over_vc = (vv + V_CIRC_LSR) / V_CIRC_LSR
        best = min(farey_ratios, key=lambda pq: abs(pq[2] - vphi_over_vc))
        d = abs(best[2] - vphi_over_vc)
        print(f"  {name:>15}  {vv:7.1f}  {vphi_over_vc:9.4f}  {best[0]}:{best[1]:<5d}  {d:9.4f}")

    # Median distance to nearest Farey-15 over all stars (not just peaks)
    vphi_over_vc_all = (V_lsr + V_CIRC_LSR) / V_CIRC_LSR
    vphi_in_range = vphi_over_vc_all[(vphi_over_vc_all > 0.5) & (vphi_over_vc_all < 1.5)]
    nearest_dist = np.array([
        min(abs(pq[2] - v) for pq in farey_ratios) for v in vphi_in_range[:5000]
    ])
    print(f"\nFor random in-range stars (n={len(nearest_dist)}):")
    print(f"  median distance to nearest Farey-15 ratio: {np.median(nearest_dist):.4f}")
    print(f"  fraction within 0.01 of any Farey-15: {(nearest_dist < 0.01).mean():.3f}")
    print(f"  for uniform null, expected fraction within ±0.01: {len(farey_ratios) * 0.02:.3f}")

    # ------------- plots ------------
    fig, ax = plt.subplots(figsize=(9, 7))
    Hp = H.T  # transpose for imshow
    im = ax.imshow(np.log10(Hp + 1), origin="lower", extent=extent,
                   aspect="auto", cmap="viridis")
    plt.colorbar(im, ax=ax, label=r"$\log_{10}(N+1)$")
    # Overlay peak markers
    for i, (u_v, v_v, h) in enumerate(peaks[:12]):
        ax.plot(u_v, v_v, "r+", markersize=12, mew=2)
        ax.text(u_v + 3, v_v + 1, f"{i+1}", fontsize=9, color="red")
    # Overlay known moving groups
    for name, (uv, vv) in KNOWN.items():
        ax.plot(uv, vv, "yo", markersize=8, mfc="none", mew=1.5)
        ax.text(uv + 3, vv - 5, name, fontsize=8, color="yellow")
    # Mark V values at Farey-low ratios
    for p, q, ratio in farey_ratios[:10]:
        v_at = (ratio - 1) * V_CIRC_LSR
        if -150 < v_at < 100:
            ax.axhline(v_at, color="white", lw=0.4, alpha=0.3)
    ax.set_xlabel(r"$U_{\rm LSR}$  (km s$^{-1}$, → Galactic center)")
    ax.set_ylabel(r"$V_{\rm LSR}$  (km s$^{-1}$, → Galactic rotation)")
    ax.set_title(f"Gaia DR3 within 100 pc:  (U, V)$_{{\\rm LSR}}$ of {len(U_lsr)} stars")
    fig.tight_layout()
    fig.savefig("c3_uv_plane.png", dpi=140)

    # V-only histogram with Farey markers
    fig, ax = plt.subplots(figsize=(10, 4.5))
    ax.hist(V_lsr, bins=200, range=(-150, 100), color="#16a085", alpha=0.85,
            edgecolor="white", linewidth=0.4)
    for p, q, ratio in farey_ratios:
        v_at = (ratio - 1) * V_CIRC_LSR
        if -150 < v_at < 100:
            ax.axvline(v_at, color="#c0392b", lw=0.5, alpha=0.4)
            ax.text(v_at, ax.get_ylim()[1] * 0.92, f"{p}:{q}", fontsize=7, ha="center",
                    color="#c0392b", rotation=90)
    # Mark known moving group V values
    for name, (uv, vv) in KNOWN.items():
        ax.axvline(vv, color="#f39c12", lw=1.5, alpha=0.7)
        ax.text(vv, ax.get_ylim()[1] * 0.78, name, fontsize=8, ha="center",
                color="#d35400", rotation=90)
    ax.set_xlabel(r"$V_{\rm LSR}$  (km s$^{-1}$)")
    ax.set_ylabel("count")
    ax.set_title(r"$V_{\rm LSR}$ distribution with Farey-15 fractions of $v_{\rm circ}$ marked")
    fig.tight_layout()
    fig.savefig("c3_v_histogram.png", dpi=140)

    print("\nplots: c3_uv_plane.png  c3_v_histogram.png")


if __name__ == "__main__":
    main()
