#!/usr/bin/env python3
"""Step 2.  Direct Gaia DR3 flyby census (linear motion) and its comparison with
the RV-free flux estimate of step 1.

Observed: Gaia DR3 stars within 100 pc (GCNS members) with a clean DR3 radial
velocity whose linear-motion perihelion lies within d of the Sun and within
|t| < T.

Predicted for exactly that sample: for a star now at distance r with speed v
the velocity direction is isotropic with respect to the line of sight
(sky-averaged).  With mu = |cos(theta)| uniform on [0,1], linear motion gives
d_ph = r sin(theta) and |t_ph| = r mu / v, so

    P(d_ph < d, |t_ph| < T | r, v) = max(0, min(1, v T / r) - mu0),
    mu0 = sqrt(1 - (d/r)^2)  (0 if r < d).

N_pred = sum_i E_v[P], with the speed distribution measured from the RV
sample.  The ratio N_obs/N_pred tests the linear-motion perihelia, the d^2
geometry and the RV data quality without any completeness model.  It is an
internal consistency check on the clean-RV subsample (it shares the isotropy
assumption and the speed distribution with the rate estimator); it does not
measure the density or speed of the stars without RVs.  The fraction of all
expected encounters that have RVs, C = sum_{RV stars}/sum_{all GCNS systems},
is the RV completeness.

Model-free lower bound (revision): the raw number of clean-RV Gaia DR3
sources, at any distance, with d_ph < d and |t_ph| < T, divided by 2 T d^2,
is a lower limit on the rate within 1 pc that uses no completeness model, no
isotropy assumption and no speed distribution; it is a lower limit because
every star without a clean RV is missing from the count.
"""
import os

import numpy as np
import pandas as pd

from common import (DATA, KMS_TO_PCMYR, RAW, gcns_systems, linear_perihelion, load_encounter_candidates,
                    load_gcns, load_rv_sample, poisson_interval, save_json, tenpc_systems)

D_LIST = [0.5, 1.0, 2.0, 3.0, 5.0, 7.0, 10.0]
R_MAX = 100.0
T_MAIN = 1.0      # Myr


def clean_mask(df, ruwe_cut=None):
    m = (df["radial_velocity"].notna() & (df["radial_velocity_error"] < 5.0)
         & (df["rv_expected_sig_to_noise"].fillna(0) >= 5.0))
    if ruwe_cut is not None:
        m &= df["ruwe"] < ruwe_cut
    return m


class SpeedDist:
    """Weighted empirical speed distribution with cumulative weight and
    cumulative first moment, for exact expectation values."""

    def __init__(self, v, w):
        o = np.argsort(v)
        self.v = np.asarray(v, float)[o]
        w = np.asarray(w, float)[o]
        w = w / w.sum()
        self.cw = np.concatenate([[0.0], np.cumsum(w)])
        self.cm = np.concatenate([[0.0], np.cumsum(w * self.v)])

    def F(self, x):      # P(v <= x)
        return self.cw[np.searchsorted(self.v, x, side="right")]

    def M(self, x):      # E[v ; v <= x]
        return self.cm[np.searchsorted(self.v, x, side="right")]


def expected_encounters(r, d, T, sd, vmax=np.inf):
    """sum_i E_v[ max(0, min(1, v T / r_i) - mu0_i) ; v < vmax ]
    (v in km/s, T in Myr, r and d in pc), evaluated exactly on the empirical
    speed distribution sd."""
    r = np.asarray(r, float)
    mu0 = np.sqrt(np.clip(1.0 - (d / r) ** 2, 0.0, None))
    k = KMS_TO_PCMYR * T / r                 # v*k = v T / r
    vlo = np.minimum(mu0 / k, vmax)          # below this speed the encounter is outside the window
    vhi = np.minimum(1.0 / k, vmax)          # above this speed min(1, vT/r) = 1
    vtop = np.full_like(r, vmax)
    e = (k * (sd.M(vhi) - sd.M(vlo)) + (sd.F(vtop) - sd.F(vhi)) - mu0 * (sd.F(vtop) - sd.F(vlo)))
    return float(np.sum(np.clip(e, 0.0, None)))


