"""Print (and save to data/report_numbers.txt) every number quoted in the report, from the CSV products.

Both astrometric cut sets are read: data/ ("final": RUWE < 1.2, parallax/error >= 50; tightened after a
first look at the tail) and data/origcuts/ ("orig": RUWE < 1.4, parallax/error >= 20; pre-declared).
Also writes data/headline_limits.csv (the limits table of the report) and data/warm_swarm_limit.csv.
"""
import io
import os

import numpy as np
import pandas as pd

import dimlib as L
from common import DATA_MAIN, DATA_ORIG

SURVEYS = ("apogee", "lamost", "gspspec")
SETS = (("final", DATA_MAIN), ("orig", DATA_ORIG))
pd.set_option("display.width", 250)
pd.set_option("display.max_columns", 60)
pd.set_option("display.max_rows", 900)
pd.set_option("display.max_colwidth", 90)


def rd(name, root=DATA_MAIN):
    return pd.read_csv(os.path.join(root, name))


def unresolved_table(root):
    u = rd("candidates_unresolved.csv", root)
    u["f_max"] = 0.0
    for f in L.F_GRID:
        u.loc[u["flag_%g" % f] & (u["status_%g" % f] == "unresolved"), "f_max"] = f
    return u


def main():
    buf = io.StringIO()

    def P(*a):
        print(*a)
        print(*a, file=buf)

    def g3(x):
        return "%.3g" % x

    P("=" * 30, "1. cut-by-cut star counts")
    for cs, root in SETS:
        for sv in SURVEYS:
            P("-- [%s cuts]" % cs, sv)
            P(rd("cuts_%s.csv" % sv, root).to_string(index=False))

    for cs, root in SETS:
        P("-- [%s cuts] combined sample (unique stars, best survey per star)" % cs)
        P(rd("combined_sample_stats.csv", root).T.to_string(header=False, float_format=g3))

    P("\n" + "=" * 30, "2. core scatter of D and tails (clean samples, final cuts)")
    for sv in SURVEYS:
        t = rd("core_tail_%s.csv" % sv)
        P(t[["survey", "subset", "n", "sigma_core", "sigma_faint", "median_sigma_i", "frac_lt_m0p2", "n_lt_m0p3",
             "n_lt_m0p6", "n_lt_m1"]].to_string(index=False, float_format=g3))

    P("\n" + "=" * 30, "3. sensitive stars, flagged stars and analytic injection recovery, by class (final cuts)")
    for sv in SURVEYS:
        t = rd("limits_raw_%s.csv" % sv)
        P(t.to_string(index=False, float_format=g3))

    P("\n" + "=" * 30, "4. modelled selection efficiency eps(f) per survey and Teff class")
    for cs, root in SETS:
        P("-- [%s cuts]" % cs)
        P(rd("selection_eff.csv", root).to_string(index=False, float_format=g3))
    P("\n-- model-free twin cross-check of eps (APOGEE, LAMOST); eps_model is the parent-sample model")
    for cs, root in SETS:
        tw = rd("selection_eff_twin.csv", root)
        a = tw[tw.cls == "all"].pivot_table(index=["survey", "model"], columns="f", values="eps_twin", sort=False)
        P("-- [%s cuts] all classes" % cs)
        P(a.to_string(float_format=g3))
        t0 = tw[(tw.model == tw.model.iloc[0]) & (tw.cls != "all")]
        P(t0[["survey", "cls", "f", "eps_twin", "eps_model", "n_twins", "frac_weight_beyond_1kpc"]].to_string(index=False, float_format=g3))

    P("\n" + "=" * 30, "5. end-to-end injection: correction C = R_end_to_end / R_analytic")
    P(rd("injection_end_to_end.csv").to_string(index=False, float_format=lambda x: "%.4g" % x))

    P("\n" + "=" * 30, "6. vetting summary and vetting efficiency")
    for cs, root in SETS:
        P("-- [%s cuts]" % cs)
        P(rd("vetting_summary.csv", root).to_string(index=False))
        P(rd("vetting_efficiency.csv", root).to_string(index=False, float_format=g3))

    P("\n" + "=" * 30, "7. limits")
    lims = {}
    for cs, root in SETS:
        lims[cs] = rd("limits.csv", root)
        P("-- [%s cuts]" % cs)
        P(lims[cs].to_string(index=False, float_format=g3))
    P("\n-- HEADLINE: combined sample, three-way criterion, both cut sets (also data/headline_limits.csv)")
    hl = []
    for f in L.F_GRID:
        row = dict(f=f)
        for cs in ("final", "orig"):
            r = lims[cs][(lims[cs]["sample"] == "combined") & (lims[cs].tier == "three-way") & (lims[cs].f == f)].iloc[0]
            for k_, v_ in (("n_clean", r.n_clean), ("n_sens", r.n_sens), ("R", r.mean_R), ("C", r.mean_C), ("eps", r.mean_eps),
                           ("eps_twin", r.mean_eps_twin), ("V", r.mean_V), ("n_flagged", r.n_flagged),
                           ("n_unresolved", r.n_unresolved), ("UL95_sample", r.UL95_sample_vetted),
                           ("UL95_pop", r.UL95_pop_vetted), ("UL95_pop_eps_lo", r.UL95_pop_vetted_eps_sysmax),
                           ("UL95_pop_eps_hi", r.UL95_pop_vetted_eps_sysmin), ("UL95_pop_twin", r.UL95_pop_vetted_eps_twin),
                           ("UL95_pop_twin_lo", r.UL95_pop_vetted_eps_twin_max), ("UL95_pop_twin_hi", r.UL95_pop_vetted_eps_twin_min),
                           ("UL95_pop_without_C", r.UL95_pop_vetted_without_C), ("UL95_pop_zero_cand", r.UL95_pop_zero_candidates)):
                row["%s_%s" % (k_, cs)] = v_
        row["UL95_pop_low"] = min(row["UL95_pop_final"], row["UL95_pop_orig"])
        row["UL95_pop_high"] = max(row["UL95_pop_final"], row["UL95_pop_orig"])
        row["ratio_orig_over_final"] = row["UL95_pop_orig"] / row["UL95_pop_final"]
        hl.append(row)
    hl = pd.DataFrame(hl)
    hl.to_csv(os.path.join(DATA_MAIN, "headline_limits.csv"), index=False)
    P(hl.T.to_string(float_format=g3))
    for cs in ("final", "orig"):
        t = lims[cs][(lims[cs].tier == "three-way")]
        P("-- per-sample population limits [%s cuts] (unresolved candidates)" % cs)
        piv = t.assign(x=t.UL95_pop_vetted.map(g3) + " (" + t.n_unresolved.astype(str) + ")").pivot_table(
            index="sample", columns="f", values="x", aggfunc="first", sort=False)
        P(piv.to_string())
        t = lims[cs][(lims[cs].tier == "spectroscopic-only")]
        P("-- spectroscopic criterion only [%s cuts]: flagged / unresolved / limit" % cs)
        piv = t.assign(x=t.n_flagged.astype(str) + " / " + t.n_unresolved.astype(str) + " / " + t.UL95_pop_vetted.map(g3)).pivot_table(
            index="sample", columns="f", values="x", aggfunc="first", sort=False)
        P(piv.to_string())

    P("\n" + "=" * 30, "8. which stars carry the limit (combined sample, three-way): exposure shares")
    for cs, root in SETS:
        P("-- [%s cuts]" % cs)
        P(rd("limits_scope.csv", root).to_string(index=False, float_format=g3))

    P("\n" + "=" * 30, "9. unresolved three-way candidates")
    for cs, root in SETS:
        u = unresolved_table(root)
        P("-- [%s cuts] unresolved rows %d, distinct stars %d; by survey: %s; by highest level: %s" % (
            cs, len(u), u.source_id.nunique(), u.survey.value_counts().to_dict(), u.f_max.value_counts().sort_index().to_dict()))
        P("   AllWISE excess flag: %d; Gaia variable: %d; Gaia non-single-star: %d" % (
            int(u.ir_excess.sum()), int(u.variable.sum()), int((u.nss.fillna(0) > 0).sum())))
        r = u.plx_ratio_dr2.dropna()
        P("   DR2/DR3 parallax ratio: min %.3f, max %.3f (n = %d)" % (r.min(), r.max(), len(r)))
        P("   tangential velocity (km/s): median %.0f, max %.0f" % (u.vtan.median(), u.vtan.max()))
        P("   parallax/error: median %.0f; fraction below 70: %.2f; RUWE median %.2f" % (
            u.poe.median(), np.mean(u.poe < 70), u.ruwe.median()))
        if "galex_FUVexp" in u.columns:
            P("   GALEX AIS: with FUV exposure %d, FUV detections %d" % (int((u.galex_FUVexp > 0).sum()), int(u.galex_FUV.notna().sum())))
        if "irsa_w3snr" in u.columns:
            P("   IRSA AllWISE: matched %d, W3 S/N>3: %d, W4 S/N>3: %d" % (
                int(u.irsa_w1mpro.notna().sum()), int((u.irsa_w3snr > 3).sum()), int((u.irsa_w4snr > 3).sum())))
        if "simbad_otype" in u.columns:
            P("   SIMBAD: in SIMBAD %d; object types %s" % (int(u.simbad_main_id.fillna("").ne("").sum()),
                                                            u.simbad_otype.fillna("").replace("", np.nan).dropna().value_counts().to_dict()))
        if "rave_teff" in u.columns:
            P("   RAVE DR6 MADERA parameters for %d stars" % int(u.rave_teff.notna().sum()))
        gal = u[u.teff_galah.notna()]
        P("   GALAH DR3 for %d unresolved stars; still below the logg-free main-sequence floor at GALAH Teff, [Fe/H]: %d" % (
            len(gal), int((gal.m_ms_galah < 0).sum())))
        if cs == "final":
            cols = ["source_id", "survey", "f_max", "ra", "dec", "teff", "logg", "feh", "G", "dist", "D_G", "D_K", "D", "sigma_i",
                    "m_ms", "m_cmd", "z_bprp", "z_gk", "ruwe", "poe", "vtan", "plx_ratio_dr2", "n_rep", "D_rep_med", "teff_gphot",
                    "variable", "ir_excess"]
            P(u.sort_values(["f_max", "D"], ascending=[False, True])[cols].to_string(index=False, float_format=g3))
            cols2 = ["source_id", "survey", "f_max", "simbad_main_id", "simbad_otype", "simbad_sptype", "rave_teff", "rave_logg",
                     "rave_mh", "rave_qual", "D_rave", "teff_galah", "logg_galah", "feh_galah", "D_galah", "m_ms_galah",
                     "irsa_w3snr", "irsa_w4snr", "galex_FUVexp", "galex_FUV", "galex_NUV"]
            cols2 = [c for c in cols2 if c in u.columns]
            P(u.sort_values(["f_max", "D"], ascending=[False, True])[cols2].to_string(index=False, float_format=g3))
    ca = rd("candidates_all.csv")
    gal = ca[ca.teff_galah.notna()]
    P("-- GALAH DR3 (information only; log g is parallax-based, so D_galah is NOT an independent test)")
    P("   flagged stars of any tier with GALAH: %d; mean GALAH logg - survey logg = %.2f; min D_galah = %.2f" % (
        len(gal), (gal.logg_galah - gal.logg).mean(), gal.D_galah.min()))
    P("   of these, still below the logg-free main-sequence floor at GALAH Teff and [Fe/H]: %d" % int((gal.m_ms_galah < 0).sum()))
    h = {sv: rd("hist_D_%s.csv" % sv) for sv in SURVEYS}
    for sv in SURVEYS:
        q = h[sv].D_hi <= -0.3
        P("%s: stars with D < -0.3 before the colour-consistency cut: %d ; after: %d" % (
            sv, int(h[sv].n_before_colour_cut[q].sum()), int(h[sv].n_clean[q].sum())))

    P("\n" + "=" * 30, "10. robustness variants (population 95% upper limits; n_unresolved / n_flagged)")
    r = rd("robustness.csv")
    for sv in SURVEYS:
        t = r[r.survey == sv]
        piv = t.pivot_table(index="variant", columns="f", values="UL95_pop", sort=False)
        cnt = t.assign(nn=t.n_unresolved.astype(str) + "/" + t.n_flagged.astype(str)).pivot_table(
            index="variant", columns="f", values="nn", aggfunc="first", sort=False)
        ncl = t.groupby("variant", sort=False).n_clean.first()
        out = pd.concat([ncl, piv.add_prefix("UL95 f="), cnt.add_prefix("unres/flag f=")], axis=1)
        P("--", sv)
        P(out.to_string(float_format=g3))
    b = r[(r.variant == "baseline")]
    P("-- flagged stars in the fainter half of each sample (baseline, final cuts): flagged in fainter half / flagged")
    P(b.assign(x=b.n_flagged_fainter_half.astype(str) + " / " + b.n_flagged.astype(str)).pivot_table(
        index="survey", columns="f", values="x", aggfunc="first", sort=False).to_string())

    P("\n" + "=" * 30, "11. candidate rate (three-way flagged, before vetting) versus astrometric quality")
    P("-- under the ORIGINAL cuts (RUWE < 1.4, parallax/error >= 20)")
    d = rd("astrometry_diagnostic.csv")
    P(d[d.f <= 0.5].to_string(index=False, float_format=g3))
    P("-- inside the FINAL cuts (RUWE < 1.2, parallax/error >= 50)")
    d = rd("astrometry_diagnostic_final.csv")
    P(d[d.f <= 0.5].to_string(index=False, float_format=g3))

    P("\n" + "=" * 30, "12. colour-magnitude test without Teff or log g")
    for sv in SURVEYS:
        P(rd("cmd_only_%s.csv" % sv).to_string(index=False))
    P(rd("cmd_check.csv").to_string(index=False, float_format=g3))

    P("\n" + "=" * 30, "13. warmest swarm for which the Ks band is still 'grey' (blackbody estimate)")
    P("A swarm absorbing a fraction f and re-radiating it as a blackbody of temperature T adds Ks-band (2.16 um) flux")
    P("R = f/(1-f) * (T*/T)^4 * (exp(x*)-1)/(exp(x)-1), x = hc/(lambda k T), relative to the transmitted starlight.")
    P("T_max is the temperature at which this brightens Ks by 0.10 mag (comparable to the colour-consistency tolerance).")
    hc_k = 14387.77 / 2.159                                 # hc/(lambda k) in K for lambda = 2.159 micron
    rows = []
    for Ts in (5000.0, 5800.0, 6300.0):
        for f in L.F_GRID:
            T = np.linspace(150.0, 1500.0, 5401)
            R = f / (1 - f) * (Ts / T) ** 4 * np.expm1(hc_k / Ts) / np.expm1(hc_k / T)
            dK = 2.5 * np.log10(1 + R)
            rows.append(dict(T_star=Ts, f=f, T_max_K=float(T[np.searchsorted(dK, 0.10)]),
                             dK_at_300K=float(np.interp(300.0, T, dK)), dK_at_600K=float(np.interp(600.0, T, dK))))
    tm = pd.DataFrame(rows)
    tm.to_csv(os.path.join(DATA_MAIN, "warm_swarm_limit.csv"), index=False)
    P(tm.to_string(index=False, float_format=g3))

    with open(os.path.join(DATA_MAIN, "report_numbers.txt"), "w") as fh:
        fh.write(buf.getvalue())


if __name__ == "__main__":
    main()