"""Selection efficiency for grey-dimmed stars relative to undimmed stars, per survey.

A star dimmed by a fraction f is fainter by dm = -2.5 log10(1-f) in every band but keeps its
colours, spectrum and parallax.  Relative to an identical undimmed star it is (a) less likely to
have been targeted / to have reached the spectroscopic S/N cut (survey selection is on apparent
magnitude), and (b) less likely to pass our parallax-S/N and 2MASS cuts.

Method: a 2 per cent random sample of Gaia DR3 (random_index) plays the role of the parent
population.  For each survey the empirical selection fraction S(m | colour class) is the ratio of
the clean survey sample to the parent sample passing the same non-spectroscopic cuts, in 0.25 mag
bins of the survey's selection magnitude (G for LAMOST and Gaia RVS, 2MASS H for APOGEE).  Then

    eps(f) = sum_j S(m_j + dm) C_j(dm)  /  sum_j S(m_j) C_j(0)

with the sum over parent stars j in the survey footprint and C_j the cuts that depend on apparent
magnitude (parallax S/N re-scaled with the empirical sigma_parallax(G) relation; 2MASS Ks limit).
S is set to zero at magnitudes where a colour class has no clean survey stars (beyond the range
sampled by undimmed stars inside our distance cut), which can only under-estimate eps.
"closure" = sum_j S(m_j) C_j(0) / N_clean should be 1 (it is 0.98-1.00).

This eps is a MODEL: it assumes that survey inclusion depends only on apparent magnitude and on three
colour classes, averaged over the footprint.

Independent cross-check (twin_eps; proposed by the referee of the first version): a star dimmed by f
at distance d has the same apparent magnitude, colours, spectrum and sky position as an undimmed twin
at distance k d with k = (1-f)^-1/2, hence the same probability of being observed and of passing every
cut that depends on apparent quantities.  The number of dimmed stars expected in the sample is therefore
the sum over the REAL survey stars at k d, each weighted by the ratio of the numbers of stars available
at the two distances, k^-3 n(d)/n(k d), with n a smooth disc density model, and required to pass the
parallax cuts at the nearer distance.  This uses no parent sample and no selection function; it needs
survey stars beyond 1 kpc, which exist in the downloads for APOGEE and LAMOST (not for GSP-Spec).

Writes data/selection_eff.csv, data/selection_function_<survey>.csv, data/selection_eff_twin.csv
"""
import os
import sys
import warnings

import numpy as np
import pandas as pd

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

warnings.filterwarnings("ignore")
MBINS = np.arange(4.0, 19.0001, 0.25)
SELMAG = dict(apogee="H", lamost="G", gspspec="G")


def parent_cuts(p):
    cs = L.cstar(p.BP - p.RP, p.bprp_excess)
    return (p.solved.isin([31, 95]) & p.G.notna() & p.BP.notna() & p.RP.notna() & (p.dup == 0) & (p.plx > 1.0)
            & (p.sfd < L.SFD_MAX) & (p.ruwe < L.RUWE_MAX) & (p.ipd_fmp <= 2) & (p.ipd_gofha < 0.1)
            & (np.abs(cs) < 5 * L.cstar_sigma(p.G)))


def twomass_ok(p):
    return (p.K.notna() & (p.e_K < 0.1) & (p.d2m < 1.0) & (p.n2m == 1)
            & p.qfl.astype(str).str[2:3].isin(["A", "B"]))


def median_relation(x, y, bins):
    idx = np.digitize(x, bins)
    med = np.array([np.nanmedian(y[idx == i]) if np.sum(idx == i) > 20 else np.nan for i in range(1, len(bins))])
    cen = 0.5 * (bins[1:] + bins[:-1])
    ok = np.isfinite(med)
    return lambda q: np.interp(q, cen[ok], med[ok])


