"""Per-survey analysis: cuts, spectroscopic-luminosity calibration, luminosity deficits D,
tail statistics, injection-recovery and candidate lists.

    python3 analysis.py apogee lamost gspspec

The astrometric cut set is taken from the environment (DIM_STARS_CUTSET = final | orig, see dimlib);
products of the "orig" set go to data/origcuts/ and WORK/<survey>_D_origcuts.pkl.

Writes  WORK/<survey>_D.pkl               every star passing the cuts, with D (large; not in shared folder)
        data/cuts_<survey>.csv            cut-by-cut star counts
        data/cmd_only_<survey>.csv        stars far below the main sequence at fixed colour and [Fe/H]
        data/core_tail_<survey>.csv       core scatter and tail fractions by Teff/logg class
        data/hist_D_<survey>.csv          histogram of D (for figures)
        data/tail_<survey>.csv            every star with D < -0.15 (before the colour cut)
        data/limits_raw_<survey>.csv      N_sens, candidate counts and recovery per dimming fraction f
"""
import os
import sys
import warnings

import numpy as np
import pandas as pd

import dimlib as L
from common import DATA, WORK, dpkl, log

warnings.filterwarnings("ignore")

NSIG_SENS, NSIG_THR = L.NSIG_SENS, L.NSIG_THR      # see dimlib.flag_mask
# columns not used by any later stage; dropped from the stored table to save disk space
DROP = ["ral", "decl", "snrr", "snri", "subClass", "e_feh", "e_rv_l", "teff_med", "logg_med", "feh_med", "rv_spread",
        "ipd_fow", "e_G", "nbp_blend", "nrp_blend", "e_rv_gaia", "n_rv", "grvs", "rvs_sn", "logg_gphot", "ag_gphot",
        "ra2000", "dec2000", "tmass", "J", "e_J", "e_H", "dustfrac", "r_MG", "s_MG", "r_MK", "s_MK", "r_bprp0",
        "s_bprp0", "r_gk0", "s_gk0", "sig_D", "Dmed_nb", "Tel", "Field", "e_mh", "feh_fe", "am", "Vsini", "TeffSp",
        "loggSp", "sflag", "aflag", "nentry", "Hmin", "Hmax", "jkmin", "jkmax", "AkT", "Survey", "EFlag", "Chi2",
        "K_xm", "phot_g_mean_flux_over_error", "teff_gspspec_lower", "teff_gspspec_upper", "logg_gspspec_lower",
        "logg_gspspec_upper", "mh_gspspec_lower", "mh_gspspec_upper", "alphafe_gspspec", "logchisq_gspspec",
        "flags_gspspec"] + L.FLOOR_COLS
CLASSES = (("K", 4000, 5200), ("G", 5200, 6000), ("F", 6000, 6500.01))


def teff_class(teff):
    out = np.full(len(teff), "", dtype=object)
    for name, lo, hi in CLASSES:
        out[(teff >= lo) & (teff < hi)] = name
    return out


def local_faint_sigma(D, idx):
    """Robust faint-side scatter of D among each star's K nearest neighbours in parameter space."""
    nb = D[idx]
    med = np.nanmedian(nb, axis=1)
    dev = med[:, None] - nb                        # positive on the faint side
    dev = np.where(dev > 0, dev, np.nan)
    return 1.4826 * np.nanmedian(dev, axis=1), med


