"""Vetting of faint-tail outliers.

For every star flagged at any dimming fraction f in any survey:
  1. parameters from the other spectroscopic surveys (APOGEE / LAMOST / Gaia GSP-Spec) -> deficit D
     recomputed with the other survey's own parameters and calibration ("reproduced" or not);
  2. LAMOST repeat spectra (each individual spectrum's parameters -> D);
  3. Gaia DR3 neighbours within 8 arcsec (fibre contamination by a brighter neighbour);
  4. astrometric reliability (RUWE, excess-noise significance, non_single_star, DR2-vs-DR3 parallax);
  5. variability flag; 6. AllWISE mid-infrared excess (secondary discriminator only).
A candidate is RESOLVED only on positive evidence for a mundane cause (rules in classify()).
Everything else stays an unresolved candidate and is counted in the limits.

Information only (never used to resolve a candidate):
  * GALAH DR3 parameters.  GALAH DR3 surface gravities are computed from the parallax-based luminosity,
    not from the spectrum, so a genuinely dimmed star would also get a high GALAH log g and "the deficit
    vanishes with GALAH parameters" proves nothing.  What is reported instead is whether the star is still
    below the main-sequence floor at GALAH's Teff and [Fe/H] with no log g used (m_ms_galah < 0).
  * SIMBAD main identifier, object type and spectral type (SIMBAD TAP) and RAVE DR6 MADERA parameters
    (VizieR III/283; purely spectroscopic) for the unresolved three-way candidates.

Writes data/candidates_all.csv, data/vetting_summary.csv, data/candidates_unresolved.csv,
       data/wise_control.csv
"""
import io
import os
import sys
import urllib.parse
import urllib.request
import warnings

import numpy as np
import pandas as pd
from scipy.spatial import cKDTree

import dimlib as L
import download_data as dd
from common import DATA, WORK, RAW, GAIA_TAP, IRSA_TAP, VIZIER_TAP, cache as cpath, dpkl, log, xmatch_upload, tap_sync

warnings.filterwarnings("ignore")
SURVEYS = ("apogee", "lamost", "gspspec")
SC = L.KNN_SCALE
SIMBAD_TAP = "https://simbad.cds.unistra.fr/simbad/sim-tap/sync"


def flag_levels(c):
    """Per dimming fraction: (three-way flag, spectroscopic-only flag); see dimlib.flag_mask."""
    out = {}
    for f in L.F_GRID:
        out[f] = (L.flag_mask(c, f)[1], L.flag_mask(c, f, require=("spec",))[1])
    return out


class KnnPredictor:
    """Median M_G, M_K of the K nearest clean stars of a survey in (Teff, logg, [Fe/H])."""

    def __init__(self, clean):
        self.X = np.vstack([clean.teff.values / SC["teff"], clean.logg.values / SC["logg"],
                            clean.feh.values / SC["feh"]]).T
        self.tree = cKDTree(self.X)
        self.MG, self.MK, self.D = clean.MG.values, clean.MK.values, clean.D.values

    def deficit(self, teff, logg, feh, MG, MK, k=L.KNN_K):
        X = np.vstack([np.asarray(teff, float) / SC["teff"], np.asarray(logg, float) / SC["logg"],
                       np.asarray(feh, float) / SC["feh"]]).T
        ok = np.isfinite(X).all(axis=1)
        D = np.full(len(X), np.nan)
        sig = np.full(len(X), np.nan)
        rk = np.full(len(X), np.nan)
        if ok.any():
            dist, idx = self.tree.query(X[ok], k=k)
            pg, pk = np.median(self.MG[idx], axis=1), np.median(self.MK[idx], axis=1)
            D[ok] = np.maximum(-0.4 * (np.asarray(MG, float)[ok] - pg), -0.4 * (np.asarray(MK, float)[ok] - pk))
            nb = self.D[idx]
            med = np.median(nb, axis=1)
            dev = np.where(med[:, None] - nb > 0, med[:, None] - nb, np.nan)
            sig[ok] = np.maximum(1.4826 * np.nanmedian(dev, axis=1), 0.015)
            rk[ok] = dist[:, -1]
        return D, sig, rk