def run(surveys):
    par = pd.read_pickle(os.path.join(WORK, "parent_stars.pkl"))
    base = parent_cuts(par)
    # empirical error-magnitude relations (parent sample)
    e_plx_of_G = median_relation(par.G[base].values, par.e_plx[base].values, np.arange(5, 17.6, 0.5))
    # 2MASS: probability that a star of magnitude Ks passes the Ks quality cut, from Gaia-selected parent stars.
    # Ks of stars without a valid 2MASS entry is predicted from G and BP-RP via the matched stars.
    has = base & twomass_ok(par)
    col_to_gk = median_relation((par.BP - par.RP)[has].values, (par.G - par.K)[has].values, np.arange(0.3, 2.6, 0.1))
    par["K_pred"] = np.where(par.K.notna(), par.K, par.G - col_to_gk((par.BP - par.RP).values))
    kb = np.arange(6, 18.01, 0.25)
    kidx = np.digitize(par.K_pred[base], kb)
    p2m = np.array([np.mean(twomass_ok(par[base])[kidx == i]) if np.sum(kidx == i) > 20 else np.nan
                    for i in range(1, len(kb))])
    kc = 0.5 * (kb[1:] + kb[:-1])
    okk = np.isfinite(p2m)
    p2m_of_K = lambda q: np.interp(q, kc[okk], p2m[okk], left=p2m[okk][0], right=0.0)
    pd.DataFrame(dict(K=kc, p_2mass_ok=p2m)).to_csv(os.path.join(DATA, "twomass_completeness.csv"), index=False)

    rows = []
    for sv in surveys:
        s = pd.read_pickle(dpkl(sv))
        c = s[s.clean]
        del s
        mcol = SELMAG[sv]
        # footprint: HEALPix level-5 pixels with >= 3 clean survey stars
        pix, cnt = np.unique(c.hpx5.values, return_counts=True)
        foot = set(pix[cnt >= 3])
        infoot = par.hpx5.isin(foot).values
        # CMD region occupied by the survey sample
        xb, yb = np.arange(0.3, 2.0001, 0.05), np.arange(-1.0, 10.0001, 0.25)
        Hs = np.histogram2d(c.bprp0, c.MG, [xb, yb])[0]
        order = np.sort(Hs.ravel())[::-1]
        thr = order[np.searchsorted(np.cumsum(order), 0.995 * Hs.sum())]
        mask = Hs >= max(thr, 1)
        ix = np.clip(np.digitize(par.bprp0, xb) - 1, 0, len(xb) - 2)
        iy = np.clip(np.digitize(par.MG, yb) - 1, 0, len(yb) - 2)
        incmd = mask[ix, iy] & par.bprp0.between(xb[0], xb[-1]).values & par.MG.between(yb[0], yb[-1]).values
        # colour-class boundaries from the survey's own colour-Teff relation
        def col_at(T):
            q = c.teff.between(T - 50, T + 50)
            return np.median(c.bprp0[q])
        b_KG, b_GF = col_at(5200.0), col_at(6000.0)
        pcls = np.where(par.bprp0 >= b_KG, "K", np.where(par.bprp0 >= b_GF, "G", "F"))
        P = par[base & infoot & incmd].copy()
        P["cls"] = pcls[base & infoot & incmd]
        P["m"] = P[mcol] if mcol != "G" else P["G"]
        if mcol == "H":
            P = P[P.H.notna()]
        log("%s: parent stars in footprint/CMD/cuts: %d  (footprint %d level-5 pixels; colour bounds %.3f %.3f)"
            % (sv, len(P), len(foot), b_KG, b_GF))
        sel_rows = []

        def eps_block(cs_all, P_all, label, binw=0.25, smooth=True, save_sel=False, nboot=200):
            """eps for each colour class and f, from clean survey stars cs_all and parent stars P_all."""
            mb = np.arange(4.0, 19.0001, binw)
            cen = 0.5 * (mb[1:] + mb[:-1])
            for cname, lo, hi in CLASSES:
                cs_ = cs_all[cs_all.cls == cname]
                Pc = P_all[P_all.cls == cname]
                if len(cs_) < 200 or len(Pc) < 100:
                    continue
                m0 = Pc.m.values
                pass0 = (Pc.poe.values >= L.POE_MIN) & twomass_ok(Pc).values
                ns = np.histogram(cs_[mcol].values, mb)[0].astype(float)
                npar = np.histogram(m0[pass0], mb)[0].astype(float)
                if smooth:                                    # 3-bin boxcar on both histograms before the ratio
                    ker = np.ones(3) / 3.0
                    ns_s, np_s = np.convolve(ns, ker, "same"), np.convolve(npar, ker, "same")
                else:
                    ns_s, np_s = ns, npar
                S = np.where(np_s > 0, ns_s / np.maximum(np_s, 1e-9), 0.0)
                S[ns_s == 0] = 0.0
                if save_sel:
                    for a, b_, n1, n2, sv_ in zip(mb[:-1], mb[1:], ns, npar, S):
                        sel_rows.append(dict(cls=cname, m_lo=a, m_hi=b_, n_survey=n1, n_parent=n2, S_smoothed=sv_))
                Sf = lambda mm: np.interp(mm, cen, S, left=0.0, right=0.0)
                wA = Sf(m0) * pass0
                w0 = wA.sum()
                for f in L.F_GRID:
                    dm = -2.5 * np.log10(1.0 - f)
                    poe_f = Pc.poe.values * e_plx_of_G(Pc.G.values) / e_plx_of_G(Pc.G.values + dm)
                    Kp = Pc.K_pred.values       # probability to keep a good 2MASS Ks at the dimmed magnitude
                    p2 = np.where(twomass_ok(Pc).values,
                                  np.minimum(1.0, p2m_of_K(Kp + dm) / np.maximum(p2m_of_K(Kp), 1e-6)), 0.0)
                    wf_sel = np.sum(Sf(m0 + dm) * pass0)                      # survey selection only
                    wB = Sf(m0 + dm) * (poe_f >= L.POE_MIN) * p2               # + parallax S/N and 2MASS
                    rng = np.random.default_rng(12345)
                    bs = []
                    for _ in range(nboot):
                        ii = rng.integers(0, len(m0), len(m0))
                        bs.append(wB[ii].sum() / max(wA[ii].sum(), 1e-9))
                    rows.append(dict(survey=sv, variant=label, cls=cname, f=f, dm=dm, n_parent=len(Pc),
                                     n_survey=len(cs_), eps=wB.sum() / w0,
                                     eps_lo=np.percentile(bs, 16) if nboot else np.nan,
                                     eps_hi=np.percentile(bs, 84) if nboot else np.nan,
                                     eps_survey_selection_only=wf_sel / w0, eps_uniform_density=10 ** (-0.6 * dm),
                                     closure=w0 / max(len(cs_), 1)))

        eps_block(c, P, "baseline", save_sel=True)
        # systematics: sky regions (the selection function differs from plate to plate), binning, smoothing
        babs_c, babs_P = np.abs(c.glat.values), np.abs(P.glat.values)
        for lo, hi in ((0, 35), (35, 55), (55, 90)):
            eps_block(c[(babs_c >= lo) & (babs_c < hi)], P[(babs_P >= lo) & (babs_P < hi)],
                      "|b| %d-%d deg" % (lo, hi), nboot=0)
        eps_block(c, P, "0.5 mag bins", binw=0.5, nboot=0)
        eps_block(c, P, "no smoothing", smooth=False, nboot=0)
        pd.DataFrame(sel_rows).to_csv(os.path.join(DATA, "selection_function_%s.csv" % sv), index=False)
    allv = pd.DataFrame(rows)
    allv.to_csv(os.path.join(DATA, "selection_eff_variants.csv"), index=False)
    out = allv[allv.variant == "baseline"].drop(columns=["variant"]).reset_index(drop=True)
    # systematic range of eps over the variants (min and max), attached to the baseline table
    g = allv.groupby(["survey", "cls", "f"]).eps.agg(["min", "max"]).reset_index().rename(
        columns={"min": "eps_sys_min", "max": "eps_sys_max"})
    out = out.merge(g, on=["survey", "cls", "f"], how="left")
    out.to_csv(os.path.join(DATA, "selection_eff.csv"), index=False)
    pd.set_option("display.width", 250)
    print(out.round(4).to_string(index=False))
    return out


