#!/usr/bin/env python3
"""Search for coherently dimmed multi-galaxy regions (spheres of radius R in redshift space).

Statistic.  Each galaxy gets a standardised FP residual zeta_i (I-direction residual, corrected for the
expected mean and variance of the magnitude-truncated distribution, winsorised at +-3.5, minus the smooth
survey-wide redshift trend of the mean).  For every sphere
of radius R centred on a sample galaxy and containing N >= N_MIN sample galaxies,
        S = sum(zeta) / sqrt(N)         (approximately N(0,1) if residuals are spatially uncorrelated)
A region dimmed by f gives every surviving member a shift log10(1-f)/sd_i and removes the members that
would fall out of the survey.

Limit.  Inject dimmed spheres at random positions in the real catalogue.
Rule A: an injection counts as detected if it would have produced a sphere with S below the most negative S
actually observed (for that R).  Since nothing more extreme than the observed minimum was observed, the 95%
Poisson limit on the number density of such regions is  n_95 = 3.0 / V_sens,  V_sens = V_draw * (detected
fraction), and the volume-filling fraction limit is  phi_95 = n_95 * (4/3) pi R^3.  The threshold of rule A is
set by the data, so rule A is not a rigorous frequentist limit.
Rule B (rigorous): fixed, pre-specified threshold S < T_FIX = -5; every independent observed sphere below it is
counted as a candidate (k_obs) and the limit is PoissonUL95(k_obs) / V_sens(T_FIX).

V_sens is the part of the survey volume that was actually tested (where a dimmed sphere would have been
detected); the filling-fraction limits extend it to the whole survey ONLY under the assumption that regions are
placed uniformly at random in volume.  The galaxy-centred injections give the complementary, galaxy-weighted
statement (fraction of sample galaxies that can lie inside such regions).
"""
import json
import os

import numpy as np
import pandas as pd
from scipy import sparse, stats
from scipy.spatial import cKDTree

import analysis_individual as A
import fpk_common as F

R_LIST = [5.0, 10.0, 20.0]          # Mpc (H0 = 70)
F_LIST = [0.3, 0.5, 0.75]
N_MIN = 3
T_FIX = -5.0           # alternative, pre-specified "5 sigma" threshold used with candidate counting
N_SHUFFLE = 300
N_INJECT_GAL = 30000                                  # galaxy-centred injections
N_INJECT_VOL = {5.0: 300000, 10.0: 100000, 20.0: 30000}   # volume-random injections (most small spheres are empty,
#                                                       so many more are needed for a usable Monte Carlo error)
WINSOR = 3.5
SEED = 20261011


def standardise(df):
    """zeta: residual standardised with the truncated-normal mean and variance implied by the survey limit."""
    z = df.dI.values / df.dI_sd.values
    a = (df.dI.values - 0.4 * (df.mlim.values - df.msel.values)) / df.dI_sd.values     # truncation point
    a = np.clip(a, -8, 6)
    lam = np.exp(stats.norm.logpdf(a) - stats.norm.logsf(a))
    var = np.clip(1 + a * lam - lam ** 2, 0.05, None)
    return lam, np.sqrt(var), a


def zeta_of(dI, sd, lam, sv):
    return np.clip((dI / sd - lam) / sv, -WINSOR, WINSOR)


def z_trend(zdist, zeta, nbin=12):
    """Smooth survey-wide redshift trend of the standardised residual (mean in quantile bins, interpolated).
    A zero-point drift with redshift (aperture/k/evolution corrections, curvature of the FP combined with the
    changing luminosity range) is a property of whole redshift shells holding ~10^4 galaxies; it cannot be
    caused by regions <= 20 Mpc in radius, and is removed before the sphere scan."""
    qb = np.quantile(zdist, np.linspace(0, 1, nbin + 1))
    xm, ym = [], []
    for a, b in zip(qb[:-1], qb[1:]):
        m = (zdist >= a) & (zdist <= b)
        xm.append(zdist[m].mean())
        ym.append(zeta[m].mean())
    return np.interp(zdist, xm, ym), (xm, ym)


def adjacency(xyz, R):
    tree = cKDTree(xyz)
    pairs = tree.query_pairs(R, output_type="ndarray")
    n = len(xyz)
    i = np.concatenate([pairs[:, 0], pairs[:, 1], np.arange(n)])
    j = np.concatenate([pairs[:, 1], pairs[:, 0], np.arange(n)])
    Aij = sparse.csr_matrix((np.ones(len(i), np.float32), (i, j)), shape=(n, n))
    return tree, Aij


