"""Other catalogues: what they can and cannot say about FGK dwarfs.
APOGEE DR17 (Ce only), Gaia DR3 GSP-Spec (Ce, Nd, Zr), GALAH DR3 (independent pipeline, same spectra),
plus the fate of known chemically peculiar (Renson & Manfroid 2009) stars and of Przybylski's Star."""
import os, csv, json, urllib.request, urllib.parse
import numpy as np
import pandas as pd
from astropy.io import fits
from astropy.coordinates import SkyCoord
import astropy.units as u
from common import *
import s03_candidates as s3

GAIA_TAP = 'https://gea.esac.esa.int/tap-server/tap/sync'
VIZ_TAP = 'https://tapvizier.cds.unistra.fr/TAPVizieR/tap/sync'


def tap(url, query, out, timeout=600):
    if os.path.exists(out) and os.path.getsize(out) > 200:
        return out
    data = urllib.parse.urlencode({'REQUEST': 'doQuery', 'LANG': 'ADQL', 'FORMAT': 'csv', 'QUERY': query}).encode()
    with urllib.request.urlopen(urllib.request.Request(url, data=data), timeout=timeout) as r:
        open(out, 'wb').write(r.read())
    return out


def fgk(teff, logg):
    return (logg >= LOGG_MIN) & (teff >= TEFF_MIN) & (teff < TEFF_MAX)


def apogee(d4):
    p = RAW + 'allStarLite-dr17-synspec_rev1.fits'
    h = fits.open(p, memmap=True); t = h[1].data
    teff = np.array(t['TEFF'], float); logg = np.array(t['LOGG'], float); ce = np.array(t['CE_FE'], float)
    cef = np.array(t['CE_FE_FLAG']); snr = np.array(t['SNR'], float); aspcap = np.array(t['ASPCAPFLAG'])
    gid = np.array(t['GAIAEDR3_SOURCE_ID']).astype('i8')
    ok = np.isfinite(ce) & (cef == 0) & (snr > 70) & ((aspcap & (1 << 23)) == 0)
    dw = fgk(teff, logg) & ok; gi = (logg < 3.5) & np.isfinite(teff) & ok
    res = dict(n_rows=int(len(t)), n_ce_all=int(np.isfinite(ce).sum()), n_ce_dwarf_good=int(dw.sum()),
               n_ce_giant_good=int(gi.sum()))
    for thr in (0.5, 1.0):
        res['dwarf_ce_gt_%.1f' % thr] = int((dw & (ce > thr)).sum())
        res['giant_ce_gt_%.1f' % thr] = int((gi & (ce > thr)).sum())
    res['dwarf_teff_median'] = float(np.median(teff[dw])); res['dwarf_teff_p10_p90'] = [float(v) for v in np.percentile(teff[dw], [10, 90])]
    # overlap with GALAH DR4 FGK dwarfs
    g4 = d4['gaiadr3_source_id'].astype('i8')
    o = np.argsort(g4); pos = np.clip(np.searchsorted(g4[o], gid), 0, len(o) - 1); j = o[pos]
    common = (g4[j] == gid) & (gid > 0) & dw
    res['dwarf_ce_in_galah_dr4'] = int(common.sum())
    x4 = d4['ce_fe'][j].astype(float); f4 = d4['flag_ce_fe'][j]
    both = common & (f4 == 0) & np.isfinite(x4)
    res['dwarf_ce_both_detected'] = int(both.sum())
    if both.sum() > 5:
        dif = x4[both] - ce[both]
        res['median_galah_minus_apogee_ce'] = float(np.median(dif)); res['mad_sigma'] = float(1.4826 * np.median(np.abs(dif - np.median(dif))))
    # is Przybylski's Star (HD 101065) in APOGEE DR17?
    ra = np.array(t['RA'], float); de = np.array(t['DEC'], float)
    near = (np.abs(de + 46.70968829912) < 1.0) & (np.abs(ra - 174.40433733673) < 1.5)
    if near.sum():
        sep = SkyCoord(ra[near] * u.deg, de[near] * u.deg).separation(SkyCoord(174.40433733673 * u.deg, -46.70968829912 * u.deg)).arcsec
        res['przybylski_nearest_apogee_arcsec'] = float(sep.min())
    else:
        res['przybylski_nearest_apogee_arcsec'] = None   # nothing within ~1 degree
    # APOGEE dwarfs with Ce > 1: list (NOT vetted at spectrum level; described at catalogue level only)
    hi = np.where(dw & (ce > 1.0))[0]
    mh = np.array(t['M_H'], float)
    res['dwarf_ce_gt_1_mh_lt_m1'] = int((mh[hi] < -1.0).sum())
    res['dwarf_ce_gt_1_mh_lt_m0p4'] = int((mh[hi] < -0.4).sum())
    res['dwarf_ce_gt_1_mh_missing'] = int((~np.isfinite(mh[hi])).sum())
    res['dwarf_ce_gt_1_in_galah_dr4'] = int(sum(((g4[j[i]] == gid[i]) and gid[i] > 0) for i in hi))
    with open(DATA + 'apogee_dwarfs_ce_gt1.csv', 'w', newline='') as fh:
        w = csv.writer(fh); w.writerow(['apogee_id', 'ra', 'dec', 'teff', 'logg', 'm_h', 'ce_fe', 'ce_fe_err', 'snr', 'c_fe', 'galah_dr4_ce_fe', 'galah_dr4_flag_ce', 'galah_dr4_ba_fe'])
        for i in hi:
            inG = (g4[j[i]] == gid[i]) and gid[i] > 0
            w.writerow([t['APOGEE_ID'][i], f"{t['RA'][i]:.5f}", f"{t['DEC'][i]:.5f}", f'{teff[i]:.0f}', f'{logg[i]:.2f}',
                        f"{t['M_H'][i]:.2f}", f'{ce[i]:.2f}', f"{t['CE_FE_ERR'][i]:.2f}", f'{snr[i]:.0f}', f"{t['C_FE'][i]:.2f}",
                        f'{d4["ce_fe"][j[i]]:.2f}' if inG else '', int(d4['flag_ce_fe'][j[i]]) if inG else '',
                        f'{d4["ba_fe"][j[i]]:.2f}' if inG else ''])
    return res


