"""Shared constants and helpers for the flyby study."""
import json
import os

import numpy as np
import pandas as pd

HERE = os.path.dirname(os.path.abspath(__file__))
STUDY = os.path.dirname(HERE)
RAW = os.environ.get("FLYBY_RAW", "fermi_raw/flyby")
DATA = os.path.join(STUDY, "data")
FIGS = os.path.join(STUDY, "figures")
os.makedirs(DATA, exist_ok=True)
os.makedirs(FIGS, exist_ok=True)

KMS_TO_PCMYR = 1.0227121650537077   # 1 km/s in pc/Myr
K_PM = 4.740470463533349            # km/s per (mas/yr * kpc) ; vT = K * mu[mas/yr] / plx[mas]
SEED = 20261011

# ICRS -> Galactic rotation matrix (Hipparcos/Gaia convention)
A_G = np.array([
    [-0.0548755604162154, -0.8734370902348850, -0.4838350155487132],
    [+0.4941094278755837, -0.4448296299600112, +0.7469822444972189],
    [-0.8676661490190047, -0.1980763734312015, +0.4559837761750669],
])


def save_json(name, obj):
    path = os.path.join(DATA, name)
    with open(path, "w") as f:
        json.dump(obj, f, indent=1, default=float)
    return path


def load_json(name):
    with open(os.path.join(DATA, name)) as f:
        return json.load(f)


def galactic_xyz_uvw(ra, dec, plx, pmra, pmdec, rv):
    """Heliocentric Galactic position [pc] and velocity [km/s] (U toward the
    Galactic centre, V toward rotation, W toward the NGP).  rv may be NaN."""
    ra = np.radians(np.asarray(ra, float))
    dec = np.radians(np.asarray(dec, float))
    r = 1000.0 / np.asarray(plx, float)
    ca, sa, cd, sd = np.cos(ra), np.sin(ra), np.cos(dec), np.sin(dec)
    n = np.stack([cd * ca, cd * sa, sd], axis=1)          # line of sight
    e_a = np.stack([-sa, ca, np.zeros_like(ra)], axis=1)  # increasing RA
    e_d = np.stack([-sd * ca, -sd * sa, cd], axis=1)      # increasing Dec
    va = K_PM * np.asarray(pmra, float) / np.asarray(plx, float)
    vd = K_PM * np.asarray(pmdec, float) / np.asarray(plx, float)
    v_icrs = (np.asarray(rv, float)[:, None] * n + va[:, None] * e_a + vd[:, None] * e_d)
    pos = (r[:, None] * n) @ A_G.T
    vel = v_icrs @ A_G.T
    return pos, vel


def linear_perihelion(plx, pmra, pmdec, rv):
    """Linear-motion perihelion: time t_ph [Myr, negative = past], distance
    d_ph [pc], speed v [km/s], tangential speed vT [km/s]."""
    plx = np.asarray(plx, float)
    r = 1000.0 / plx
    vt = K_PM * np.hypot(np.asarray(pmra, float), np.asarray(pmdec, float)) / plx
    rv = np.asarray(rv, float)
    v = np.hypot(vt, rv)
    d_ph = r * vt / v
    t_ph = -r * rv / v**2 / KMS_TO_PCMYR
    return t_ph, d_ph, v, vt


def load_gcns():
    g = pd.read_csv(os.path.join(RAW, "gcns_table1c.csv"))
    g = g.rename(columns={"GaiaEDR3": "source_id"})
    g["dist"] = g["Dist50"] * 1000.0                      # pc (GCNS posterior median)
    g["MG"] = g["Gmag"] - 5.0 * np.log10(g["dist"]) + 5.0
    g["vT"] = K_PM * np.hypot(g["pmRA"], g["pmDE"]) / g["Plx"]
    return g


WIDE_SEP_PC = 0.1   # maximum projected separation for supplementary pair merging


