#!/usr/bin/env python3
"""Step 1.  Local density of star systems and their speed distribution relative
to the Sun, and the resulting RV-free estimate of the stellar encounter rate.

Key identity.  For a population of number density n whose velocity (relative to
the Sun) is statistically independent of position, the rate at which members
pass within distance d of the Sun is  Gamma(d) = pi d^2 n <v>  (gravitational
focusing is < 1e-4 for d >= 0.1 pc at v ~ 50 km/s).  For a volume-limited,
all-sky sample the line of sight is isotropic with respect to any fixed
velocity vector, so <v_T> = (pi/4) <v> exactly, whatever the anisotropy of the
velocity ellipsoid.  Hence

    Gamma(d) = 4 d^2 * (1/V) * sum_i v_T,i

which needs only parallaxes and proper motions of a volume-complete sample.
No radial velocities are needed, so there is no RV-incompleteness correction.
"""
import os

import numpy as np
import pandas as pd

from common import (DATA, KMS_TO_PCMYR, SEED, galactic_xyz_uvw, gcns_systems, linear_perihelion,
                    load_gcns, load_rv_sample, save_json, tenpc_systems)

VT_CAP = 600.0      # km/s; sources above this are treated as astrometric outliers
R_MAIN = 50.0       # pc; volume used for the headline density and <v>
MISSING_CORR = 1.045      # correction for systems absent from GCNS (see report)
MISSING_CORR_ERR = 0.045  # half-range: 1.00 (none missing at 50 pc) .. 1.09 (as at 10 pc)


def sphere_volume(r):
    return 4.0 / 3.0 * np.pi * r**3


def tenpc_check(g):
    """Compare GCNS systems with the literature-complete 10 pc sample."""
    t, prim = tenpc_systems()
    inside = prim[prim["plx"] >= 100.0]
    stellar = inside[inside["rk"] < 3]
    V = sphere_volume(10.0)
    res = dict(n_systems_all=int(len(inside)), n_systems_stellar=int(len(stellar)),
               n_stellar_objects=int(((t["rk"] < 3) & (t["plx"] >= 100.0)).sum()),
               dens_systems_all=len(inside) / V, dens_systems_stellar=len(stellar) / V,
               dens_systems_stellar_err=np.sqrt(len(stellar)) / V)
    tt = t[t["GaiaEDR3"].notna() & t["Sys"].isin(stellar.index)].copy()
    tt["sid"] = tt["GaiaEDR3"].astype("int64")
    mp = tt.merge(g[["source_id", "sys_id"]], left_on="sid", right_on="source_id")
    res["stellar_systems_with_gcns_member"] = int(mp["Sys"].nunique())
    res["distinct_gcns_systems_matched"] = int(mp["sys_id"].nunique())
    res["split_systems"] = int((mp.groupby("Sys")["sys_id"].nunique() > 1).sum())
    res["false_merges"] = int((mp.groupby("sys_id")["Sys"].nunique() > 1).sum())
    f = res["stellar_systems_with_gcns_member"] / len(stellar)
    res["gcns_system_completeness_10pc"] = f
    res["gcns_system_completeness_10pc_err"] = float(np.sqrt(f * (1 - f) / len(stellar)))

    # What a GCNS "system" is: Gaia-visible brown dwarfs are GCNS sources too.
    # Count the GCNS sources and systems within 10 pc whose 10 pc sample
    # classification (of the source / of the brightest GCNS member) is BD.
    ta = t[t["GaiaEDR3"].notna()].copy()
    ta["sid"] = ta["GaiaEDR3"].astype("int64")
    g10 = g[g["dist"] < 10.0]
    ma = g10.merge(ta[["sid", "rk", "Sys"]], left_on="source_id", right_on="sid", how="left")
    res["gcns_sources_10pc"] = int(len(g10))
    res["gcns_sources_10pc_matched"] = int(ma["rk"].notna().sum())
    res["gcns_sources_10pc_brown_dwarfs"] = int((ma["rk"] == 3).sum())
    pm_ = ma[ma["is_primary"]]
    res["gcns_systems_10pc"] = int(len(pm_))
    res["gcns_systems_10pc_bd_primary"] = int((pm_["rk"] == 3).sum())

    # Independent rate from the 10 pc sample alone: its stellar-system density
    # times the mean 3D speed of its own members (literature RVs and proper
    # motions), sharing nothing with the GCNS estimator.
    q = stellar[stellar["RV"].notna() & stellar["pmRA"].notna() & stellar["pmDE"].notna()]
    _, _, v3, _ = linear_perihelion(q["plx"], q["pmRA"], q["pmDE"], q["RV"])
    res["tenpc_own_speed_N"] = int(len(q))
    res["tenpc_own_speed_mean_kms"] = float(v3.mean())
    res["tenpc_own_speed_err_kms"] = float(v3.std() / np.sqrt(len(q)))
    gam = np.pi * res["dens_systems_stellar"] * v3.mean() * KMS_TO_PCMYR
    res["gamma1pc_tenpc_own_speed"] = float(gam)
    res["gamma1pc_tenpc_own_speed_err"] = float(gam * np.hypot(1 / np.sqrt(len(stellar)),
                                                               res["tenpc_own_speed_err_kms"] / v3.mean()))
    return res