class FloorPredictor:
    """Main-sequence floor at (Teff, [Fe/H]) from the main-sequence stars (logg >= LOGG_MS) of a clean
    sample, exactly as dimlib.knn_floor but for external parameters; log g of the star is not used."""

    def __init__(self, clean):
        ms = clean[clean.logg >= L.LOGG_MS]
        self.tree = cKDTree(np.vstack([ms.teff.values / SC["teff"], ms.feh.values / SC["feh"]]).T)
        self.MG, self.MK = ms.MG.values, ms.MK.values

    def margin(self, teff, feh, MG, MK, k=L.KNN_K_FLOOR):
        X = np.vstack([np.asarray(teff, float) / SC["teff"], np.asarray(feh, float) / SC["feh"]]).T
        ok = np.isfinite(X).all(axis=1)
        m = np.full(len(X), np.nan)
        if ok.any():
            dist, idx = self.tree.query(X[ok], k=k)
            out = []
            for y, ref in ((np.asarray(MG, float)[ok], self.MG), (np.asarray(MK, float)[ok], self.MK)):
                q50, q84, q95 = np.percentile(ref[idx], [50, 84, 95], axis=1)
                edge = np.maximum((q95 - q50) + L.EDGE_MULT * (q95 - q84), L.EDGE_MIN_MAG)
                out.append(-0.4 * (y - q50 - edge))
            mm = np.maximum(out[0], out[1])
            mm[dist[:, -1] > 3.0 * L.KNN_RMAX] = np.nan          # outside the populated (Teff, [Fe/H]) plane
            m[ok] = mm
        return m


def gaia_sync(adql):
    out = tap_sync(GAIA_TAP + "/sync", adql, timeout=600)
    return pd.read_csv(io.BytesIO(out))


def gaia_extra(ids):
    """non_single_star, variability flag and DR2 parallax for a list of DR3 source ids (cached)."""
    cache = os.path.join(RAW, "cand_gaia_extra.csv")
    have = pd.read_csv(cache) if os.path.exists(cache) else pd.DataFrame(columns=["source_id"])
    need = [i for i in ids if i not in set(have.source_id.values)]
    res = [have]
    for i in range(0, len(need), 1500):
        lst = ",".join(str(x) for x in need[i:i + 1500])
        log("  Gaia archive: extra columns for %d candidates" % len(need[i:i + 1500]))
        a = gaia_sync("SELECT source_id, non_single_star, phot_variable_flag, has_epoch_photometry "
                      "FROM gaiadr3.gaia_source WHERE source_id IN (%s)" % lst)
        b = gaia_sync("SELECT n.dr3_source_id AS source_id, n.angular_distance AS dr2_sep, n.magnitude_difference AS dr2_dmag, "
                      "g2.parallax AS plx_dr2, g2.parallax_error AS e_plx_dr2 "
                      "FROM gaiadr3.dr2_neighbourhood AS n JOIN gaiadr2.gaia_source AS g2 "
                      "ON g2.source_id = n.dr2_source_id WHERE n.dr3_source_id IN (%s)" % lst)
        b = b.sort_values("dr2_sep").drop_duplicates("source_id")
        res.append(a.merge(b, on="source_id", how="left"))
    out = pd.concat(res, ignore_index=True).drop_duplicates("source_id")
    out.to_csv(cache, index=False)
    return out


def neighbours(cand):
    cache = cpath("cand_neighbours.csv")
    if os.path.exists(cache):
        nb = pd.read_csv(cache)
        if set(cand.source_id) <= set(nb.cand_id):
            return nb
    up = cand[["source_id", "ra", "dec"]].rename(columns={"source_id": "cand_id", "ra": "ra_c", "dec": "dec_c"})
    log("  XMatch: Gaia neighbours within 8 arcsec for %d candidates" % len(up))
    nb = xmatch_upload(up, "ra_c", "dec_c", "vizier:I/355/gaiadr3", 8.0, cols2="Source,Gmag,RPmag,Plx,pmRA,pmDE")
    nb.to_csv(cache, index=False)
    return nb


def allwise(tab, tag):
    cache = cpath("allwise_%s.csv" % tag)
    if os.path.exists(cache) and os.path.exists(cache + ".ids") and \
            set(map(str, tab.source_id)) <= set(open(cache + ".ids").read().split()):
        return pd.read_csv(cache)
    # AllWISE epoch ~2010.5: move Gaia (2016.0) positions back by 5.5 yr
    dt = -5.5
    ra = tab.ra.values + np.nan_to_num(tab.pmra.values) * dt / 3.6e6 / np.cos(np.radians(tab.dec.values))
    de = tab.dec.values + np.nan_to_num(tab.pmde.values) * dt / 3.6e6
    up = pd.DataFrame(dict(cand_id=tab.source_id.values, ra_w=ra, dec_w=de))
    log("  XMatch: AllWISE for %d stars (%s)" % (len(up), tag))
    w = xmatch_upload(up, "ra_w", "dec_w", "vizier:II/328/allwise", 2.0,
                      cols2="AllWISE,W1mag,e_W1mag,W2mag,e_W2mag,W3mag,e_W3mag,W4mag,e_W4mag,ccf,ex,qph")
    w = w.sort_values("angDist").drop_duplicates("cand_id")
    w.to_csv(cache, index=False)
    open(cache + ".ids", "w").write(" ".join(map(str, tab.source_id)))
    return w