def greedy_top(S, N, xyz, R, k=5, sign=-1):
    order = np.argsort(sign * -S) if sign > 0 else np.argsort(S)
    chosen = []
    for idx in order:
        if not np.isfinite(S[idx]):
            continue
        if all(np.linalg.norm(xyz[idx] - xyz[c]) > 2 * R for c in chosen):
            chosen.append(idx)
        if len(chosen) >= k:
            break
    return chosen


def coherent_rms(zeta, xyz, R, sd_dex):
    """Excess (spatially coherent) rms of the mean residual in non-overlapping cubic cells of side
    (4 pi/3)^(1/3) R (same volume as a sphere of radius R).  Returns dex, with bootstrap error."""
    side = (4 * np.pi / 3) ** (1 / 3) * R
    cell = np.floor(xyz / side).astype(np.int64)
    key = (cell[:, 0] * 1000003 + cell[:, 1]) * 1000003 + cell[:, 2]
    u, inv, cnt = np.unique(key, return_inverse=True, return_counts=True)
    sums = np.bincount(inv, weights=zeta)
    sums2 = np.bincount(inv, weights=zeta ** 2)
    m = cnt >= 2
    n = cnt[m].astype(float)
    mean = sums[m] / n
    # unbiased per-cell estimate of the coherent variance: mean^2 - s^2/n, with the global unit variance
    s2 = np.var(zeta)
    est = mean ** 2 - s2 / n
    w = n                                              # galaxy-weighted
    val = np.sum(w * est) / np.sum(w)
    rng = np.random.default_rng(SEED)
    bs = []
    for _ in range(300):
        k = rng.integers(0, len(est), len(est))
        bs.append(np.sum(w[k] * est[k]) / np.sum(w[k]))
    err = np.std(bs)
    to_dex = lambda v: np.sign(v) * np.sqrt(abs(v)) * sd_dex
    return dict(var_zeta=float(val), var_zeta_err=float(err), rms_dex=float(to_dex(val)),
                rms_dex_hi95=float(to_dex(val + 1.645 * err)), n_cells=int(m.sum()))


def env_trend(xyz, zdist, zeta):
    """Mean standardised residual as a function of local galaxy over-density (neighbours within 5 Mpc,
    normalised by the mean count at that redshift).  Returns the trend evaluated for every galaxy."""
    tree = cKDTree(xyz)
    n5 = np.array([len(v) for v in tree.query_ball_point(xyz, 5.0)], float)
    zb = np.quantile(zdist, np.linspace(0, 1, 11))
    k = np.clip(np.searchsorted(zb, zdist, side="right") - 1, 0, 9)
    nbar = np.array([n5[k == i].mean() for i in range(10)])
    od = np.log10(n5 / nbar[k])
    qb = np.quantile(od, np.linspace(0, 1, 13))
    xm, ym = [], []
    for a, b in zip(qb[:-1], qb[1:]):
        m = (od >= a) & (od <= b)
        if m.sum() > 20:
            xm.append(od[m].mean())
            ym.append(zeta[m].mean())
    return np.interp(od, xm, ym), od, (xm, ym)


def photometry_crosscheck(df):
    """SDSS only.  (i) residual of the r(SDSS deV) - J(2MASS XSC total) colour after removing its dependence on
    sigma, z, size, g-r and magnitude: if SDSS photometry were biased faint in a region, the colour residual
    would be positive by the same amount; (ii) the log-distance ratio published by Howlett+22 for the same
    galaxies (their own FP fit, group redshifts)."""
    if "devmag_r" not in df:
        return None
    import vet_candidates as V
    x = V.xmatch(df, "vizier:VII/233/xsc", os.path.join(F.RAW, "xmatch_2massxsc_sdss.csv"), radius=5.0)
    x = x.sort_values("angDist").drop_duplicates("id")
    j = df[["id"]].merge(x[["id", "J.ext"]], on="id", how="left")["J.ext"].values
    col = (df.devmag_r.values - df.ext_r.values) - j
    z, ls, lt = df.zdist.values - 0.08, df.logs.values - 2.25, np.log10(df.theta_e.values)
    X = np.column_stack([np.ones(len(df)), ls, ls ** 2, z, z ** 2, lt, lt ** 2, df.gr.values, df.msel.values - 16.5])
    m = np.isfinite(col) & np.all(np.isfinite(X), axis=1)
    for _ in range(4):
        c = np.linalg.lstsq(X[m], col[m], rcond=None)[0]
        res = col - X @ c
        m = np.isfinite(res) & (np.abs(res) < 3 * F.mad_sigma(res[m]))
    h = pd.read_csv(os.path.join(F.DATA, "sample_howlett.csv"))[["id", "logdist"]]
    ld = df[["id"]].merge(h, on="id", how="left")["logdist"].values
    return dict(dcolJ=res, logdist=ld)