def gspspec(d4=None):
    lists = []
    out = RAW + 'gspspec_ncap.csv'
    q = ("SELECT a.source_id, g.ra, g.dec, g.phot_g_mean_mag, g.bp_rp, g.parallax, g.ruwe, g.rvs_spec_sig_to_noise, "
         "g.non_single_star, a.teff_gspspec, a.logg_gspspec, a.mh_gspspec, a.alphafe_gspspec, a.fem_gspspec, "
         "a.flags_gspspec, a.logchisq_gspspec, a.cefe_gspspec, a.cefe_gspspec_lower, a.cefe_gspspec_upper, "
         "a.cefe_gspspec_nlines, a.cefe_gspspec_linescatter, a.ndfe_gspspec, a.ndfe_gspspec_lower, a.ndfe_gspspec_upper, "
         "a.ndfe_gspspec_nlines, a.ndfe_gspspec_linescatter, a.zrfe_gspspec, a.zrfe_gspspec_lower, a.zrfe_gspspec_upper, "
         "a.zrfe_gspspec_nlines, a.zrfe_gspspec_linescatter FROM gaiadr3.astrophysical_parameters AS a JOIN "
         "gaiadr3.gaia_source AS g ON g.source_id=a.source_id WHERE a.teff_gspspec IS NOT NULL AND "
         "(a.cefe_gspspec IS NOT NULL OR a.ndfe_gspspec IS NOT NULL OR a.zrfe_gspspec IS NOT NULL)")
    tap(GAIA_TAP, q, out)
    g = pd.read_csv(out, dtype={'flags_gspspec': str})
    dw = fgk(g.teff_gspspec, g.logg_gspspec)
    res = dict(n_rows=int(len(g)), n_dwarf_any=int(dw.sum()))
    # flag chain positions (Recio-Blanco et al. 2023, Table 2): Zr 34-35, Ce 36-37, Nd 38-39 (1-based)
    for el, p0 in (('zr', 33), ('ce', 35), ('nd', 37)):
        x = g[el + 'fe_gspspec']; ok = x.notna()
        up = g.flags_gspspec.str[p0].astype(int); un = g.flags_gspspec.str[p0 + 1].astype(int)
        good = ok & (up <= 1) & (un <= 1)
        gd = good & dw
        res[el] = dict(n=int(ok.sum()), n_dwarf=int((ok & dw).sum()), n_goodflag=int(good.sum()),
                       n_goodflag_dwarf=int(gd.sum()),
                       n_goodflag_giant_gt1=int((good & ~dw & (x > 1.0)).sum()),
                       n_goodflag_dwarf_gt1=int((gd & (x > 1.0)).sum()),
                       # revision: these 'good-flag' dwarf values are almost all AT the upper edge of the
                       # GSP-Spec abundance grid (+2.0 dex) and rest on a single line
                       n_goodflag_dwarf_at_grid_edge=int((gd & (x >= 1.99)).sum()),
                       n_goodflag_dwarf_single_line=int((gd & (g[el + 'fe_gspspec_nlines'] <= 1)).sum()),
                       n_goodflag_dwarf_with_public_rvs_snr=int((gd & g.rvs_spec_sig_to_noise.notna()).sum()))
        if gd.sum():
            o = g[gd][['source_id', 'ra', 'dec', 'phot_g_mean_mag', 'teff_gspspec', 'logg_gspspec', 'mh_gspspec', 'flags_gspspec',
                       el + 'fe_gspspec', el + 'fe_gspspec_lower', el + 'fe_gspspec_upper', el + 'fe_gspspec_nlines',
                       'rvs_spec_sig_to_noise', 'ruwe', 'non_single_star']].copy()
            o.insert(0, 'element', el)
            o.columns = ['element', 'source_id', 'ra', 'dec', 'G', 'teff', 'logg', 'm_h', 'flags_gspspec', 'x_fe', 'x_fe_lower',
                         'x_fe_upper', 'n_lines', 'rvs_snr', 'ruwe', 'non_single_star']
            lists.append(o)
    if lists:
        L = pd.concat(lists, ignore_index=True)
        if d4 is not None:
            g4 = set(d4['gaiadr3_source_id'].astype('i8').tolist())
            L['in_galah_dr4'] = [int(int(sid) in g4) for sid in L.source_id]
            res['goodflag_dwarfs_in_galah_dr4'] = int(L.in_galah_dr4.sum())
        L.to_csv(DATA + 'gspspec_goodflag_dwarfs_unvetted.csv', index=False)
    return res