# ------------------------------------------------------------------------------------------------
# Independent cross-check of eps: undimmed twins at larger distance
# ------------------------------------------------------------------------------------------------
R0 = 8200.0                                   # pc, Galactocentric distance of the Sun
# (label, thin-disc scale height, thick-disc scale height, thick/thin local normalisation, radial scale length)
DISC_MODELS = (("h 300/900 pc, 12% thick, hR 2.6 kpc", 300.0, 900.0, 0.12, 2600.0),
               ("thin 250 pc", 250.0, 900.0, 0.12, 2600.0),
               ("thin 350 pc", 350.0, 900.0, 0.12, 2600.0),
               ("thick 4%", 300.0, 900.0, 0.04, 2600.0),
               ("h 200/700 pc, 20% thick", 200.0, 700.0, 0.20, 2600.0),
               ("no radial gradient", 300.0, 900.0, 0.12, np.inf),
               ("hR 3.5 kpc", 300.0, 900.0, 0.12, 3500.0))
TWIN_POE_MIN = 5.0                            # the twin itself needs a usable parallax to be placed


def disc_density(dist, glon, glat, h1, h2, fthick, hR):
    b, l = np.radians(glat), np.radians(glon)
    z = np.abs(dist * np.sin(b))
    x = dist * np.cos(b)
    R = np.sqrt(R0 ** 2 + x ** 2 - 2.0 * R0 * x * np.cos(l))
    rad = np.exp(-(R - R0) / hR) if np.isfinite(hR) else 1.0
    return rad * (np.exp(-z / h1) + fthick * np.exp(-z / h2))


