"""End-to-end injection test of the recovery R (final cut set only).

The recovery R used in the limits is "analytic": the grey shift log10(1-f) is added to the three
statistics (D, m_ms, m_cmd) of every sensitive star and the flag is re-applied.  Here the dimming is put
into the catalogue instead: the G and Ks magnitudes of a random fraction P_INJ of the clean stars are
made fainter by -2.5 log10(1-f) (colours, parallax and spectroscopic parameters untouched), the whole
calibration is redone with the dimmed stars inside it, and the flags are recomputed.

For ONE injected star the two routes agree exactly, because every reference a star is compared with
(local median, local scatter, main-sequence floors) is leave-one-out.  They differ only because injected
stars contaminate each other's reference sets, so the shortfall grows with the injected fraction.  The
correction factor

    C(survey, f) = R_end_to_end / R_analytic     measured at P_INJ = 2e-3

is applied to R in limits.py.  2e-3 is at or above the largest limit quoted in the report, so C is the
right size for the weakest limit and conservative for the others.  A second set of runs at higher
injected fractions (f = 0.25 only) documents the scaling with density.

Writes data/injection_runs.csv (one row per run and level) and data/injection_end_to_end.csv (summed).
"""
import os
import warnings

import numpy as np
import pandas as pd

import dimlib as L
from common import DATA_MAIN, WORK, dpkl, log
from robustness import NEEDED, flags_full, prepare

warnings.filterwarnings("ignore")
SURVEYS = ("apogee", "lamost", "gspspec")
P_INJ = 2e-3
GROUPS = ((0.25,), (0.5,), (0.75, 0.9))                      # levels injected together in one run
NREAL = dict(apogee=4, lamost=1, gspspec=1)                  # realisations per group (small samples need more)
SCALING = dict(apogee=((8e-3, 3), (2e-2, 3)), lamost=((8e-3, 1),), gspspec=())   # (P_INJ, realisations), f = 0.25
MIN_REC = 30                                                 # analytic recoveries needed for a per-survey C


def one_run(dd, sv, base, levels, p_inj, rng):
    clean = np.flatnonzero(base.clean.values)
    n = max(1, int(round(p_inj * len(clean))))
    pick = rng.choice(clean, size=n * len(levels), replace=False)
    cols = [c for c in ("MG", "MK", "G", "K") if c in dd.columns]
    saved = dd.loc[pick, cols].values.copy()
    sets = {}
    for j, f in enumerate(levels):
        ii = pick[j * n:(j + 1) * n]
        dd.loc[ii, cols] = dd.loc[ii, cols].values - 2.5 * np.log10(1.0 - f)
        sets[f] = ii
    s2 = prepare(dd, sv, precut=True)
    dd.loc[pick, cols] = saved                                # restore the catalogue
    rows = []
    for f, ii in sets.items():
        sens0, rec0 = L.flag_mask(base, f, inject=True)
        sens0, rec0 = sens0 & base.clean.values, rec0 & base.clean.values
        _, fl2 = flags_full(s2, f)
        a, b = rec0[ii], fl2[ii] & sens0[ii]
        rows.append(dict(survey=sv, f=f, p_inj=p_inj, n_injected=len(ii), n_sensitive=int(sens0[ii].sum()),
                         n_rec_analytic=int(a.sum()), n_rec_end_to_end=int(b.sum()),
                         n_lost=int((a & ~b).sum()), n_gained=int((b & ~a).sum())))
    return rows


def run(surveys=SURVEYS):
    rows = []
    for sv in surveys:
        d = pd.read_pickle(os.path.join(WORK, "%s_stars.pkl" % sv))
        m, _ = L.apply_cuts(d, sv, verbose=False)
        dd = d.loc[m, [c for c in NEEDED if c in d.columns]].reset_index(drop=True)
        del d
        base = pd.read_pickle(dpkl(sv, "final"))[["source_id", "D", "sigma_i", "m_ms", "m_cmd", "clean"]]
        assert np.array_equal(dd.source_id.values, base.source_id.values)
        rng = np.random.default_rng(20261012 + SURVEYS.index(sv))
        for r in range(NREAL[sv]):
            for levels in GROUPS:
                log("injection: %s  f = %s  P_INJ = %g  realisation %d" % (sv, levels, P_INJ, r + 1))
                rows += one_run(dd, sv, base, levels, P_INJ, rng)
        for p_inj, nreal in SCALING[sv]:
            for r in range(nreal):
                log("injection: %s  f = 0.25  P_INJ = %g  realisation %d (density scaling)" % (sv, p_inj, r + 1))
                rows += one_run(dd, sv, base, (0.25,), p_inj, rng)
        del dd, base
    runs = pd.DataFrame(rows)
    runs.to_csv(os.path.join(DATA_MAIN, "injection_runs.csv"), index=False)
    g = runs.groupby(["survey", "f", "p_inj"], sort=False)[
        ["n_injected", "n_sensitive", "n_rec_analytic", "n_rec_end_to_end", "n_lost", "n_gained"]].sum().reset_index()
    g["R_analytic"] = g.n_rec_analytic / g.n_sensitive.clip(lower=1)
    g["R_end_to_end"] = g.n_rec_end_to_end / g.n_sensitive.clip(lower=1)
    g["C"] = g.n_rec_end_to_end / g.n_rec_analytic.clip(lower=1)
    g["C_err"] = np.sqrt(g.n_lost + g.n_gained) / g.n_rec_analytic.clip(lower=1)      # paired comparison
    # pooled over surveys at each (f, p_inj); used when a survey has too few sensitive injected stars
    pool = g.groupby(["f", "p_inj"])[["n_rec_analytic", "n_rec_end_to_end"]].sum()
    pool = (pool.n_rec_end_to_end / pool.n_rec_analytic.clip(lower=1)).rename("C_pooled").reset_index()
    g = g.merge(pool, on=["f", "p_inj"], how="left")
    g["C_used"] = np.minimum(1.0, np.where(g.n_rec_analytic >= MIN_REC, g.C, g.C_pooled))
    g.to_csv(os.path.join(DATA_MAIN, "injection_end_to_end.csv"), index=False)
    pd.set_option("display.width", 250)
    print(g.to_string(index=False, float_format=lambda x: "%.4g" % x))
    return g


def correction(survey, f):
    """C(survey, f) at P_INJ for limits.py and robustness.py."""
    g = pd.read_csv(os.path.join(DATA_MAIN, "injection_end_to_end.csv"))
    r = g[(g.survey == survey) & np.isclose(g.f, f) & np.isclose(g.p_inj, P_INJ)]
    return float(r.C_used.iloc[0])


if __name__ == "__main__":
    import sys
    run(tuple(sys.argv[1:]) or SURVEYS)