RVS_LINE = {'ce': 851.375, 'nd': 859.389}     # nm, vacuum: the single lines GSP-Spec uses (Contursi et al. 2023, 2024)


def gspspec_rvs_check():
    """Revision: quick spectrum-level look at the 'good-flag' GSP-Spec dwarfs that sit at the +2.0 dex
    grid edge.  Only a minority have a public Gaia RVS mean spectrum; for those the absorption at
    the line position (mean flux within +-0.04 nm relative to the 90th percentile of the flux within
    +-0.6 nm) is compared with that of up to four comparison dwarfs of the same GSP-Spec
    Teff/logg/[M/H] taken in order of Gaia random_index.  A +2 dex enhancement should give a line
    far deeper than in the comparison stars."""
    import io
    lp = DATA + 'gspspec_goodflag_dwarfs_unvetted.csv'
    if not os.path.exists(lp):
        return {}
    L = pd.read_csv(lp)
    T = L[L.rvs_snr.notna() & (L.x_fe >= 1.99)].reset_index(drop=True)
    os.makedirs(RAW + 'rvs', exist_ok=True)

    def spec(sid):
        p_ = RAW + 'rvs/%d.csv' % sid
        if not (os.path.exists(p_) and os.path.getsize(p_) > 1000):
            url = ('https://gea.esac.esa.int/data-server/data?ID=Gaia+DR3+%d&RETRIEVAL_TYPE=RVS&DATA_STRUCTURE=INDIVIDUAL&FORMAT=csv' % sid)
            try:
                with urllib.request.urlopen(url, timeout=120) as r:
                    open(p_, 'wb').write(r.read())
            except Exception:
                return None
        try:
            t = pd.read_csv(p_)
            return t.wavelength.values.astype(float), t.flux.values.astype(float)
        except Exception:
            return None

    def depth(sp, w0):
        w, f = sp
        core = np.abs(w - w0) <= 0.04; wide = np.abs(w - w0) <= 0.6
        if core.sum() < 3 or wide.sum() < 20:
            return np.nan
        return float(1.0 - np.nanmean(f[core]) / np.nanpercentile(f[wide], 90))

    rows = []
    for _, r in T.iterrows():
        cp = RAW + 'rvs/comp_%d.csv' % r.source_id
        if not (os.path.exists(cp) and os.path.getsize(cp) > 50):
            q = ("SELECT TOP 4 g.source_id, g.random_index FROM gaiadr3.gaia_source AS g JOIN gaiadr3.astrophysical_parameters AS a "
                 "ON a.source_id=g.source_id WHERE g.has_rvs='t' AND g.rvs_spec_sig_to_noise>80 AND a.teff_gspspec BETWEEN %d AND %d "
                 "AND a.logg_gspspec BETWEEN %.2f AND %.2f AND a.mh_gspspec BETWEEN %.2f AND %.2f AND a.flags_gspspec LIKE '0000000000000%%' "
                 "AND g.random_index BETWEEN 0 AND 600000000 ORDER BY g.random_index"
                 % (r.teff - 120, r.teff + 120, r.logg - 0.25, r.logg + 0.25, r.m_h - 0.15, r.m_h + 0.15))
            try:
                tap(GAIA_TAP, q, cp)
            except Exception:
                pass
        comp = pd.read_csv(cp) if os.path.exists(cp) else pd.DataFrame(dict(source_id=[]))
        w0 = RVS_LINE[r.element]
        s_ = spec(int(r.source_id))
        dc = [depth(x, w0) for x in (spec(int(c_)) for c_ in comp.source_id) if x is not None]
        dc = [x for x in dc if np.isfinite(x)]
        rows.append(dict(element=r.element, source_id=int(r.source_id), teff=r.teff, logg=r.logg, m_h=r.m_h, rvs_snr=r.rvs_snr,
                         catalogue_x_fe=r.x_fe, line_nm_vacuum=w0, depth_at_line=depth(s_, w0) if s_ is not None else np.nan,
                         n_comparison=len(dc), comparison_depth_median=float(np.median(dc)) if dc else np.nan,
                         comparison_depth_max=float(np.max(dc)) if dc else np.nan))
    out = pd.DataFrame(rows)
    # how many FGK dwarfs/subgiants did GSP-Spec analyse at all?  (the trials factor behind the grid-edge values)
    n_dw = None
    try:
        q0 = ("SELECT COUNT(*) AS n FROM gaiadr3.astrophysical_parameters WHERE teff_gspspec >= %d AND teff_gspspec < %d "
              "AND logg_gspspec >= %.1f" % (TEFF_MIN, TEFF_MAX, LOGG_MIN))
        n_dw = int(pd.read_csv(tap(GAIA_TAP, q0, RAW + 'gspspec_fgk_dwarf_count.csv', timeout=280)).iloc[0]['n'])
    except Exception:
        pass
    if len(out):
        out['excess_depth'] = out.depth_at_line - out.comparison_depth_median
        out.to_csv(DATA + 'gspspec_rvs_check.csv', index=False, float_format='%.4f')
        ok = out[np.isfinite(out.excess_depth)]
        return dict(n_fgk_dwarfs_with_gspspec_parameters=n_dw,
                    n_grid_edge_dwarfs=int((L.x_fe >= 1.99).sum()), n_with_public_rvs=int(len(T)), n_checked=int(len(ok)),
                    median_depth_at_line=float(ok.depth_at_line.median()) if len(ok) else None,
                    median_comparison_depth=float(ok.comparison_depth_median.median()) if len(ok) else None,
                    max_depth_at_line=float(ok.depth_at_line.max()) if len(ok) else None,
                    max_excess_depth_over_comparison=float(ok.excess_depth.max()) if len(ok) else None,
                    median_excess_depth=float(ok.excess_depth.median()) if len(ok) else None)
    return dict(n_fgk_dwarfs_with_gspspec_parameters=n_dw, n_grid_edge_dwarfs=int((L.x_fe >= 1.99).sum()), n_with_public_rvs=0, n_checked=0)