def twin_eps(surveys=("apogee", "lamost")):
    rows = []
    for sv in surveys:
        d = pd.read_pickle(os.path.join(WORK, "%s_stars.pkl" % sv))
        m = np.asarray(L.survey_cuts(d, sv)[0][1], bool).copy()
        for name, cut in L.common_cuts(d):
            if not name.startswith("parallax"):          # every cut except the parallax one
                m &= np.asarray(cut, bool)
        e = d.loc[m, ["plx", "poe", "glon", "glat", "teff"]].copy()
        del d
        plxc = e.plx.values + 0.017
        ok0 = (plxc > 0.2) & (e.poe.values > TWIN_POE_MIN)
        e, plxc = e[ok0], plxc[ok0]
        dist = 1000.0 / plxc
        plx, poe = e.plx.values, e.poe.values
        cls = np.where(e.teff.values < 5200, "K", np.where(e.teff.values < 6000, "G", "F"))
        base = (plx > 1.0) & (poe >= L.POE_MIN)          # the undimmed sample (before the last two cuts)
        log("twin cross-check %s: %d stars pass all non-parallax cuts, %d also the parallax cuts" % (sv, len(e), base.sum()))
        for label, h1, h2, ft, hR in DISC_MODELS:
            n_far = disc_density(dist, e.glon.values, e.glat.values, h1, h2, ft, hR)
            for f in L.F_GRID:
                k = (1.0 - f) ** -0.5
                ok = (plx * k > 1.0) & (poe * k >= L.POE_MIN)   # the dimmed star, at dist / k, passes the parallax cuts
                w = k ** -3 * disc_density(dist / k, e.glon.values, e.glat.values, h1, h2, ft, hR) / n_far * ok
                for cname in ("all", "K", "G", "F"):
                    q = np.ones(len(e), bool) if cname == "all" else (cls == cname)
                    rows.append(dict(survey=sv, model=label, cls=cname, f=f, eps_twin=w[q].sum() / base[q].sum(),
                                     n_twins=int((ok & q).sum()),
                                     frac_weight_beyond_1kpc=w[q & (plx <= 1.0)].sum() / max(w[q].sum(), 1e-12)))
    out = pd.DataFrame(rows)
    # attach the model eps for comparison
    mod = pd.read_csv(os.path.join(DATA, "selection_eff.csv"))[["survey", "cls", "f", "eps"]].rename(columns={"eps": "eps_model"})
    out = out.merge(mod, on=["survey", "cls", "f"], how="left")
    out.to_csv(os.path.join(DATA, "selection_eff_twin.csv"), index=False)
    piv = out[out.cls == "all"].pivot_table(index=["survey", "model"], columns="f", values="eps_twin", sort=False)
    print(piv.round(3).to_string())
    print(out[(out.model == DISC_MODELS[0][0]) & (out.cls != "all")].round(3).to_string(index=False))
    return out


if __name__ == "__main__":
    args = [a for a in sys.argv[1:] if a != "--twin-only"]
    if "--twin-only" not in sys.argv:
        run(args or ["apogee", "lamost", "gspspec"])
    twin_eps([a for a in (args or ["apogee", "lamost"]) if a in ("apogee", "lamost")])