def stellar_pop_prediction(df):
    """SDSS only.  Residuals of D4000 and g-r at fixed (sigma, z), and the part of the FP residual that is
    linearly predictable from them (older/redder populations are fainter at fixed sigma and R_e).
    Grey dimming would change neither D4000 nor colour."""
    if "d4000" not in df:
        return None
    ls, z = df.logs.values - 2.25, df.zdist.values - np.median(df.zdist.values)
    B = np.column_stack([np.ones(len(df)), ls, ls ** 2, z, z ** 2])
    out = {}
    for nm, col in (("dD4000", "d4000"), ("dgr", "gr")):
        v = df[col].values.astype(float)
        m = np.isfinite(v)
        c = np.linalg.lstsq(B[m], v[m], rcond=None)[0]
        out[nm] = v - B @ c
    P = np.column_stack([np.ones(len(df)), out["dD4000"], out["dgr"]])
    m = np.all(np.isfinite(P), axis=1) & (np.abs(df.dI.values) < 0.35)
    c = np.linalg.lstsq(P[m], df.dI.values[m], rcond=None)[0]
    out["pred"] = P @ c
    out["coef"] = c
    out["dls"] = df.logs.values - df.logs_port.values
    return out


def run(name, df, cfg, zmin, zmax, rng, detrend=False):
    df = df[(df.zdist >= zmin) & (df.zdist <= zmax)].reset_index(drop=True)
    lam, sv, a = standardise(df)
    zeta = zeta_of(df.dI.values, df.dI_sd.values, lam, sv)
    ztr, ztr_curve = z_trend(df.zdist.values, zeta)
    zeta = zeta - ztr
    sd_dex = float(np.median(df.dI_sd.values * sv))
    dist = F.comoving_distance(df.zdist.values)
    xyz = F.radec_to_xyz(df.ra.values, df["dec"].values, dist)
    spop = stellar_pop_prediction(df)
    phot = photometry_crosscheck(df)
    trend, od, trend_curve = env_trend(xyz, df.zdist.values, zeta)
    if not detrend:
        trend = np.zeros(len(df))
    zeta = zeta - trend
    # bounding region for random centres
    ra, dec = df.ra.values, df["dec"].values
    sd_lo, sd_hi = np.sin(np.radians(dec.min())), np.sin(np.radians(dec.max()))
    d_lo, d_hi = F.comoving_distance(zmin), F.comoving_distance(zmax)
    v_draw = (2 * np.pi) * (sd_hi - sd_lo) * (d_hi ** 3 - d_lo ** 3) / 3.0   # full RA range, dec band
    summary = dict(sample=name, detrend=bool(detrend), n_gal=int(len(df)), zmin=zmin, zmax=zmax,
                   v_draw_mpc3=float(v_draw), zeta_mean=float(zeta.mean()), zeta_std=float(zeta.std()), sd_dex=sd_dex,
                   z_trend_z=[float(v) for v in ztr_curve[0]], z_trend_mean_zeta=[float(v) for v in ztr_curve[1]],
                   spop_coef_const_dD4000_dgr=[float(v) for v in spop["coef"]] if spop else None,
                   env_trend_logoverdensity=[float(v) for v in trend_curve[0]],
                   env_trend_mean_zeta=[float(v) for v in trend_curve[1]])
    rows, patches, coh = [], [], []
    # pre-compute the dimmed twins
    twin = {}
    for f in F_LIST:
        ok, g = A.survive(df, f, cfg)
        twin[f] = (ok, zeta_of(df.dI.values + np.log10(1 - f), df.dI_sd.values, lam, sv) - ztr - trend)
        summary["survive_frac_f%.2f" % f] = float(ok.mean())
    for R in R_LIST:
        tree, Aij = adjacency(xyz, R)
        N = np.asarray(Aij.sum(axis=1)).ravel()
        sumz = Aij @ zeta
        good = N >= N_MIN
        S = np.where(good, sumz / np.sqrt(np.maximum(N, 1)), np.nan)
        s_min, s_max = float(np.nanmin(S)), float(np.nanmax(S))
        # number of independent observed spheres below the fixed threshold (greedy, centres > 2R apart)
        below = np.where(np.nan_to_num(S, nan=0.0) < T_FIX)[0]
        below = below[np.argsort(S[below])]
        indep = []
        for idx in below:
            if all(np.linalg.norm(xyz[idx] - xyz[c]) > 2 * R for c in indep):
                indep.append(idx)
        k_fix = len(indep)
        # shuffle null (destroys all spatial coherence, keeps positions and the one-point distribution)
        null_min, null_max, null_std = [], [], []
        for _ in range(N_SHUFFLE):
            zs = rng.permutation(zeta)
            Ss = (Aij @ zs)[good] / np.sqrt(N[good])
            null_min.append(Ss.min())
            null_max.append(Ss.max())
            null_std.append(Ss.std())
        null_min, null_max = np.array(null_min), np.array(null_max)
        p_min = float((np.sum(null_min <= s_min) + 1) / (N_SHUFFLE + 1))
        p_max = float((np.sum(null_max >= s_max) + 1) / (N_SHUFFLE + 1))
        c = coherent_rms(zeta, xyz, R, sd_dex)
        c.update(sample=name, R_mpc=R, S_std_obs=float(np.nanstd(S)), S_std_shuffle=float(np.mean(null_std)))
        coh.append(c)
        for sign, lab in ((-1, "most negative (dim)"), (1, "most positive (bright)")):
            for idx in greedy_top(S, N, xyz, R, k=5, sign=sign):
                nb = Aij[idx].indices
                patches.append(dict(sample=name, R_mpc=R, kind=lab, S=float(S[idx]), N=int(N[idx]),
                                    mean_dI_dex=float(df.dI.values[nb].mean()),
                                    mean_dI_corrected_dex=float(zeta[nb].mean() * sd_dex),
                                    equiv_dimming_f=float(1 - 10 ** df.dI.values[nb].mean()),
                                    ra=float(df.ra[idx]), dec=float(df["dec"][idx]), z=float(df.zdist[idx]),
                                    median_ext_r=float(np.median(df.loc[nb, "ext_r"])) if "ext_r" in df else np.nan,
                                    n_plates=int(df.loc[nb, "plate"].nunique()) if "plate" in df else -1,
                                    mean_overdensity_log=float(od[nb].mean()),
                                    mean_dD4000=float(np.nanmean(spop["dD4000"][nb])) if spop else np.nan,
                                    sem_dD4000=float(np.nanstd(spop["dD4000"][nb]) / np.sqrt(len(nb))) if spop else np.nan,
                                    dI_predicted_from_D4000_colour=float(np.nanmean(spop["pred"][nb])) if spop else np.nan,
                                    mean_dlogsigma_two_codes=float(np.nanmean(spop["dls"][nb])) if spop else np.nan,
                                    frac_in_howlett=float(df.loc[nb, "in_howlett"].mean()) if "in_howlett" in df else np.nan,
                                    n_2mass=int(np.isfinite(phot["dcolJ"][nb]).sum()) if phot else 0,
                                    mean_r_minus_J_resid_mag=float(np.nanmean(phot["dcolJ"][nb])) if phot else np.nan,
                                    sem_r_minus_J_resid_mag=float(np.nanstd(phot["dcolJ"][nb]) / np.sqrt(
                                        max(np.isfinite(phot["dcolJ"][nb]).sum(), 1))) if phot else np.nan,
                                    howlett_mean_logdist=float(np.nanmean(phot["logdist"][nb])) if phot and np.isfinite(
                                        phot["logdist"][nb]).any() else np.nan,
                                    mean_dFJ_dex=float(df.dFJ.values[nb].mean())))
        # ---------------- injections
        n_v, n_g = N_INJECT_VOL[R], N_INJECT_GAL
        u = rng.random((n_v, 3))
        r = (d_lo ** 3 + u[:, 0] * (d_hi ** 3 - d_lo ** 3)) ** (1 / 3)
        rr = np.degrees(2 * np.pi * u[:, 1])
        dd = np.degrees(np.arcsin(sd_lo + u[:, 2] * (sd_hi - sd_lo)))
        cen_vol = F.radec_to_xyz(rr, dd, r)
        z_vol = np.interp(r, F._dc, F._zg)
        i_gal = rng.integers(0, len(xyz), n_g)
        cen_gal = xyz[i_gal]
        z_gal = df.zdist.values[i_gal]
        vR = 4 * np.pi / 3 * R ** 3
        ul5 = float(F.poisson_ul(k_fix))
        for mode, cen, zcen in (("volume", cen_vol, z_vol), ("galaxy", cen_gal, z_gal)):
            n_c = len(cen)
            G_all = tree.query_ball_point(cen, R)
            C_all = tree.query_ball_point(cen, 2 * R)
            n_in = np.array([len(g) for g in G_all])
            for f in F_LIST:
                ok, zt = twin[f]
                det = np.zeros(n_c, bool)
                det5 = np.zeros(n_c, bool)
                n_surv = np.zeros(n_c, int)
                for k in range(n_c):
                    G = np.asarray(G_all[k], dtype=int)
                    if len(G) == 0:
                        continue
                    alive = ok[G]
                    n_surv[k] = alive.sum()
                    if n_surv[k] < N_MIN:
                        continue
                    C = np.asarray(C_all[k], dtype=int)
                    sub = Aij[C][:, G]
                    dz = np.where(alive, zt[G], 0.0) - zeta[G]
                    dn = alive.astype(float) - 1.0
                    newN = N[C] + sub @ dn
                    newS = (sumz[C] + sub @ dz) / np.sqrt(np.maximum(newN, 1))
                    # dimmed-out centres cannot serve as centres
                    cen_alive = ~np.isin(C, G[~alive])
                    valid = (newN >= N_MIN) & cen_alive
                    if valid.any():
                        m = newS[valid].min()
                        det[k] = m < s_min
                        # fixed threshold: the injection itself must push a sphere below T_FIX
                        pre = np.nanmin(np.where(N[C] >= N_MIN, sumz[C] / np.sqrt(np.maximum(N[C], 1)), np.nan))
                        det5[k] = (m < T_FIX) and not (pre < T_FIX)
                nA, nB = int(det.sum()), int(det5.sum())
                frac, frac5 = nA / n_c, nB / n_c
                err = np.sqrt(max(frac * (1 - frac), 1e-12) / n_c)
                zq = lambda mk, q: float(np.percentile(zcen[mk], q)) if mk.any() else np.nan
                row = dict(sample=name, R_mpc=R, f=f, mode=mode, n_inject=n_c, S_min_obs=s_min, p_shuffle_min=p_min,
                           S_max_obs=s_max, p_shuffle_max=p_max,
                           n_detected=nA, frac_detected=float(frac), frac_err=float(err),
                           T_fix=T_FIX, k_obs_fix=k_fix, n_detected_fix=nB, frac_detected_fix=float(frac5),
                           # Monte Carlo (binomial) relative 1-sigma error of each efficiency = of each limit
                           mc_relerr=float(1 / np.sqrt(nA)) if nA > 0 else np.nan,
                           mc_relerr_fix=float(1 / np.sqrt(nB)) if nB > 0 else np.nan,
                           z_det_median=zq(det, 50), z_det_p90=zq(det, 90),
                           z_det_fix_median=zq(det5, 50), z_det_fix_p90=zq(det5, 90),
                           mean_n_in=float(n_in.mean()), frac_nsurv_ge_min=float((n_surv >= N_MIN).mean()))
                if mode == "volume":
                    v_sens = v_draw * frac
                    v5 = v_draw * frac5
                    row.update(v_sens_fix_mpc3=float(v5),
                               n95_fix_per_gpc3=float(ul5 / v5 * 1e9) if v5 > 0 else np.nan,
                               fill95_fix=float(ul5 * vR / v5) if v5 > 0 else np.nan)
                    row.update(v_sens_mpc3=float(v_sens),
                               n95_per_mpc3=float(3.0 / v_sens) if v_sens > 0 else np.nan,
                               n95_per_gpc3=float(3.0 / v_sens * 1e9) if v_sens > 0 else np.nan,
                               fill95=float(3.0 * vR / v_sens) if v_sens > 0 else np.nan)
                else:
                    # galaxy-weighted placement: regions centred on a random subset of the sample galaxies.
                    # At most UL/(p * N_gal) of the sample galaxies can be centres of such regions, and (union
                    # bound) at most that times the mean membership of a galaxy-centred sphere can lie inside one.
                    ng = float(len(df))
                    row.update(frac_gal_centres95=float(3.0 / (frac * ng)) if frac > 0 else np.nan,
                               frac_gal_centres95_fix=float(ul5 / (frac5 * ng)) if frac5 > 0 else np.nan,
                               frac_gal_inside95=float(3.0 * n_in.mean() / (frac * ng)) if frac > 0 else np.nan,
                               frac_gal_inside95_fix=float(ul5 * n_in.mean() / (frac5 * ng)) if frac5 > 0 else np.nan)
                rows.append(row)
        print(name, "detrend" if detrend else "raw", "R=%g done: S_min %.2f (p=%.3f)  S_max %.2f (p=%.3f)" % (R, s_min, p_min, s_max, p_max),
              flush=True)
    lab = "env-detrended" if detrend else "raw"
    rows, patches, coh = pd.DataFrame(rows), pd.DataFrame(patches), pd.DataFrame(coh)
    for t in (rows, patches, coh):
        t.insert(1, "variant", lab)
    return summary, rows, patches, coh, dict(xyz=xyz, zeta=zeta, df=df)