def przybylski():
    q = ("SELECT g.source_id, g.phot_g_mean_mag, g.bp_rp, g.radial_velocity, g.has_rvs, g.teff_gspphot, "
         "a.teff_gspspec, a.cefe_gspspec, a.ndfe_gspspec FROM gaiadr3.gaia_source g JOIN "
         "gaiadr3.astrophysical_parameters a USING (source_id) WHERE g.source_id=5372587514128271232")
    out = tap(GAIA_TAP, q, RAW + 'przybylski_gaia.csv')
    r = pd.read_csv(out).iloc[0].to_dict()
    q2 = ("SELECT COUNT(*) AS n, COUNT(radial_velocity) AS n_rv FROM gaiadr3.gaia_source WHERE "
          "phot_g_mean_mag BETWEEN 7.5 AND 8.5 AND bp_rp BETWEEN 0.6 AND 0.85")
    r2 = pd.read_csv(tap(GAIA_TAP, q2, RAW + 'bright_fg_rv_counts.csv')).iloc[0].to_dict()
    q3 = ("SELECT COUNT(*) AS n, COUNT(a.teff_gspspec) AS n_gspspec FROM gaiadr3.gaia_source AS g JOIN "
          "gaiadr3.astrophysical_parameters AS a ON a.source_id=g.source_id WHERE g.phot_g_mean_mag BETWEEN 7.5 AND 8.5 "
          "AND g.bp_rp BETWEEN 0.6 AND 0.85")
    r3 = pd.read_csv(tap(GAIA_TAP, q3, RAW + 'bright_fg_gspspec_counts.csv')).iloc[0].to_dict()
    return dict(gaia={k: (None if (isinstance(v, float) and np.isnan(v)) else (v.item() if hasattr(v, 'item') else v)) for k, v in r.items()},
                bright_fg_n=int(r2['n']), bright_fg_with_rv=int(r2['n_rv']), bright_fg_with_gspspec=int(r3['n_gspspec']))