def galah(uniq):
    """GALAH DR3 parameters for the candidates (CDS XMatch, 2 arcsec; cached). Information only."""
    cache = cpath("cand_galah.csv")
    if os.path.exists(cache) and os.path.exists(cache + ".ids") and \
            set(map(str, uniq.source_id)) <= set(open(cache + ".ids").read().split()):
        return pd.read_csv(cache)
    up = uniq[["source_id", "ra", "dec"]].rename(columns={"source_id": "cand_id", "ra": "ra_c", "dec": "dec_c"})
    log("  XMatch: GALAH DR3 for %d candidates" % len(up))
    g = xmatch_upload(up, "ra_c", "dec_c", "vizier:J/MNRAS/506/150/stars", 2.0)
    if len(g):
        g = g.sort_values("angDist").drop_duplicates("cand_id")[["cand_id", "GALAH", "Teff", "logg", "[Fe/H]"]]
        g.columns = ["source_id", "galah_id", "teff_galah", "logg_galah", "feh_galah"]
    else:
        g = pd.DataFrame(columns=["source_id", "galah_id", "teff_galah", "logg_galah", "feh_galah"])
    g.to_csv(cache, index=False)
    open(cache + ".ids", "w").write(" ".join(map(str, uniq.source_id)))
    return g


def galex(unres):
    """GALEX AIS (GR6+7) photometry of the unresolved candidates: a far-UV detection would betray a hot
    (white-dwarf or subdwarf) companion contaminating the spectrum.  CDS XMatch, 3 arcsec; cached."""
    cache = cpath("galex_unresolved.csv")
    u = unres.drop_duplicates("source_id")
    if os.path.exists(cache) and os.path.exists(cache + ".ids") and \
            set(map(str, u.source_id)) <= set(open(cache + ".ids").read().split()):
        return pd.read_csv(cache)
    up = u[["source_id", "ra", "dec"]].rename(columns={"source_id": "cand_id", "ra": "ra_c", "dec": "dec_c"})
    log("  XMatch: GALEX AIS for %d unresolved candidates" % len(up))
    cols = ["source_id", "galex_sep", "galex_FUVexp", "galex_NUVexp", "galex_FUV", "galex_e_FUV", "galex_NUV", "galex_e_NUV"]
    try:
        g = xmatch_upload(up, "ra_c", "dec_c", "vizier:II/335/galex_ais", 3.0)
    except Exception as e:
        log("  GALEX XMatch failed: %r" % e)
        return pd.DataFrame(columns=cols)
    if len(g):
        g = g.sort_values("angDist").drop_duplicates("cand_id")[
            ["cand_id", "angDist", "FUVexp", "NUVexp", "FUVmag", "e_FUVmag", "NUVmag", "e_NUVmag"]]
        g.columns = cols
    else:
        g = pd.DataFrame(columns=cols)
    g.to_csv(cache, index=False)
    open(cache + ".ids", "w").write(" ".join(map(str, u.source_id)))
    return g


def irsa_allwise(unres):
    """AllWISE source catalogue (IRSA TAP, allwise_p3as_psd) for the unresolved three-way candidates:
    an independent copy of the mid-infrared photometry with S/N and quality flags (cached)."""
    cache = os.path.join(RAW, "irsa_allwise_unresolved.csv")
    have = pd.read_csv(cache) if os.path.exists(cache) else pd.DataFrame(columns=["source_id"])
    rows = [have]
    nfail = 0
    for _, r in unres.drop_duplicates("source_id").iterrows():
        if r.source_id in set(have.source_id.values):
            continue
        if nfail >= 3:                                  # IRSA down: stop asking; the CDS copy of AllWISE is used anyway
            log("  IRSA TAP unavailable: skipping the remaining candidates (they are queried on the next run)")
            break
        dt = -5.5                                       # Gaia epoch 2016.0 -> AllWISE ~2010.5
        ra = r.ra + np.nan_to_num(r.pmra) * dt / 3.6e6 / np.cos(np.radians(r.dec))
        de = r.dec + np.nan_to_num(r.pmde) * dt / 3.6e6
        q = ("SELECT designation, ra, dec, w1mpro, w1sigmpro, w2mpro, w2sigmpro, w3mpro, w3sigmpro, w3snr, "
             "w4mpro, w4sigmpro, w4snr, cc_flags, ext_flg, ph_qual FROM allwise_p3as_psd "
             "WHERE CONTAINS(POINT('ICRS',ra,dec),CIRCLE('ICRS',%.7f,%.7f,%.6f))=1" % (ra, de, 3.0 / 3600.0))
        try:
            t = pd.read_csv(io.BytesIO(tap_sync(IRSA_TAP, q, timeout=90, retries=2)), dtype={"cc_flags": str, "ph_qual": str})
            nfail = 0
        except Exception as e:                           # IRSA unreachable: keep going, the XMatch copy is used anyway
            log("  IRSA TAP failed for %d: %r" % (r.source_id, e))
            nfail += 1
            continue
        if len(t) == 0:
            t = pd.DataFrame([dict(designation="")])
        else:
            t["irsa_sep"] = 3600 * np.hypot((t.ra - ra) * np.cos(np.radians(de)), t.dec - de)
            t = t.sort_values("irsa_sep").head(1)
        t.insert(0, "source_id", r.source_id)
        rows.append(t.drop(columns=[x for x in ("ra", "dec") if x in t.columns]))
    out = pd.concat(rows, ignore_index=True)
    out.to_csv(cache, index=False)
    return out