def main():
    rng = np.random.default_rng(SEED)
    S = A.load_samples()
    out_rows, out_patches, out_coh, summaries, keep = [], [], [], [], {}
    for name, fn, zmin, zmax in (("sdss", "resid_sdss.csv", 0.02, 0.12),
                                 ("6df", "resid_6df.csv", 3000 / F.C_KMS, 16120 / F.C_KMS)):
        df = pd.read_csv(F.resid_path(name))
        if name == "6df":
            df["mfib"] = A.fibre_mag_6df(df)
        else:
            df["mfib"] = df.fibermag_r
        cfg = S[name][1]
        for detrend in (False, True):
            summ, rows, patches, coh, extra = run(name, df, cfg, zmin, zmax, rng, detrend=detrend)
            summaries.append(summ)
            out_rows.append(rows)
            out_patches.append(patches)
            out_coh.append(coh)
            if not detrend:
                keep[name] = extra
                pd.DataFrame({"id": extra["df"]["id"].values, "zeta": extra["zeta"]}).to_csv(
                    os.path.join(F.DATA, "zeta_%s.csv" % name), index=False, float_format="%.5g")
    rows = pd.concat(out_rows, ignore_index=True)
    rows.to_csv(os.path.join(F.DATA, "limits_regions.csv"), index=False, float_format="%.6g")
    pd.concat(out_patches, ignore_index=True).to_csv(os.path.join(F.DATA, "region_extreme_patches.csv"), index=False,
                                                    float_format="%.6g")
    pd.concat(out_coh, ignore_index=True).to_csv(os.path.join(F.DATA, "region_coherent_rms.csv"), index=False,
                                                 float_format="%.6g")
    with open(os.path.join(F.DATA, "region_summary.json"), "w") as fh:
        json.dump(summaries, fh, indent=1)
    pd.set_option("display.width", 250)
    print(rows[rows["mode"] == "volume"][["sample", "variant", "R_mpc", "f", "S_min_obs", "p_shuffle_min", "n_detected",
                                         "v_sens_mpc3", "n95_per_gpc3", "fill95", "k_obs_fix", "n_detected_fix",
                                         "fill95_fix", "z_det_p90", "z_det_fix_p90"]].to_string(index=False))
    print(rows[rows["mode"] == "galaxy"][["sample", "variant", "R_mpc", "f", "frac_detected", "frac_detected_fix",
                                         "mean_n_in", "frac_gal_centres95", "frac_gal_inside95",
                                         "frac_gal_inside95_fix"]].to_string(index=False))
    print(pd.concat(out_coh)[["sample", "variant", "R_mpc", "rms_dex", "rms_dex_hi95", "S_std_obs", "S_std_shuffle",
                              "n_cells"]].to_string(index=False))
    return keep


if __name__ == "__main__":
    main()