def renson(d4):
    out = RAW + 'renson2009.csv'
    tap(VIZ_TAP, 'SELECT * FROM "III/260/catalog"', out)
    r = pd.read_csv(out)
    cg = SkyCoord(d4['ra'] * u.deg, d4['dec'] * u.deg)
    cr = SkyCoord(r.RAJ2000.values * u.deg, r.DEJ2000.values * u.deg)
    idx, sep, _ = cr.match_to_catalog_sky(cg)
    m = sep.arcsec < 10
    o = r[m].copy().reset_index(drop=True); j = idx[m]
    for col in ['sobject_id', 'teff', 'logg', 'fe_h', 'vsini', 'flag_sp', 'chi2_sp', 'snr_px_ccd2']:
        o[col] = d4[col][j]
    for e in ['y', 'zr', 'ba', 'la', 'ce', 'nd', 'sm', 'eu']:
        o[e + '_fe'] = d4[e + '_fe'][j]; o['flag_' + e] = d4['flag_' + e + '_fe'][j]
    cols = ['HD', 'SpType', 'Vmag', 'sobject_id', 'teff', 'logg', 'fe_h', 'vsini', 'flag_sp', 'chi2_sp', 'snr_px_ccd2'] + \
           [k for e in ['y', 'zr', 'ba', 'la', 'ce', 'nd', 'sm', 'eu'] for k in (e + '_fe', 'flag_' + e)]
    o[cols].round(3).to_csv(DATA + 'renson_cp_stars_in_galah_dr4.csv', index=False)
    lan = np.zeros(len(o), int)
    for e in ['la', 'ce', 'nd']:
        lan += (((o['flag_' + e] == 0) | (o['flag_' + e] == 8)) & ((o[e + '_fe'] >= 0.9) | (o['flag_' + e] == 8))).astype(int)
    anyfit = (o.flag_la != 2) | (o.flag_ce != 2) | (o.flag_nd != 2)
    pz = SkyCoord(174.40433733673 * u.deg, -46.70968829912 * u.deg)
    return dict(n_renson=int(len(r)), n_in_galah_dr4=int(m.sum()), n_teff_lt_7200=int((o.teff < 7200).sum()),
                n_any_lanthanide_fitted=int(anyfit.sum()), n_lanthanide_ge_0p9=int((lan >= 1).sum()),
                n_two_lanthanides_ge_0p9=int((lan >= 2).sum()),
                przybylski_nearest_galah_arcsec=float(pz.separation(cg).arcsec.min()))