def simbad(unres):
    """SIMBAD main identifier, object type and spectral type by Gaia DR3 identifier (cached by id)."""
    cache = os.path.join(RAW, "simbad_unresolved.csv")
    cols = ["source_id", "simbad_main_id", "simbad_otype", "simbad_sptype"]
    have = pd.read_csv(cache) if os.path.exists(cache) else pd.DataFrame(columns=cols)
    ids = [int(i) for i in unres.source_id.drop_duplicates() if i not in set(have.source_id.values)]
    rows = [have]
    for i in range(0, len(ids), 150):
        part = ids[i:i + 150]
        q = ("SELECT i.id, b.main_id, b.otype, b.sp_type FROM ident AS i JOIN basic AS b ON b.oid = i.oidref "
             "WHERE i.id IN (%s)" % ",".join("'Gaia DR3 %d'" % x for x in part))
        try:
            t = pd.read_csv(io.BytesIO(tap_sync(SIMBAD_TAP, q, timeout=180)))
        except Exception as e:
            log("  SIMBAD TAP failed: %r" % e)
            continue
        t["source_id"] = t["id"].astype(str).str.replace("Gaia DR3 ", "").astype(np.int64)
        t = t.rename(columns={"main_id": "simbad_main_id", "otype": "simbad_otype", "sp_type": "simbad_sptype"})[cols]
        # stars not in SIMBAD get an empty row so that they are not asked for again
        miss = [x for x in part if x not in set(t.source_id.values)]
        rows += [t, pd.DataFrame(dict(source_id=miss, simbad_main_id="", simbad_otype="", simbad_sptype=""))]
    out = pd.concat(rows, ignore_index=True).drop_duplicates("source_id")
    out.to_csv(cache, index=False)
    return out


def rave(unres):
    """RAVE DR6 MADERA parameters (spectroscopic, no parallax used) within 3 arcsec (cached by id).
    RAVE covers the southern sky only, so only candidates with dec < +5 deg are queried."""
    cache = os.path.join(RAW, "rave_unresolved.csv")
    cols = ["source_id", "rave_obsid", "rave_snr", "rave_teff", "rave_logg", "rave_mh", "rave_qual", "rave_n"]
    have = pd.read_csv(cache) if os.path.exists(cache) else pd.DataFrame(columns=cols)
    rows = [have]
    for _, r in unres.drop_duplicates("source_id").iterrows():
        if r.source_id in set(have.source_id.values):
            continue
        row = dict(source_id=r.source_id, rave_n=0)
        if r.dec < 5.0:
            dt = -16.0                                   # Gaia epoch 2016.0 -> J2000 catalogue positions
            ra = r.ra + np.nan_to_num(r.pmra) * dt / 3.6e6 / np.cos(np.radians(r.dec))
            de = r.dec + np.nan_to_num(r.pmde) * dt / 3.6e6
            q = ('SELECT r.ObsID, r."S/N" AS snr, m.TeffmC, m.loggmC, m."[m/H]mC" AS mh, m.Qual '
                 'FROM "III/283/ravedr6" AS r JOIN "III/283/madera" AS m ON m.ObsID = r.ObsID '
                 "WHERE 1=CONTAINS(POINT('ICRS', r.RAJ2000, r.DEJ2000), CIRCLE('ICRS', %.7f, %.7f, %.6f))"
                 % (ra, de, 3.0 / 3600.0))
            try:
                t = pd.read_csv(io.BytesIO(tap_sync(VIZIER_TAP, q, timeout=180)))
            except Exception as e:
                log("  VizieR TAP (RAVE DR6) failed for %d: %r" % (r.source_id, e))
                continue
            if len(t):
                b = t.sort_values("snr", ascending=False).iloc[0]
                row.update(rave_obsid=str(b.ObsID).strip(), rave_snr=b.snr, rave_teff=b.TeffmC, rave_logg=b.loggmC,
                           rave_mh=b.mh, rave_qual=b.Qual, rave_n=len(t))
        rows.append(pd.DataFrame([row]))
    out = pd.concat(rows, ignore_index=True).drop_duplicates("source_id")
    for c in cols:
        if c not in out.columns:
            out[c] = np.nan
    out[cols].to_csv(cache, index=False)
    return out[cols]