def run_survey(survey):
    log("=== analysis:", survey)
    d = pd.read_pickle(os.path.join(WORK, "%s_stars.pkl" % survey))
    m, rows = L.apply_cuts(d, survey)
    s, cal = L.compute_deficits(d[m])
    s["sigma_i"], s["Dmed_nb"] = local_faint_sigma(s.D.values, cal.idx)
    s["sigma_i"] = np.maximum(s.sigma_i, 0.015)                     # floor: photometric + parallax errors
    s["cls"] = teff_class(s.teff.values)
    s["ms"] = s.logg >= L.LOGG_MS
    pop = s.rk < L.KNN_RMAX
    rows.append(("populated parameter space: distance to %d-th neighbour < %g (scaled units)" % (L.KNN_K, L.KNN_RMAX),
                 int(pop.sum())))
    col = (np.abs(s.z_bprp) < L.COLOUR_ZMAX) & (np.abs(s.z_gk) < L.COLOUR_ZMAX)
    s["pop"] = pop
    s["colour_ok"] = col
    s["clean"] = pop & col
    rows.append(("colours consistent with spectroscopic Teff: |z(BP-RP)|<%g and |z(G-Ks)|<%g" % (L.COLOUR_ZMAX, L.COLOUR_ZMAX),
                 int(s.clean.sum())))
    print("   %9d  %s\n   %9d  %s" % (rows[-2][1], rows[-2][0], rows[-1][1], rows[-1][0]))
    pd.DataFrame(rows, columns=["cut", "n_remaining"]).to_csv(os.path.join(DATA, "cuts_%s.csv" % survey), index=False)
    s["survey"] = survey
    s.drop(columns=[x for x in DROP if x in s.columns]).to_pickle(dpkl(survey))

    c = s[s.clean]
    # ---------------- core scatter and tails
    out = []
    def stats(lbl, q):
        x = c.D[q].values
        if len(x) < 30:
            return
        out.append(dict(survey=survey, subset=lbl, n=len(x), median_D=np.median(x), sigma_core=L.robust_sigma(x),
                        sigma_faint=1.4826 * np.median(np.median(x) - x[x < np.median(x)]),
                        median_sigma_i=np.median(c.sigma_i[q]),
                        p0p1=np.percentile(x, 0.1), p1=np.percentile(x, 1), p99=np.percentile(x, 99),
                        frac_lt_m0p125=np.mean(x < -0.125), frac_lt_m0p2=np.mean(x < -0.2),
                        frac_lt_m0p3=np.mean(x < -0.301), frac_lt_m0p6=np.mean(x < -0.602),
                        frac_lt_m1=np.mean(x < -1.0), n_lt_m0p3=int(np.sum(x < -0.301)),
                        n_lt_m0p6=int(np.sum(x < -0.602)), n_lt_m1=int(np.sum(x < -1.0)),
                        frac_gt_p0p2=np.mean(x > 0.2)))
    stats("all", np.ones(len(c), bool))
    for name, lo, hi in CLASSES:
        stats("%s dwarfs (logg>=%.1f)" % (name, L.LOGG_MS), (c.cls == name).values & c.ms.values)
        stats("%s subgiants (logg<%.1f)" % (name, L.LOGG_MS), (c.cls == name).values & ~c.ms.values)
    pd.DataFrame(out).to_csv(os.path.join(DATA, "core_tail_%s.csv" % survey), index=False)
    print(pd.DataFrame(out)[["subset", "n", "sigma_core", "median_sigma_i", "frac_lt_m0p2", "n_lt_m0p3", "n_lt_m0p6",
                             "n_lt_m1"]].round(4).to_string(index=False))

    # ---------------- colour-magnitude test without Teff or logg: stars far below the main sequence
    # D_cmd = log10(L / L_ridge) with L_ridge the median of the 200 nearest main-sequence stars in
    # (BP-RP, [Fe/H]); the less deficient of G and Ks, so "x mag below" means below in both bands.
    cm = []
    for x in (0.31, 0.75, 1.0, 1.25, 1.5, 2.0, 2.5):
        cm.append(dict(survey=survey, mag_below_ridge=x, n_clean=int(len(c)),
                       n_below_clean=int((c.D_cmd < -0.4 * x).sum()),
                       n_before_colour_cut=int(s["pop"].sum()),
                       n_below_before_colour_cut=int((s.D_cmd[s["pop"]] < -0.4 * x).sum())))
    pd.DataFrame(cm).to_csv(os.path.join(DATA, "cmd_only_%s.csv" % survey), index=False)
    print(pd.DataFrame(cm).to_string(index=False))

    # ---------------- histograms
    bins = np.arange(-2.0, 1.0001, 0.01)
    h_all = np.histogram(s.D[s["pop"]], bins)[0]
    h_clean = np.histogram(c.D, bins)[0]
    h_z = np.histogram(np.clip(c.D / c.sigma_i, -40, 40), np.arange(-40, 40.01, 0.25))[0]
    pd.DataFrame(dict(D_lo=bins[:-1], D_hi=bins[1:], n_before_colour_cut=h_all, n_clean=h_clean)).to_csv(
        os.path.join(DATA, "hist_D_%s.csv" % survey), index=False)
    pd.DataFrame(dict(z_lo=np.arange(-40, 40, 0.25), n_clean=h_z)).to_csv(
        os.path.join(DATA, "hist_z_%s.csv" % survey), index=False)

    # ---------------- tail table (all stars with D < -0.15, before colour cut)
    keep = ["source_id", "survey", "ra", "dec", "glon", "glat", "teff", "logg", "feh", "snr", "plx", "e_plx", "poe",
            "dist", "ruwe", "sepsi", "ipd_fmp", "G", "bprp0", "K", "sfd", "MG", "MK", "D_G", "D_K", "D", "sigma_i",
            "D_ms", "m_ms", "floor_ms", "D_cmd", "m_cmd", "floor_cmd",
            "z_bprp", "z_gk", "rk", "pop", "colour_ok", "clean", "cls", "ms", "teff_gphot", "rvamp", "varflag"]
    extra = [x for x in ("apogee_id", "obsid", "nspec", "teff_spread", "logg_spread", "nvis", "vscatter", "Prog",
                         "fib_sep", "non_single_star", "main_survey") if x in s.columns]
    s.loc[s.D < -0.15, keep + extra].sort_values("D").to_csv(os.path.join(DATA, "tail_%s.csv" % survey), index=False)

    # ---------------- per-f sensitive sample, flagged stars, statistical recovery
    lim = []
    for f in L.F_GRID:
        Df = np.log10(1.0 - f)
        sens, fl = L.flag_mask(c, f)
        _, fl_spec = L.flag_mask(c, f, require=("spec",))
        _, rec = L.flag_mask(c, f, inject=True)
        _, rec_spec = L.flag_mask(c, f, inject=True, require=("spec",))
        for name in ("all",) + tuple(x[0] for x in CLASSES) + ("MS", "SG"):
            q = np.ones(len(c), bool)
            if name in ("K", "G", "F"):
                q = (c.cls == name).values
            elif name == "MS":
                q = c.ms.values
            elif name == "SG":
                q = ~c.ms.values
            ns = int((sens & q).sum())
            lim.append(dict(survey=survey, f=f, D_f=Df, cls=name, n_clean=int(q.sum()), n_sens=ns,
                            n_flagged=int((fl & q).sum()), recovery=(rec & q).sum() / max(ns, 1),
                            n_flagged_spec_only=int((fl_spec & q).sum()),
                            recovery_spec_only=(rec_spec & q).sum() / max(ns, 1)))
    lim = pd.DataFrame(lim)
    lim.to_csv(os.path.join(DATA, "limits_raw_%s.csv" % survey), index=False)
    print(lim.round(4).to_string(index=False))
    return s


if __name__ == "__main__":
    for sv in (sys.argv[1:] or ["apogee", "lamost", "gspspec"]):
        run_survey(sv)