def cp_onset(d4):
    """Fraction of logg>=3.5, flag_sp=0, S/N>50 stars with >=2 of (La,Ce,Nd) reported >=0.9 dex, vs Teff."""
    teff = d4['teff'].astype(float); logg = d4['logg'].astype(float)
    ok = (logg >= 3.5) & (d4['flag_sp'] == 0) & (d4['snr_px_ccd2'] > 50)
    multi = np.zeros(len(teff), int)
    for e in ['nd', 'la', 'ce']:
        x = d4[e + '_fe'].astype(float); f = d4['flag_' + e + '_fe']
        multi += (((f == 0) | (f == 8)) & ((x >= 0.9) | (f == 8))).astype(int)
    # revision (referee major point 3): the raw fraction is dominated by sensitivity.  n_sensitive is
    # the number of stars whose own 90% detection thresholds are at or below +0.9 dex for at least
    # two of La, Ce, Nd (only defined below 6500 K, where the detection model of s02 was fitted);
    # the fraction among sensitive stars is the meaningful one.  Above 6500 K the catalogue values
    # are unvetted and the sensitivity is not modelled, so those rows are raw counts only.
    nsens = np.zeros(len(teff), int)
    for e in ['nd', 'la', 'ce']:
        if 'x90_' + e in d4:
            nsens += (d4['x90_' + e].astype(float) <= 0.9).astype(int)
    rows = []
    for t0 in np.arange(4500, 7751, 250):
        m = ok & (teff >= t0) & (teff < t0 + 250)
        k = int((m & (multi >= 2)).sum()); n = int(m.sum())
        if t0 < TEFF_MAX:
            ns = int((m & (nsens >= 2)).sum()); ks = int((m & (nsens >= 2) & (multi >= 2)).sum())
            fs = ks / max(ns, 1); us = binom_upper(ks, ns) if ns else np.nan
        else:
            ns = ks = ''; fs = us = np.nan
        rows.append((t0, t0 + 250, n, k, k / max(n, 1), binom_upper(k, n) if n else np.nan, ns, ks, fs, us))
    with open(DATA + 'cp_onset_vs_teff.csv', 'w', newline='') as fh:
        w = csv.writer(fh); w.writerow(['teff_lo', 'teff_hi', 'n_stars', 'n_two_lanthanides_ge_0p9', 'raw_fraction', 'raw_fraction_95ul',
                                        'n_sensitive_to_two_lanthanides', 'n_found_among_sensitive', 'fraction_among_sensitive',
                                        'fraction_among_sensitive_95ul'])
        for r in rows:
            w.writerow([r[0], r[1], r[2], r[3], f'{r[4]:.6f}', f'{r[5]:.6f}', r[6], r[7],
                        f'{r[8]:.6f}' if np.isfinite(r[8]) else '', f'{r[9]:.6f}' if np.isfinite(r[9]) else ''])
    return rows