def main():
    out = {}
    g = gcns_systems(load_gcns())
    prim = g[g["is_primary"] & (g["dist"] < R_MAX)].copy()
    rv = load_rv_sample()
    cols = ["source_id", "radial_velocity", "radial_velocity_error", "rv_expected_sig_to_noise",
            "ruwe", "rv_nb_transits", "grvs_mag", "bp_rp", "phot_g_mean_mag", "parallax",
            "pmra", "pmdec", "rv_template_teff"]
    m = prim.merge(rv[cols], on="source_id", how="left")
    m["clean"] = clean_mask(m)
    t, dph, v, vt = linear_perihelion(m["Plx"], m["pmRA"], m["pmDE"], m["radial_velocity"])
    m["t_ph"], m["d_ph"], m["v"] = t, dph, v
    out["parent_systems_100pc"] = int(len(m))
    out["parent_with_any_rv"] = int(m["radial_velocity"].notna().sum())
    out["parent_with_clean_rv"] = int(m["clean"].sum())
    print("parent systems <100 pc:", len(m), " with RV:", out["parent_with_any_rv"],
          " clean RV:", out["parent_with_clean_rv"])

    # RV availability map (for the report / figure)
    m["MGbin"] = np.floor(m["MG"])
    m["dbin"] = np.floor(m["dist"] / 10) * 10
    fmap = m.groupby(["MGbin", "dbin"])["clean"].agg(["mean", "size"]).reset_index()
    fmap.to_csv(os.path.join(DATA, "rv_availability_map.csv"), index=False)
    fm = m.groupby("MGbin")["clean"].agg(["mean", "size"])
    print(fm.round(3).T.to_string())

    # ---- speed distribution: volume-limited 50 pc RV sample, LF-weighted ----
    vs = pd.read_csv(os.path.join(DATA, "velocity_sample_50pc.csv"))
    sp = np.sqrt(vs["U"]**2 + vs["V"]**2 + vs["W"]**2).values
    vg_rv = SpeedDist(sp, np.ones(len(sp)))       # RV-sample speeds
    vg_all = SpeedDist(sp, vs["weight"].values)   # reweighted to all systems

    # ---- observed versus predicted ---------------------------------------
    rows = []
    for T in [0.5, 1.0, 1.5]:
        for ruwe_cut in [None, 1.4]:
            cl = clean_mask(m, ruwe_cut)
            rcl = m.loc[cl, "dist"].values
            for vmax in [np.inf, 150.0]:
                for d in D_LIST:
                    obs = int((cl & (m["d_ph"] < d) & (m["t_ph"].abs() < T) & (m["v"] < vmax)).sum())
                    pred = expected_encounters(rcl, d, T, vg_rv, vmax)
                    pred_all = expected_encounters(m["dist"].values, d, T, vg_all)
                    lo, hi = poisson_interval(obs)
                    rows.append(dict(T_Myr=T, ruwe_cut=ruwe_cut if ruwe_cut else 0, vmax=vmax, d_pc=d,
                                     N_obs=obs, N_pred=pred, ratio=obs / pred if pred > 0 else np.nan,
                                     ratio_lo=lo / pred, ratio_hi=hi / pred,
                                     N_pred_all_systems=pred_all,
                                     rv_completeness=pred / pred_all if np.isinf(vmax) else np.nan))
    cen = pd.DataFrame(rows)
    cen.to_csv(os.path.join(DATA, "census_obs_vs_pred.csv"), index=False)
    main_tab = cen[(cen.T_Myr == T_MAIN) & (cen.ruwe_cut == 0) & np.isinf(cen.vmax)]
    print(main_tab.round(3).to_string(index=False))
    print(cen[(cen.T_Myr == T_MAIN) & (cen.ruwe_cut == 1.4) & np.isinf(cen.vmax)].round(3).to_string(index=False))
    print(cen[(cen.T_Myr == T_MAIN) & (cen.ruwe_cut == 0) & (cen.vmax == 150)].round(3).to_string(index=False))
    print(cen[(cen.T_Myr == 0.5) & (cen.ruwe_cut == 0) & np.isinf(cen.vmax)].round(3).to_string(index=False))

    # power-law slope of the observed cumulative counts N(<d) ~ d^alpha
    # (maximum likelihood for a truncated power law between d_min and d_max)
    def ml_slope(dvals, dmin, dmax):
        x = dvals[(dvals >= dmin) & (dvals < dmax)]
        from scipy.optimize import brentq
        def score(a):
            # d/da log L for pdf a x^(a-1)/(dmax^a - dmin^a)
            return (len(x) / a + np.log(x).sum()
                    - len(x) * (dmax**a * np.log(dmax) - dmin**a * np.log(dmin)) / (dmax**a - dmin**a))
        a = brentq(score, 0.2, 6.0)
        # Fisher information numerically
        h = 1e-3
        info = -(score(a + h) - score(a - h)) / (2 * h)
        return a, 1 / np.sqrt(info), len(x)

    enc = m[m["clean"] & (m["t_ph"].abs() < T_MAIN)]
    a, ae, na = ml_slope(enc["d_ph"].values, 0.0 + 1e-3, 10.0)
    out["slope_dph_alpha"] = a
    out["slope_dph_alpha_err"] = ae
    out["slope_dph_N"] = na
    print(f"cumulative N(<d) ~ d^alpha: alpha = {a:.3f} +- {ae:.3f} (N={na}); note the r<100 pc"
          " truncation is distance independent so alpha=2 is expected")

    r5 = main_tab[main_tab.d_pc == 5.0].iloc[0]
    out["main_T_Myr"] = T_MAIN
    out["obs_over_pred_5pc"] = float(r5["ratio"])
    out["obs_over_pred_5pc_lo"] = float(r5["ratio_lo"])
    out["obs_over_pred_5pc_hi"] = float(r5["ratio_hi"])
    out["N_obs_5pc"] = int(r5["N_obs"])
    out["rv_completeness_5pc_1Myr_100pc"] = float(r5["rv_completeness"])
    r10 = main_tab[main_tab.d_pc == 10.0].iloc[0]
    out["obs_over_pred_10pc"] = float(r10["ratio"])
    out["obs_over_pred_10pc_lo"] = float(r10["ratio_lo"])
    out["obs_over_pred_10pc_hi"] = float(r10["ratio_hi"])
    out["N_obs_10pc"] = int(r10["N_obs"])
    r1 = main_tab[main_tab.d_pc == 1.0].iloc[0]
    out["N_obs_1pc"] = int(r1["N_obs"]); out["N_pred_1pc"] = float(r1["N_pred"])

    # ---- all-distance candidate list: comparison with Bailer-Jones 2022 ----
    c = load_encounter_candidates()
    tc, dc, vc, vtc = linear_perihelion(c["parallax"], c["pmra"], c["pmdec"], c["radial_velocity"])
    c["t_ph"], c["d_ph"], c["v"] = tc, dc, vc
    c["clean"] = clean_mask(c)
    out["candidates_rows"] = int(len(c))
    tab = []
    for T in [1.0, 2.0, 5.0, 15.0]:
        for d in [1.0, 2.0, 5.0]:
            sel = (c["d_ph"] < d) & (c["t_ph"].abs() < T)
            tab.append(dict(T_Myr=T, d_pc=d, N_raw=int(sel.sum()), N_clean=int((sel & c["clean"]).sum()),
                            N_clean_ruwe=int((sel & c["clean"] & (c["ruwe"] < 1.4)).sum())))
    tab = pd.DataFrame(tab)
    tab.to_csv(os.path.join(DATA, "census_all_distance_counts.csv"), index=False)
    print(tab.to_string(index=False))

    # ---- model-free lower bound on the rate (no completeness correction) ----
    # Primary window fixed in advance as d < 5 pc, |t| < 0.25 Myr (the window
    # the referee used); the other windows show the trend.  Counts are Gaia
    # sources (not merged systems): a wide pair would count twice.
    lb = []
    for T in [0.1, 0.25, 0.5, 1.0]:
        for d in [2.0, 5.0, 10.0]:
            sel = c["clean"] & (c["d_ph"] < d) & (c["t_ph"].abs() < T)
            n = int(sel.sum())
            lo, hi = poisson_interval(n)
            norm = 2.0 * T * d**2
            lb.append(dict(T_Myr=T, d_pc=d, N_clean=n, gamma1pc_lower=n / norm,
                           gamma1pc_lower_err=np.sqrt(n) / norm, lo68=lo / norm, hi68=hi / norm,
                           median_dist_now_pc=float(np.median(1000.0 / c.loc[sel, "parallax"])) if n else np.nan))
    lb = pd.DataFrame(lb)
    lb.to_csv(os.path.join(DATA, "rate_lower_bound.csv"), index=False)
    print(lb.round(3).to_string(index=False))
    r0 = lb[(lb.T_Myr == 0.25) & (lb.d_pc == 5.0)].iloc[0]
    out["lower_bound_window"] = "d < 5 pc, |t| < 0.25 Myr, clean Gaia DR3 RV, any distance"
    out["lower_bound_N"] = int(r0["N_clean"])
    out["lower_bound_gamma1pc"] = float(r0["gamma1pc_lower"])
    out["lower_bound_gamma1pc_err"] = float(r0["gamma1pc_lower_err"])

    bj = pd.read_csv(os.path.join(RAW, "bailerjones2022_table12.csv"))
    x = bj.merge(c[["source_id", "t_ph", "d_ph", "v", "clean", "ruwe"]], left_on="GaiaDR3",
                 right_on="source_id", how="left")
    out["bj22_rows"] = int(len(bj))
    out["bj22_matched_in_candidates"] = int(x["source_id"].notna().sum())
    xm = x[x["source_id"].notna()]
    # t_ph in BJ22 is in kyr
    out["bj22_median_abs_dd_pc"] = float(np.median(np.abs(xm["d_ph"] - xm["dphmed"])))
    out["bj22_median_abs_dt_kyr"] = float(np.median(np.abs(xm["t_ph"] * 1000 - xm["tphmed"])))
    near = xm[xm["tphmed"].abs() < 2000]
    out["bj22_within2Myr_N"] = int(len(near))
    out["bj22_within2Myr_median_abs_dd_pc"] = float(np.median(np.abs(near["d_ph"] - near["dphmed"])))
    out["bj22_within2Myr_max_abs_dd_pc"] = float(np.max(np.abs(near["d_ph"] - near["dphmed"])))
    x[["GaiaDR3", "tphmed", "dphmed", "t_ph", "d_ph", "v", "clean", "ruwe", "Flag"]].to_csv(
        os.path.join(DATA, "bj22_comparison.csv"), index=False)
    print("BJ22 match:", out["bj22_matched_in_candidates"], "of", len(bj),
          " median |d_lin - d_BJ22| =", round(out["bj22_median_abs_dd_pc"], 3), "pc (all);",
          round(out["bj22_within2Myr_median_abs_dd_pc"], 3), "pc for |t|<2 Myr, N =", len(near))
    my1 = c[(c["d_ph"] < 1.0)]
    out["my_linear_within1pc_any_time_plx_gt2"] = int(len(my1))
    out["my_linear_within1pc_in_bj22"] = int(my1["source_id"].isin(bj["GaiaDR3"]).sum())

    # closest encounters within +-1.5 Myr (clean)
    top = c[c["clean"] & (c["t_ph"].abs() < 1.5)].sort_values("d_ph").head(25)
    top = top.assign(dist_pc=1000 / top["parallax"])
    top[["source_id", "dist_pc", "phot_g_mean_mag", "radial_velocity", "radial_velocity_error",
         "v", "t_ph", "d_ph", "ruwe", "rv_expected_sig_to_noise"]].to_csv(
        os.path.join(DATA, "closest_encounters_1p5Myr.csv"), index=False)
    print(top[["source_id", "dist_pc", "phot_g_mean_mag", "v", "t_ph", "d_ph", "ruwe"]].round(3).head(12).to_string(index=False))

    # ---- independent check with the literature-complete 10 pc sample -------
    _, pr = tenpc_systems()
    pr = pr[(pr["plx"] >= 100) & (pr["rk"] < 3)]
    has = pr["RV"].notna() & pr["pmRA"].notna()
    q = pr[has]
    tt, dd, vv, _ = linear_perihelion(q["plx"], q["pmRA"], q["pmDE"], q["RV"])
    r10 = 1000.0 / q["plx"].values
    rows10 = []
    for T in [0.1, 0.2, 0.3]:
        for d in [2.0, 3.0, 5.0]:
            obs = int(((dd < d) & (np.abs(tt) < T)).sum())
            pred = expected_encounters(r10, d, T, vg_all)
            lo, hi = poisson_interval(obs)
            rows10.append(dict(T_Myr=T, d_pc=d, N_obs=obs, N_pred=pred, ratio=obs / pred,
                               ratio_lo=lo / pred, ratio_hi=hi / pred))
    t10 = pd.DataFrame(rows10)
    t10.to_csv(os.path.join(DATA, "census_tenpc_check.csv"), index=False)
    out["tenpc_stellar_systems"] = int(len(pr)); out["tenpc_stellar_systems_with_rv"] = int(len(q))
    print("10 pc sample: stellar systems", len(pr), "with RV+PM", len(q))
    print(t10.round(3).to_string(index=False))

    m.loc[m["clean"] & (m["t_ph"].abs() < 2.0) & (m["d_ph"] < 15),
          ["source_id", "dist", "MG", "v", "t_ph", "d_ph", "ruwe"]].to_csv(
        os.path.join(DATA, "census_encounters_100pc.csv"), index=False)
    save_json("s2_census.json", out)
    return out


if __name__ == "__main__":
    main()