def ipf_weights(mg_bin_rv, vt_bin_rv, mg_bin_all, vt_bin_all, n_iter=20):
    """Iterative proportional fitting: weights for the RV sample so that its
    marginal distributions in M_G and in v_T match those of the complete
    sample (v_T is known for every star, so this removes the kinematic bias
    of RV availability to first order)."""
    w = np.ones(len(mg_bin_rv))
    for _ in range(n_iter):
        for b_rv, b_all in ((mg_bin_rv, mg_bin_all), (vt_bin_rv, vt_bin_all)):
            tot_all = np.bincount(b_all, minlength=b_all.max() + 1).astype(float)
            tot_rv = np.bincount(b_rv, weights=w, minlength=len(tot_all))
            ratio = np.where(tot_rv > 0, tot_all / np.maximum(tot_rv, 1e-300), 1.0)
            w = w * ratio[b_rv]
    return w / w.sum()


def main():
    rng = np.random.default_rng(SEED)
    out = {}
    g0 = load_gcns()
    g_t3 = gcns_systems(g0, supplement=False)
    g = gcns_systems(g0, supplement=True)
    out["gcns_rows"] = int(len(g))
    out["gcns_systems_total_table3_only"] = int(g_t3["is_primary"].sum())
    out["gcns_systems_total"] = int(g["is_primary"].sum())

    prim = g[g["is_primary"]].copy()
    prim_t3 = g_t3[g_t3["is_primary"]]
    # ---- density versus radius -------------------------------------------
    rows = []
    for R in [10, 15, 20, 25, 30, 40, 50, 60, 75, 100]:
        V = sphere_volume(R)
        n_src = int((g["dist"] < R).sum())
        n_sys = int((prim["dist"] < R).sum())
        rows.append(dict(R_pc=R, n_sources=n_src, n_systems_table3=int((prim_t3["dist"] < R).sum()),
                         n_systems=n_sys, dens_sources=n_src / V, dens_systems=n_sys / V,
                         dens_systems_err=np.sqrt(n_sys) / V))
    dens = pd.DataFrame(rows)
    dens.to_csv(os.path.join(DATA, "density_vs_radius.csv"), index=False)
    print(dens.round(5).to_string(index=False))

    tp = tenpc_check(g)
    out["tenpc"] = tp
    print("10 pc sample check:", {k: (round(v, 4) if isinstance(v, float) else v) for k, v in tp.items()})

    # ---- the RV-free flux sum --------------------------------------------
    sel = prim[(prim["dist"] < R_MAIN)].copy()
    n_out = int((sel["vT"] > VT_CAP).sum())
    sel = sel[sel["vT"] <= VT_CAP]
    V = sphere_volume(R_MAIN)
    n_sys = len(sel) / V
    vt = sel["vT"].values
    mean_vT = vt.mean()
    mean_v = 4.0 / np.pi * mean_vT
    boots = []
    for _ in range(400):
        boots.append(rng.choice(vt, rng.poisson(len(vt))).sum())
    flux = 4.0 * vt.sum() / V * KMS_TO_PCMYR          # = pi n <v>  [per pc^2 per Myr]
    flux_err = 4.0 * np.std(boots) / V * KMS_TO_PCMYR
    out["R_main_pc"] = R_MAIN
    out["n_systems_gcns_50pc_count"] = int(len(sel))
    out["n_systems_gcns_50pc"] = n_sys
    out["n_systems_gcns_50pc_err"] = np.sqrt(len(sel)) / V
    out["n_vT_outliers_removed"] = n_out
    out["mean_vT_kms"] = mean_vT
    out["mean_speed_kms"] = mean_v
    out["mean_speed_staterr_kms"] = float(4 / np.pi * vt.std() / np.sqrt(len(vt)))
    out["rms_speed_kms"] = float(np.sqrt(1.5 * np.mean(vt**2)))
    out["gamma1pc_gcns_raw"] = flux
    out["gamma1pc_gcns_raw_staterr"] = flux_err
    print(f"GCNS systems <{R_MAIN:.0f} pc: N = {len(sel)}, n = {n_sys:.4f} /pc^3, <vT> = {mean_vT:.2f}, "
          f"<v> = {mean_v:.2f} km/s, Gamma(1pc) raw = {flux:.2f} +- {flux_err:.2f} /Myr")

    # sensitivity of the flux sum to the volume, outlier cap and merging
    sens = []
    for label, pp in (("table3+wide", prim), ("table3 only", prim_t3)):
        for R in [15, 20, 25, 30, 40, 50, 75, 100]:
            for cap in [300.0, 600.0, 1e9]:
                s = pp[(pp["dist"] < R) & (pp["vT"] <= cap)]
                sens.append(dict(merging=label, R_pc=R, vT_cap=cap, n_sys=len(s),
                                 dens=len(s) / sphere_volume(R),
                                 mean_v=4 / np.pi * s["vT"].mean(),
                                 mean_v_err=4 / np.pi * s["vT"].std() / np.sqrt(len(s)),
                                 gamma1pc=4 * s["vT"].sum() / sphere_volume(R) * KMS_TO_PCMYR))
    sens = pd.DataFrame(sens)
    sens.to_csv(os.path.join(DATA, "flux_sensitivity.csv"), index=False)
    print(sens[(sens.vT_cap == 600)].round(4).to_string(index=False))

    # ---- luminosity function and <v> by absolute magnitude -----------------
    G_LIM = 20.0
    edges = np.arange(-2, 22, 1.0)
    rv = load_rv_sample()
    m = prim.merge(rv[["source_id", "radial_velocity", "radial_velocity_error",
                       "rv_expected_sig_to_noise"]], on="source_id", how="left")
    # "clean" RV, the same definition as in the census (step 2): error < 5 km/s
    # and expected signal-to-noise >= 5 (revision: the S/N cut is now applied
    # to the velocity sample that feeds the settlement simulations as well)
    m["has_rv"] = (m["radial_velocity"].notna() & (m["radial_velocity_error"] < 5)
                   & (m["rv_expected_sig_to_noise"].fillna(0) >= 5.0))
    m["v3d"] = np.hypot(m["vT"], m["radial_velocity"])
    lf = []
    for lo, hi in zip(edges[:-1], edges[1:]):
        dlim = min(R_MAIN, 10 ** ((G_LIM - hi + 5) / 5))
        s = m[(m["MG"] >= lo) & (m["MG"] < hi) & (m["dist"] < dlim) & (m["vT"] <= VT_CAP)]
        if len(s) == 0:
            continue
        Vb = sphere_volume(dlim)
        r = s[s["has_rv"]]
        nr = s[~s["has_rv"]]
        lf.append(dict(MG_lo=lo, MG_hi=hi, d_lim=dlim, N=len(s), phi=len(s) / Vb,
                       phi_err=np.sqrt(len(s)) / Vb,
                       mean_v=4 / np.pi * s["vT"].mean(),
                       mean_v_err=4 / np.pi * s["vT"].std() / np.sqrt(len(s)) if len(s) > 1 else np.nan,
                       flux=4 * s["vT"].sum() / Vb * KMS_TO_PCMYR,
                       N_rv=len(r), f_rv=len(r) / len(s),
                       v_from_vT_rv=4 / np.pi * r["vT"].mean() if len(r) else np.nan,
                       v_from_vT_norv=4 / np.pi * nr["vT"].mean() if len(nr) else np.nan,
                       v3d_rv=r["v3d"].mean() if len(r) else np.nan))
    lf = pd.DataFrame(lf)
    lf.to_csv(os.path.join(DATA, "luminosity_function_kinematics.csv"), index=False)
    out["n_systems_from_LF"] = float(lf["phi"].sum())
    out["gamma1pc_from_LF"] = float(lf["flux"].sum())
    print(lf.round(4).to_string(index=False))

    # ---- checks with the 6D (RV) sample ------------------------------------
    inside = m[(m["dist"] < R_MAIN) & (m["vT"] <= VT_CAP)]
    r = inside[inside["has_rv"]]
    out["check_rv_subsample_N"] = int(len(r))
    out["check_v3d_mean_kms"] = float(r["v3d"].mean())
    out["check_v3d_mean_err"] = float(r["v3d"].std() / np.sqrt(len(r)))
    out["check_4overpi_vT_kms"] = float(4 / np.pi * r["vT"].mean())
    out["check_4overpi_vT_err"] = float(4 / np.pi * r["vT"].std() / np.sqrt(len(r)))
    # kinematic bias of RV availability: bright stars (G-limited cells where
    # most stars have RVs) with and without RV
    br = inside[(inside["MG"] >= 4) & (inside["MG"] < 12)]
    a, b = br[br["has_rv"]], br[~br["has_rv"]]
    out["bias_MG4to12_N_rv"] = int(len(a)); out["bias_MG4to12_N_norv"] = int(len(b))
    out["bias_MG4to12_v_rv"] = float(4 / np.pi * a["vT"].mean())
    out["bias_MG4to12_v_norv"] = float(4 / np.pi * b["vT"].mean())
    out["bias_MG4to12_v_norv_err"] = float(4 / np.pi * b["vT"].std() / np.sqrt(len(b)))
    print(f"identity check on {len(r)} RV stars: <v3d> = {r['v3d'].mean():.2f} +- "
          f"{out['check_v3d_mean_err']:.2f}, (4/pi)<vT> = {out['check_4overpi_vT_kms']:.2f} +- "
          f"{out['check_4overpi_vT_err']:.2f} km/s")
    print(f"RV-availability bias (4<=MG<12, <50 pc): with RV {out['bias_MG4to12_v_rv']:.1f} km/s "
          f"(N={len(a)}), without RV {out['bias_MG4to12_v_norv']:.1f} +- "
          f"{out['bias_MG4to12_v_norv_err']:.1f} km/s (N={len(b)})")

    # ---- adopted numbers --------------------------------------------------
    n_ad = n_sys * MISSING_CORR
    n_err = n_ad * np.hypot(MISSING_CORR_ERR / MISSING_CORR, 0.02)   # 2%: merging definition
    # <v>: scatter between volumes of 25..100 pc in flux_sensitivity.csv is ~1.5%
    sv = sens[(sens.merging == "table3+wide") & (sens.vT_cap == 600) & (sens.R_pc >= 25)]["mean_v"]
    v_ad = mean_v
    v_err = float(max(sv.max() - sv.min(), 2 * out["mean_speed_staterr_kms"]))
    gam = np.pi * n_ad * v_ad * KMS_TO_PCMYR
    gam_err_internal = gam * np.hypot(n_err / n_ad, v_err / v_ad)
    # Revision (referee point 2.2): the uncertainty of the adopted rate is the
    # half-range of the stellar-system estimates that differ only in defensible
    # choices (which sample, which pair merging, how many systems are missing),
    # not the internal error above.
    t3 = sens[(sens.merging == "table3 only") & (sens.vT_cap == 600) & (sens.R_pc == R_MAIN)]["gamma1pc"].iloc[0]
    corr_max = MISSING_CORR + MISSING_CORR_ERR
    variants = {
        "GCNS <50 pc, table 3 + wide-pair merging, no missing-system correction": flux,
        "GCNS <50 pc, table 3 merging only, no correction": float(t3),
        f"GCNS <50 pc, table 3 + wide pairs, x {corr_max:.2f} (missing as at 10 pc)": flux * corr_max,
        f"GCNS <50 pc, table 3 only, x {corr_max:.2f}": float(t3) * corr_max,
        "10 pc sample stellar systems x 50 pc mean speed": float(np.pi * tp["dens_systems_stellar"] * v_ad * KMS_TO_PCMYR),
        "10 pc sample stellar systems x their own mean 3D speed": tp["gamma1pc_tenpc_own_speed"],
    }
    vals = np.array(list(variants.values()))
    gam_err = float(0.5 * (vals.max() - vals.min()))
    pd.DataFrame(dict(variant=list(variants.keys()), gamma1pc=vals)).to_csv(
        os.path.join(DATA, "rate_variants.csv"), index=False)
    out.update(dict(missing_corr=MISSING_CORR, missing_corr_err=MISSING_CORR_ERR,
                    n_sys_adopted=n_ad, n_sys_adopted_err=n_err,
                    mean_speed_adopted_kms=v_ad, mean_speed_adopted_err=v_err,
                    gamma1pc_adopted=gam, gamma1pc_adopted_err=gam_err,
                    gamma1pc_internal_err=gam_err_internal,
                    gamma1pc_variant_min=float(vals.min()), gamma1pc_variant_max=float(vals.max())))
    print(pd.DataFrame(dict(variant=list(variants.keys()), gamma1pc=vals)).round(2).to_string(index=False))
    print(f"ADOPTED: n_sys = {n_ad:.4f} +- {n_err:.4f} /pc^3, <v> = {v_ad:.1f} +- {v_err:.1f} km/s,"
          f" Gamma(1 pc) = {gam:.2f} +- {gam_err:.2f} systems/Myr (internal error {gam_err_internal:.2f};"
          f" variants {vals.min():.2f}-{vals.max():.2f})")

    # ---- empirical 3D velocity sample for the settlement simulation --------
    r = r[r["MG"].notna()].copy()
    _, vel = galactic_xyz_uvw(r["RA_ICRS"], r["DE_ICRS"], r["Plx"], r["pmRA"], r["pmDE"],
                              r["radial_velocity"])
    r["U"], r["V"], r["W"] = vel[:, 0], vel[:, 1], vel[:, 2]
    allp = inside
    mg_edges = np.arange(0, 13)
    vt_edges = np.array([0, 5, 10, 15, 20, 25, 30, 40, 50, 60, 80, 100, 130, 170, 250, 1e9])
    mgb = lambda x: np.clip(np.floor(x.fillna(12)).astype(int), 0, 12)
    vtb = lambda x: np.digitize(x, vt_edges) - 1
    w = ipf_weights(mgb(r["MG"]).values, vtb(r["vT"].values), mgb(allp["MG"]).values,
                    vtb(allp["vT"].values))
    r["weight"] = w
    r[["source_id", "MG", "dist", "U", "V", "W", "weight"]].to_csv(
        os.path.join(DATA, "velocity_sample_50pc.csv"), index=False)
    out["velocity_sample_N"] = int(len(r))
    out["velocity_sample_Neff"] = float(1.0 / np.sum(w**2))
    for c in "UVW":
        mu = np.sum(w * r[c])
        out[f"mean_{c}_kms"] = float(mu)
        out[f"sigma_{c}_kms"] = float(np.sqrt(np.sum(w * (r[c] - mu)**2)))
    sp = np.sqrt(r["U"]**2 + r["V"]**2 + r["W"]**2).values
    out["velocity_sample_unweighted_mean_speed"] = float(sp.mean())
    out["velocity_sample_weighted_mean_speed"] = float(np.sum(w * sp))
    out["frac_speed_gt_100kms"] = float(np.sum(w * (sp > 100)))
    out["frac_speed_gt_200kms"] = float(np.sum(w * (sp > 200)))
    out["velocity_sample_N_speed_gt_200kms"] = int((sp > 200).sum())
    print("velocity sample: N =", len(r), " N_eff =", round(out["velocity_sample_Neff"]),
          " <v> unweighted =", round(sp.mean(), 2), " weighted =",
          round(out["velocity_sample_weighted_mean_speed"], 2), "km/s; ",
          {k: round(v, 2) for k, v in out.items() if k.startswith(("mean_U", "mean_V", "mean_W", "sigma_"))})

    save_json("s1_density_kinematics.json", out)
    return out


if __name__ == "__main__":
    main()