def completeness_s_rich(d, c):
    """Catalogue completeness for REAL strong absorbers, without DR3 (the referee's independent
    check, reproduced here): among clean dwarfs that are securely rich in both Ba and Y (so their
    lanthanides must be enhanced) and whose own detection threshold lies at least 0.25 dex below
    the predicted abundance, how often does DR4 report a secure detection, and one within 0.4 dex
    of the prediction?  These are natural s-process stars, in which several elements rise together;
    the response to a single-element anomaly is not tested by this."""
    clean = c['base'] & (d['flag_sp'] == 0)
    rich = clean & c['sec']['ba'] & c['sec']['y'] & (d['ba_fe'] >= 0.6) & (d['y_fe'] >= 0.6)
    out = {}
    for e in ['nd', 'la']:
        pred = c['pred'][e]; x = d[e + '_fe'].astype(float)
        m = rich & np.isfinite(pred) & (d['x90_' + e] <= pred - 0.25)
        n = int(m.sum()); k = int((m & c['sec'][e]).sum()); k2 = int((m & c['sec'][e] & (x >= pred - 0.4)).sum())
        out[e] = dict(n=n, n_secure=k, frac_secure=k / max(n, 1), frac_secure_lower95=1 - binom_upper(n - k, n) if n else None,
                      n_secure_within_0p4=k2, frac_secure_within_0p4=k2 / max(n, 1),
                      frac_secure_within_0p4_lower95=1 - binom_upper(n - k2, n) if n else None)
    return out


def extreme_census(d, c):
    """Przybylski-like census: dwarfs/subgiants with several lanthanides reported at or above the
    top of the model range, with and without the clean-spectrum requirement (a grossly peculiar
    spectrum may itself trip flag_sp)."""
    teff = d['teff'].astype(float); logg = d['logg'].astype(float)
    L = ['la', 'ce', 'nd', 'sm', 'eu']
    rows = []
    for thr in (0.9, 1.2):
        nh = np.zeros(len(teff), int)
        for e in L:
            x = d[e + '_fe'].astype(float); f = d['flag_' + e + '_fe']
            nh += ((((f == 0) & (x >= min(thr, DR4_UPPER[e] - 0.05))) | (((f & 8) > 0) & ((f & 1) == 0)))).astype(int)
        for lab, m in [('4500-6500 K, flag_sp == 0', c['base'] & (d['flag_sp'] == 0)),
                       ('4500-6500 K, flag_sp != 0', c['base'] & (d['flag_sp'] != 0)),
                       ('6500-7000 K, any flag_sp', (teff >= 6500) & (teff < 7000) & (logg >= 3.5)),
                       ('>= 7000 K, any flag_sp', (teff >= 7000) & (logg >= 3.5))]:
            rows.append(dict(threshold_dex=thr, sample=lab, n_stars=int(m.sum()), n_ge2=int((m & (nh >= 2)).sum()),
                             n_ge3=int((m & (nh >= 3)).sum()), n_ge4=int((m & (nh >= 4)).sum())))
    pd.DataFrame(rows).to_csv(DATA + 'extreme_multi_lanthanide_census.csv', index=False)
    # list the most extreme cool stars (>= 3 lanthanides at >= +1.2 dex or above the model range)
    m = c['base'] & (nh >= 3)
    cols = {'sobject_id': d['sobject_id'][m], 'ra': d['ra'][m], 'dec': d['dec'][m], 'teff': teff[m],
            'logg': logg[m], 'fe_h': d['fe_h'][m], 'flag_sp': d['flag_sp'][m], 'n_lanthanides_ge_1p2': nh[m]}
    for e in ['ba', 'y'] + L:
        cols[e + '_fe'] = d[e + '_fe'][m]; cols['flag_' + e] = d['flag_' + e + '_fe'][m]
    pd.DataFrame(cols).round(4).to_csv(DATA + 'extreme_multi_lanthanide_stars.csv', index=False)
    return rows