def run():
    C = {}
    for sv in SURVEYS:
        s = pd.read_pickle(dpkl(sv))
        C[sv] = s[s.clean].reset_index(drop=True)
        del s
    pred = {sv: KnnPredictor(C[sv]) for sv in SURVEYS}

    # ---------------------------------------------------------------- candidates
    cands = []
    for sv in SURVEYS:
        c = C[sv]
        fl = flag_levels(c)
        anyf = np.zeros(len(c), bool)
        for f in L.F_GRID:
            c["flag_%g" % f] = fl[f][0]            # primary: spectroscopic deficit AND below both main-sequence floors
            c["flagS_%g" % f] = fl[f][1]           # spectroscopic deficit only (superset)
            anyf |= fl[f][1]
        cands.append(c[anyf])
        log("%s: flagged at any f: %d spectroscopic-only, %d three-way" % (
            sv, anyf.sum(), np.any([fl[f][0] for f in L.F_GRID], axis=0).sum()))
    cand = pd.concat(cands, ignore_index=True)
    for f in L.F_GRID:
        for p in ("flag_", "flagS_"):
            cand["%s%g" % (p, f)] = cand["%s%g" % (p, f)].fillna(False).astype(bool)

    # ---------------------------------------------------------------- 1. other surveys
    for sv in SURVEYS:
        r = pd.read_pickle(os.path.join(WORK, "%s_stars.pkl" % sv))       # one survey at a time (memory)
        r = r[r.source_id.isin(set(cand.source_id.values))]
        q = np.asarray(L.survey_cuts(r, sv)[0][1], bool) & np.asarray(L.common_cuts(r)[0][1], bool)
        r = r[q][["source_id", "teff", "logg", "feh", "snr"]].rename(
            columns={"teff": "teff_" + sv, "logg": "logg_" + sv, "feh": "feh_" + sv, "snr": "snr_" + sv})
        cand = cand.merge(r, on="source_id", how="left")
        D, sig, rk = pred[sv].deficit(cand["teff_" + sv], cand["logg_" + sv], cand["feh_" + sv], cand.MG, cand.MK)
        D[rk > L.KNN_RMAX] = np.nan
        cand["D_" + sv], cand["sig_" + sv] = D, sig
        own = (cand.survey == sv).values
        cand.loc[own, ["D_" + sv, "sig_" + sv]] = np.nan          # only *other* surveys here

    # ---------------------------------------------------------------- 1b. GALAH DR3 (information only)
    # GALAH is not one of the three calibrated samples, so its parameters are passed through the LAMOST
    # calibration.  NOT used to resolve a candidate, and D_galah is NOT an independent test: GALAH DR3
    # log g comes from the parallax-based luminosity, so D_galah is near zero for any star, dimmed or not.
    # m_ms_galah uses GALAH's Teff and [Fe/H] only (no log g): negative = still below the main-sequence floor.
    gal = galah(cand.drop_duplicates("source_id"))
    cand = cand.merge(gal, on="source_id", how="left")
    Dg, sg, rkg = pred["lamost"].deficit(cand.teff_galah, cand.logg_galah, cand.feh_galah, cand.MG, cand.MK)
    cand["D_galah"], cand["sig_galah"] = Dg, sg
    cand["m_ms_galah"] = FloorPredictor(C["lamost"]).margin(cand.teff_galah, cand.feh_galah, cand.MG, cand.MK)

    # ---------------------------------------------------------------- 2. LAMOST repeat spectra
    lam_ids = set(cand.source_id[cand.survey == "lamost"])
    parts = []
    for p in dd.get_lamost():
        t = pd.read_csv(p, dtype={"GaiaDR3": str}, low_memory=False)
        t = t[t.GaiaDR3.notna() & (t.GaiaDR3.str.strip() != "")]
        t["source_id"] = t.GaiaDR3.str.strip().astype(np.int64)
        parts.append(t[t.source_id.isin(lam_ids)])
    rep = pd.concat(parts, ignore_index=True)
    rep = rep[(rep.snrg >= 30) & (rep.offsets == 0) & (rep.e_teff < 150) & (rep.e_logg < 0.25)]
    rep = rep.merge(cand.loc[cand.survey == "lamost", ["source_id", "MG", "MK", "obsid"]], on="source_id")
    rep = rep[rep.ObsID != rep.obsid]                          # other spectra than the primary one
    rep["D_rep"], rep["sig_rep"], rk = pred["lamost"].deficit(rep.teff, rep.logg, rep.feh, rep.MG, rep.MK)
    g = rep.groupby("source_id")
    repsum = pd.DataFrame(dict(n_rep=g.size(), D_rep_med=g.D_rep.median(), D_rep_max=g.D_rep.max(),
                               sig_rep_med=g.sig_rep.median())).reset_index()
    cand = cand.merge(repsum, on="source_id", how="left")
    cand["n_rep"] = cand.n_rep.fillna(0).astype(int)

    # ---------------------------------------------------------------- 3. neighbours
    uniq = cand.drop_duplicates("source_id")
    nb = neighbours(uniq)
    nb = nb[nb.Source != nb.cand_id].merge(uniq[["source_id", "G"]].rename(columns={"source_id": "cand_id", "G": "G_c"}),
                                           on="cand_id")
    nb["dG"] = nb.Gmag - nb.G_c
    g = nb.sort_values("angDist").groupby("cand_id")
    nbsum = pd.DataFrame(dict(nb_n8=g.size(), nb_sep=g.angDist.first(), nb_dG=g.dG.first())).reset_index()
    bright = nb[((nb.angDist < 5.0) & (nb.dG < 0.0)) | ((nb.angDist < 2.5) & (nb.dG < 1.0))]
    bsum = bright.sort_values("dG").groupby("cand_id").agg(bnb_sep=("angDist", "first"), bnb_dG=("dG", "first")).reset_index()
    cand = cand.merge(nbsum.rename(columns={"cand_id": "source_id"}), on="source_id", how="left")
    cand = cand.merge(bsum.rename(columns={"cand_id": "source_id"}), on="source_id", how="left")

    # ---------------------------------------------------------------- 4/5. Gaia extras
    ex = gaia_extra(list(uniq.source_id.values))
    ex = ex.rename(columns={"non_single_star": "nss", "phot_variable_flag": "varflag_dr3"})
    cand = cand.merge(ex[["source_id", "nss", "varflag_dr3", "plx_dr2", "e_plx_dr2", "dr2_sep"]], on="source_id", how="left")
    cand["plx_ratio_dr2"] = cand.plx_dr2 / cand.plx

    # ---------------------------------------------------------------- 6. AllWISE (candidates + control)
    w = allwise(uniq, "cand").rename(columns={"cand_id": "source_id"})
    ctrl = pd.concat([C[sv].sample(n=min(15000, len(C[sv])), random_state=1) for sv in SURVEYS]).drop_duplicates("source_id")
    wc = allwise(ctrl, "control").rename(columns={"cand_id": "source_id"})
    wc = wc.merge(ctrl[["source_id", "K", "bprp0", "teff"]], on="source_id")
    okc = (wc.ccf.astype(str).str[:2] == "00") & (wc.ex == 0)
    wc = wc[okc]
    # photospheric loci: median K-W1, K-W2 and W1-W3 / W1-W4 for S/N>3 detections, versus colour
    loci = {}
    for a, b, lab in (("K", "W1mag", "K_W1"), ("K", "W2mag", "K_W2"), ("W1mag", "W3mag", "W1_W3"), ("W1mag", "W4mag", "W1_W4")):
        col = wc[a] - wc[b]
        ee = wc["e_" + b]
        good = col.notna() & ee.notna() & (ee < (0.1 if lab.startswith("K") else 0.3))
        xb = np.arange(0.4, 1.81, 0.1)
        ii = np.digitize(wc.bprp0[good], xb)
        med = np.array([np.median(col[good][ii == j]) if np.sum(ii == j) > 30 else np.nan for j in range(1, len(xb))])
        sd = np.array([L.robust_sigma(col[good][ii == j]) if np.sum(ii == j) > 30 else np.nan for j in range(1, len(xb))])
        loci[lab] = (0.5 * (xb[1:] + xb[:-1]), med, sd)
    pd.DataFrame({"bprp0": loci["K_W1"][0], **{k + "_med": v[1] for k, v in loci.items()},
                  **{k + "_sig": v[2] for k, v in loci.items()}}).to_csv(os.path.join(DATA, "wise_control.csv"), index=False)
    cand = cand.merge(w[["source_id", "AllWISE", "W1mag", "e_W1mag", "W2mag", "e_W2mag", "W3mag", "e_W3mag", "W4mag",
                         "e_W4mag", "ccf", "ex"]], on="source_id", how="left")

    def excess(a, b, lab):
        cen, med, sd = loci[lab]
        ok = np.isfinite(med)
        exp_ = np.interp(cand.bprp0, cen[ok], med[ok])
        sdv = np.interp(cand.bprp0, cen[ok], sd[ok])
        col = cand[a] - cand[b]
        err = np.sqrt(cand["e_" + b].fillna(9.0) ** 2 + sdv ** 2 + (cand.e_K.fillna(0.03) ** 2 if a == "K" else 0.0))
        return col - exp_, (col - exp_) / err
    for a, b, lab in (("K", "W1mag", "K_W1"), ("K", "W2mag", "K_W2"), ("W1mag", "W3mag", "W1_W3"), ("W1mag", "W4mag", "W1_W4")):
        cand["x_" + lab], cand["xs_" + lab] = excess(a, b, lab)
    cand["wise_ok"] = cand.W1mag.notna() & (cand.ccf.astype(str).str[:2] == "00") & (cand.ex.fillna(1) == 0)
    cand["ir_excess"] = cand.wise_ok & (((cand.xs_K_W2 > 4) & (cand.x_K_W2 > 0.15))
                                        | ((cand.xs_W1_W3 > 4) & (cand.x_W1_W3 > 0.3) & (cand.e_W3mag < 0.3))
                                        | ((cand.xs_W1_W4 > 4) & (cand.x_W1_W4 > 1.0) & (cand.e_W4mag < 0.3)))

    # ---------------------------------------------------------------- classification per f
    summ = []
    for f in L.F_GRID:
        Df = np.log10(1 - f)
        fl = cand["flagS_%g" % f].values          # every spectroscopically flagged star gets a status
        status = np.full(len(cand), "", dtype=object)
        # (a) not reproduced with another survey's parameters
        notrep = np.zeros(len(cand), bool)
        rep2 = np.zeros(len(cand), bool)
        for sv in SURVEYS:
            D2, s2 = cand["D_" + sv].values, cand["sig_" + sv].values
            has = np.isfinite(D2)
            notrep |= has & (D2 >= Df + L.NSIG_THR * s2)
            rep2 |= has & (D2 < Df + L.NSIG_THR * s2)
        # (b) not reproduced in LAMOST repeat spectra (median of the other good spectra)
        notrep_l = (cand.n_rep.values >= 1) & (cand.D_rep_med.values >= Df + L.NSIG_THR * cand.sig_rep_med.values)
        rep_l = (cand.n_rep.values >= 1) & ~notrep_l
        # (c) brighter Gaia neighbour inside/near the fibre
        blend = cand.bnb_sep.notna().values
        # (d) DR2 and DR3 parallaxes disagree by enough to remove the deficit at this f
        #     (a parallax over-estimate by factor q lowers the inferred luminosity by q^2)
        q = cand.plx.values / cand.plx_dr2.values
        dsig = (cand.plx.values - cand.plx_dr2.values) / np.hypot(cand.e_plx.values, cand.e_plx_dr2.values)
        plxbad = np.isfinite(q) & (dsig > 3.0) & (q > 1.1) & \
            (cand.D.values + 2 * np.log10(np.maximum(q, 1e-3)) >= Df + L.NSIG_THR * cand.sigma_i.values)
        status[fl] = "unresolved"
        status[fl & plxbad] = "resolved: DR2 parallax differs by >3 sigma and removes deficit"
        status[fl & blend] = "resolved: brighter Gaia neighbour at fibre"
        status[fl & notrep_l] = "resolved: not reproduced in repeat LAMOST spectra"
        status[fl & notrep] = "resolved: not reproduced with second survey parameters"
        cand["status_%g" % f] = status
        cand["confirmed2_%g" % f] = fl & (rep2 | rep_l) & ~notrep & ~notrep_l
        for tier, col in (("three-way", "flag_%g" % f), ("spectroscopic-only", "flagS_%g" % f)):
            for sv in SURVEYS + ("any",):
                m = cand[col].values if sv == "any" else cand[col].values & (cand.survey == sv).values
                vc = pd.Series(status[m]).value_counts()
                row = dict(tier=tier, survey=sv, f=f, n_flagged=int(m.sum()), n_unresolved=int(vc.get("unresolved", 0)),
                           n_unres_measured_twice=int((m & (status == "unresolved") & cand["confirmed2_%g" % f].values).sum()))
                for k, v in vc.items():
                    if k != "unresolved":
                        row[k] = int(v)
                summ.append(row)
    summ = pd.DataFrame(summ).fillna(0)
    summ.to_csv(os.path.join(DATA, "vetting_summary.csv"), index=False)
    print(summ.to_string(index=False))

    cand["astrom_suspect"] = (cand.ruwe > 1.25) | (cand.sepsi > 3) | (cand.nss.fillna(0) > 0) | \
        ((np.abs(cand.plx_ratio_dr2 - 1) > 0.15) & (cand.plx_dr2 / cand.e_plx_dr2 > 5))
    cand["variable"] = cand.varflag_dr3.astype(str).str.upper().eq("VARIABLE")
    cand["vtan"] = 4.74 * np.hypot(cand.pmra, cand.pmde) / cand.plxc
    keep = ["source_id", "survey", "ra", "dec", "glon", "glat", "teff", "logg", "feh", "snr", "G", "bprp0", "K", "plx",
            "poe", "dist", "vtan", "ruwe", "sepsi", "epsi", "sfd", "MG", "MK", "D_G", "D_K", "D", "sigma_i", "D_ms", "m_ms", "D_cmd", "m_cmd",
            "z_bprp", "z_gk", "cls", "ms"]
    keep += ["flag_%g" % f for f in L.F_GRID] + ["flagS_%g" % f for f in L.F_GRID]
    keep += ["status_%g" % f for f in L.F_GRID] + ["confirmed2_%g" % f for f in L.F_GRID]
    keep += [x for sv in SURVEYS for x in ("teff_" + sv, "logg_" + sv, "feh_" + sv, "D_" + sv, "sig_" + sv)]
    keep += ["galah_id", "teff_galah", "logg_galah", "feh_galah", "D_galah", "sig_galah", "m_ms_galah", "pmra", "pmde"]
    keep += ["n_rep", "D_rep_med", "sig_rep_med", "nb_n8", "nb_sep", "nb_dG", "bnb_sep", "bnb_dG", "nss", "varflag_dr3",
             "plx_dr2", "e_plx_dr2", "plx_ratio_dr2", "astrom_suspect", "variable", "AllWISE", "W1mag", "W2mag", "W3mag",
             "e_W3mag", "W4mag", "e_W4mag", "x_K_W1", "x_K_W2", "xs_K_W2", "x_W1_W3", "xs_W1_W3", "x_W1_W4", "xs_W1_W4",
             "wise_ok", "ir_excess"]
    keep += [x for x in ("apogee_id", "obsid", "nspec", "teff_spread", "logg_spread", "nvis", "vscatter", "Prog",
                         "rvamp", "teff_gphot", "vbroad") if x in cand.columns]
    cand = cand[keep].sort_values("D")
    cand.to_csv(os.path.join(DATA, "candidates_all.csv"), index=False)
    un3 = np.zeros(len(cand), bool)
    for f in L.F_GRID:
        un3 |= cand["flag_%g" % f].values & (cand["status_%g" % f] == "unresolved").values
    unres = cand[un3]
    iw = irsa_allwise(unres)
    if len(iw):
        iw = iw.rename(columns={c: "irsa_" + c for c in iw.columns if c not in ("source_id", "irsa_sep")})
        unres = unres.merge(iw.drop_duplicates("source_id"), on="source_id", how="left")
    gx = galex(unres)
    if len(gx):
        unres = unres.merge(gx.drop_duplicates("source_id"), on="source_id", how="left")
    unres = unres.merge(simbad(unres), on="source_id", how="left")
    unres = unres.merge(rave(unres), on="source_id", how="left")
    # the RAVE spectroscopic parameters through the star's own survey calibration (information only)
    Dr = np.full(len(unres), np.nan)
    for sv in SURVEYS:
        q = (unres.survey == sv).values & unres.rave_teff.notna().values
        if q.any():
            Dr[q] = pred[sv].deficit(unres.rave_teff[q], unres.rave_logg[q], unres.rave_mh[q], unres.MG[q], unres.MK[q])[0]
    unres["D_rave"] = Dr
    unres.to_csv(os.path.join(DATA, "candidates_unresolved.csv"), index=False)
    log("candidates: %d flagged rows (%d unique stars); %d unresolved rows" % (len(cand), cand.source_id.nunique(), len(unres)))
    return cand


if __name__ == "__main__":
    run()