def _wide_pairs(g, sep_max=WIDE_SEP_PC):
    """Co-moving pairs in GCNS with projected separation < sep_max, consistent
    distances and a tangential-velocity difference compatible with a bound
    orbit.  Supplements GCNS table 3, which misses many nearby, large-angle
    pairs (e.g. 70 Oph, 36 Oph, eta Cas)."""
    from scipy.spatial import cKDTree
    ra = np.radians(g["RA_ICRS"].values)
    de = np.radians(g["DE_ICRS"].values)
    u = np.stack([np.cos(de) * np.cos(ra), np.cos(de) * np.sin(ra), np.sin(de)], axis=1)
    d = g["dist"].values
    tree = cKDTree(u)
    th = np.minimum(sep_max / d, 0.5)
    nb = tree.query_ball_point(u, 2 * np.sin(th / 2), return_sorted=False)
    i = np.repeat(np.arange(len(g)), [len(x) for x in nb])
    j = np.concatenate([np.asarray(x, int) for x in nb])
    k = i < j
    i, j = i[k], j[k]
    dm = 0.5 * (d[i] + d[j])
    ang = np.arccos(np.clip((u[i] * u[j]).sum(1), -1, 1))
    sep = ang * np.minimum(d[i], d[j])
    rel = (g["e_Plx"] / g["Plx"]).values
    sig = np.hypot(d[i] * rel[i], d[j] * rel[j])
    ok_d = np.abs(d[i] - d[j]) < sep_max + 3 * sig + 0.02 * dm
    va = K_PM * g["pmRA"].values / g["Plx"].values
    vd = K_PM * g["pmDE"].values / g["Plx"].values
    dv = np.hypot(va[i] - va[j], vd[i] - vd[j])
    sep_au = np.maximum(sep * 206264.8, 1.0)
    ok_v = dv < 2.0 + 84.0 / np.sqrt(sep_au)       # 2 km/s + 2 x circular speed (2 Msun)
    ok = (sep < sep_max) & ok_d & ok_v
    return i[ok], j[ok]


def gcns_systems(g, supplement=True):
    """Attach a system id to every GCNS source (union-find).  Pairs come from
    GCNS table 3 (`binary` flag) and, if supplement=True, from the co-moving
    pair search above.  Returns g with columns sys_id, is_primary."""
    b = pd.read_csv(os.path.join(RAW, "gcns_resolved_binaries.csv"))
    b = b[b["binary"] == 1]
    idx = pd.Series(np.arange(len(g)), index=g["source_id"].values)
    a1 = idx.reindex(b["SourceID1"].values).values
    a2 = idx.reindex(b["SourceID2"].values).values
    ok = ~np.isnan(a1) & ~np.isnan(a2)
    a1, a2 = a1[ok].astype(int), a2[ok].astype(int)
    if supplement:
        w1, w2 = _wide_pairs(g)
        a1, a2 = np.r_[a1, w1], np.r_[a2, w2]
    from scipy.sparse import coo_matrix
    from scipy.sparse.csgraph import connected_components
    n = len(g)
    m = coo_matrix((np.ones(len(a1)), (a1, a2)), shape=(n, n))
    _, lab = connected_components(m, directed=False)
    g = g.copy()
    g["sys_id"] = lab
    # primary = brightest member (missing G sorts last)
    gm = g["Gmag"].fillna(99.0).values
    order = np.lexsort((gm, lab))
    first = np.ones(n, bool)
    first[1:] = lab[order][1:] != lab[order][:-1]
    is_prim = np.zeros(n, bool)
    is_prim[order[first]] = True
    g["is_primary"] = is_prim
    return g


def load_rv_sample():
    """Gaia DR3 sources with RV and parallax > 8 mas."""
    return pd.read_csv(os.path.join(RAW, "gaia_dr3_rv_plx8.csv"))


def load_encounter_candidates():
    return pd.read_csv(os.path.join(RAW, "gaia_dr3_encounter_candidates.csv"))


def load_tenpc():
    """10 pc sample; Gaia ids are read as exact nullable integers."""
    t = pd.read_csv(os.path.join(RAW, "tenpc_reyle2021.csv"),
                    dtype={"GaiaEDR3": "Int64", "GaiaDR2": "Int64"})
    return t


def tenpc_systems():
    """Non-planet members of the 10 pc sample and a per-system summary
    (rk: 0 = AFGK star, 1 = M dwarf, 2 = white dwarf, 3 = brown dwarf)."""
    t = load_tenpc()
    t = t[t["ObjType"] != "Planet"].copy()
    rank = {"*": 0, "LM": 1, "LM?": 1, "WD": 2, "WD?": 2, "BD": 3, "BD?": 3}
    t["rk"] = t["ObjType"].map(rank)
    prim = t.sort_values(["Sys", "rk"]).groupby("Sys").first()
    return t, prim


def poisson_interval(n, cl=0.6827):
    """Central Poisson confidence interval for an observed count n."""
    from scipy.stats import chi2
    a = 1 - cl
    lo = 0.0 if n == 0 else 0.5 * chi2.ppf(a / 2, 2 * n)
    hi = 0.5 * chi2.ppf(1 - a / 2, 2 * (n + 1))
    return lo, hi