def dr3_crosscheck(d, c):
    d3 = dict(np.load(CACHE + 'dr3.npz'))
    o3 = np.argsort(d3['sobject_id']); pos = np.clip(np.searchsorted(d3['sobject_id'][o3], d['sobject_id']), 0, len(o3) - 1)
    j = o3[pos]; common = d3['sobject_id'][j] == d['sobject_id']
    base = c['base'] & (d['flag_sp'] == 0)
    res = dict(n_dr3_rows=int(len(d3['sobject_id'])), n_common=int(common.sum()), n_common_base=int((base & common).sum()))
    for e in ['nd', 'la', 'ce', 'sm', 'eu', 'zr', 'ba', 'y']:
        x4 = d[e + '_fe'].astype(float); x3 = d3[e + '_fe'][j].astype(float); f3 = d3['flag_' + e + '_fe'][j]
        m = base & common & c['sec'][e] & (f3 == 0) & np.isfinite(x3)
        dif = (x4 - x3)[m]
        cand = s3.candidates(d, c, e, 0.75) & common
        res[e] = dict(n_both=int(m.sum()), median_dr4_minus_dr3=float(np.median(dif)) if m.sum() else None,
                      mad_sigma=float(1.4826 * np.median(np.abs(dif - np.median(dif)))) if m.sum() else None,
                      cand075_in_dr3=int(cand.sum()), cand075_dr3_flag0=int((cand & (f3 == 0)).sum()),
                      cand075_dr3_agree=int((cand & (f3 == 0) & (x3 >= x4 - 0.3)).sum()),
                      cand075_dr3_upper_limit=int((cand & (f3 == 1)).sum()))
    return res, (j, common, d3)


def main(d=None, c=None):
    if d is None:
        d = s3.load(); c = s3.compute(d)
    res = dict(apogee=apogee(d), gspspec=gspspec(d), przybylski=przybylski(), renson=renson(d))
    res['completeness_s_rich'] = completeness_s_rich(d, c)
    try:
        res['gspspec_rvs_check'] = gspspec_rvs_check()
    except Exception as ex:      # network trouble must not stop the pipeline; the check is then reported as not done
        res['gspspec_rvs_check'] = dict(error=str(ex))
    res['dr3'] = dr3_crosscheck(d, c)[0]
    res['extreme_census'] = extreme_census(d, c)
    rows = cp_onset(d)
    res['cp_onset'] = [dict(teff_lo=float(r[0]), n=r[2], k=r[3], n_sensitive=r[6], k_sensitive=r[7]) for r in rows]
    json.dump(res, open(DATA + 'other_surveys_summary.json', 'w'), indent=1, default=str)
    return res


if __name__ == '__main__':
    r = main()
    print(json.dumps({k: r[k] for k in ['apogee', 'gspspec', 'przybylski', 'renson']}, indent=1, default=str))