"""What does the spectroscopy add at large f?  A colour-magnitude test without Teff or log g.

After the three-way criterion a star is a candidate only if it lies below the faint floor of the main
sequence at its Teff, at its colour and at its [Fe/H].  At f >= 0.75 this is close to asking whether
anything sits 1.5-2.5 mag below the main sequence.  Three counts of stars more than x mag below the
main-sequence ridge in BOTH G and Ks are compared:

  (a) LAMOST clean sample, ridge at fixed colour AND spectroscopic [Fe/H]   (D_cmd of the main analysis)
  (b) LAMOST clean sample, ridge at fixed colour only                        (no spectroscopic quantity)
  (c) the Gaia DR3 2% random parent sample with the same astrometric / photometric / extinction cuts and
      the same colour and apparent-magnitude range, ridge at fixed colour only (no spectroscopy at all)

(b) versus (a) shows what the metallicity contributes; (c) versus (b) shows what membership of a
spectroscopic FGK sample contributes (it excludes white-dwarf binaries, subdwarfs and other objects that
do populate the region below the main sequence in a purely photometric sample).

Writes data/cmd_check.csv and data/cmd_check_parent_outliers.csv
"""
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 selection import parent_cuts, twomass_ok

warnings.filterwarnings("ignore")
XMAG = (0.75, 1.0, 1.25, 1.5, 2.0, 2.5)


def run(sv="lamost"):
    s = pd.read_pickle(dpkl(sv, "final"))
    c = s[s.clean]
    ms = c[c.logg >= L.LOGG_MS]
    lo, hi = np.percentile(c.bprp0, [1, 99])
    glo, ghi = np.percentile(c.G, [1, 99])
    cb = np.arange(lo, hi + 0.02, 0.02)
    idx = np.digitize(ms.bprp0.values, cb)
    cen, rg, rk = [], [], []
    for i in range(1, len(cb)):
        q = idx == i
        if q.sum() >= 50:
            cen.append(0.5 * (cb[i - 1] + cb[i]))
            rg.append(np.median(ms.MG.values[q]))
            rk.append(np.median(ms.MK.values[q]))
    cen, rg, rk = np.array(cen), np.array(rg), np.array(rk)

    def below(t):
        """mag below the colour-only ridge, in the less deviant of G and Ks"""
        return np.minimum(t.MG.values - np.interp(t.bprp0.values, cen, rg), t.MK.values - np.interp(t.bprp0.values, cen, rk))

    par = pd.read_pickle(os.path.join(WORK, "parent_stars.pkl"))
    p = par[parent_cuts(par) & (par.poe >= L.POE_MIN) & twomass_ok(par) & par.bprp0.between(lo, hi) & par.G.between(glo, ghi)]
    cc = c[c.bprp0.between(lo, hi) & c.G.between(glo, ghi)]
    log("cmd_check: %s clean %d (in colour/magnitude range %d); parent %d; colour %.2f-%.2f, G %.1f-%.1f"
        % (sv, len(c), len(cc), len(p), lo, hi, glo, ghi))
    def vtan(t):
        return 4.74 * np.hypot(t.pmra.values, t.pmde.values) / t.plxc.values

    sets = (("(a) %s clean, ridge at colour and [Fe/H]" % sv, cc, -cc.D_cmd.values / 0.4),
            ("(b) %s clean, ridge at colour only" % sv, cc, below(cc)),
            ("(c) Gaia parent sample, no spectroscopy, ridge at colour only", p, below(p)))
    rows = []
    for x in XMAG:
        for name, t, b in sets:
            q = b > x
            v = vtan(t)[q]
            rows.append(dict(sample=name, mag_below_ridge=x, n_stars=len(t), n_below=int(q.sum()), fraction=q.sum() / len(t),
                             fraction_UL95=L.poisson_ul(int(q.sum())) / len(t),
                             median_vtan_kms=float(np.median(v)) if q.any() else np.nan,
                             frac_vtan_gt_150=float(np.mean(v > 150)) if q.any() else np.nan,
                             median_vtan_all_kms=float(np.median(vtan(t)))))
    # the parent-sample stars more than 1.25 mag below the ridge (a pointer for a Gaia-only search; not analysed here)
    b = below(p)
    po = p.assign(mag_below_ridge=b, vtan=vtan(p))[b > 1.25]
    po[["source_id", "ra", "dec", "G", "bprp0", "MG", "MK", "dist", "poe", "ruwe", "mag_below_ridge", "vtan"]].to_csv(
        os.path.join(DATA_MAIN, "cmd_check_parent_outliers.csv"), index=False)
    out = pd.DataFrame(rows)
    out.to_csv(os.path.join(DATA_MAIN, "cmd_check.csv"), index=False)
    print(out.to_string(index=False, float_format=lambda v: "%.3g" % v))
    return out


if __name__ == "__main__":
    run()