"""Spectrum-level vetting of catalogue candidates with a differential 'twin' test.

For every candidate (star, element) the public GALAH DR4 spectrum (observed flux, uncertainty and
the pipeline's best-fit model) is compared with K observed twin stars of the same Teff/logg/[Fe/H]/
vsini cell whose neutron-capture abundances are ordinary.  In each of the element's line windows
(GALAH DR3 line masks, Buder et al. 2021) we measure, after matching the local continuum,

    d_obs = sum (twin_obs - target_obs) dlambda      extra absorption actually observed
    d_mod = sum (twin_model - target_model) dlambda  extra absorption implied by the catalogue fit

and combine the lines with a matched filter,  rho = sum(d_obs d_mod / s^2) / sum(d_mod^2 / s^2).
rho ~ 1: the catalogue abundance is supported by the spectrum; rho ~ 0: the best-fit model contains
lines that are not in the data (optimiser failure).  In the weak-line limit the abundance supported
by the spectrum is  x_spec = log10(10^x_twin + rho (10^x_cat - 10^x_twin))  and the supported excess
Ex_spec = x_spec - prediction.  Null controls (ordinary stars) calibrate the per-line noise, positive
controls (s-process-rich dwarfs, whose enhancement is real) calibrate the efficiency of the test.
"""
import os, csv, sys, urllib.request
from concurrent.futures import ThreadPoolExecutor
import numpy as np
from common import *
import s03_candidates as s3
import spectra

MAIN = ['nd', 'la', 'ce', 'sm', 'eu', 'zr']   # probes; Ba and Y are the anchors of the natural model
T_VET = 0.75        # every catalogue candidate with observed excess >= T_VET is vetted
T_SUB = 0.5         # candidates with T_SUB <= excess < T_VET: a random subsample is vetted
N_SUB = 200         # size of that subsample per element
K_TWIN = 3
CLIP = 3.0           # one line is dropped if it deviates by more than CLIP sigma from the fit to the others
SEED = 20261012
DR3_CTRL = ['nd', 'la']   # elements for the DR3-based completeness control
DR3_MIN = 0.8        # DR3 abundance (on the DR4 scale) defining 'strongly enhanced according to DR3'
N_DR3 = 80           # stars per class and element
MASK_URL = 'https://raw.githubusercontent.com/svenbuder/GALAH_DR3/master/analysis/abundances/mode_DR3'


def line_masks():
    p = RAW + 'lit/mode_DR3.txt'
    os.makedirs(RAW + 'lit', exist_ok=True)
    if not os.path.exists(p):
        with urllib.request.urlopen(MASK_URL, timeout=60) as r:
            open(p, 'wb').write(r.read())
    masks = {}
    for line in open(p):
        q = line.split()
        if len(q) < 6 or line.startswith(';'):
            continue
        name = q[0]
        el = ''.join(ch for ch in name if ch.isalpha()).lower()
        if el in NCAP and any(ch.isdigit() for ch in name):
            try:
                masks.setdefault(el, []).append((float(q[1]), float(q[2]), float(q[3]), float(q[4]), float(q[5])))
            except ValueError:
                pass
    with open(DATA + 'line_masks_used.csv', 'w', newline='') as fh:
        w = csv.writer(fh); w.writerow(['element', 'line_A', 'mask_start', 'mask_end', 'seg_start', 'seg_end'])
        for el in NCAP:
            for m in masks.get(el, []):
                w.writerow([el] + list(m))
    return masks


NOISE_INFLATION = {}   # (element, line) -> per-line noise inflation, set by evaluate() on the main table
TWIN_MODE = {}   # target index -> 0 twins from the same cell, 1 from a neighbouring cell, 2 nearest-neighbour fallback


def cell_key(d, i):
    return (int(round(d['teff'][i] / 100.0)), int(round(d['logg'][i] / 0.2)),
            int(round(d['fe_h'][i] / 0.1)), int(np.digitize(d['vsini'][i], [6, 12, 20])))


def pick_twins(d, c, targets):
    """For each target index choose up to K_TWIN ordinary, high-S/N stars from the same cell."""
    teff = d['teff']; logg = d['logg']; feh = d['fe_h']; vs = d['vsini']
    normal = c['base'] & (d['flag_sp'] == 0) & (d['snr_px_ccd2'] > 60) & (d['snr_px_ccd1'] > 35)
    for e in MAIN:
        x = d[e + '_fe'].astype(float)
        normal &= ~(np.isfinite(x) & (x > 0.45) & ((d['flag_' + e + '_fe'] & 3) == 0))
    normal &= np.abs(d['ba_fe']) < 0.35
    keys = np.stack([np.round(teff / 100.0), np.round(logg / 0.2), np.round(feh / 0.1),
                     np.digitize(vs, [6, 12, 20])], axis=1)
    keys = np.where(np.isfinite(keys), keys, -999).astype(int)
    order = np.argsort(-d['snr_px_ccd2'])
    pool = {}
    for j in order:
        if not normal[j]:
            continue
        k = tuple(keys[j])
        lst = pool.setdefault(k, [])
        if len(lst) < K_TWIN:
            lst.append(j)
    tw = {}
    far = []
    for i in targets:
        k = tuple(keys[i])
        lst = [j for j in pool.get(k, []) if j != i]
        if len(lst) < 2:   # widen: neighbouring cells in Teff then [Fe/H]
            for dk in [(1, 0, 0, 0), (-1, 0, 0, 0), (0, 0, 1, 0), (0, 0, -1, 0), (0, 1, 0, 0), (0, -1, 0, 0)]:
                k2 = tuple(np.array(k) + np.array(dk))
                lst += [j for j in pool.get(k2, []) if j != i and j not in lst]
                if len(lst) >= K_TWIN:
                    break
        tw[i] = lst[:K_TWIN]
        TWIN_MODE[int(i)] = 0 if all(tuple(keys[j]) == k for j in tw[i]) else 1
        if len(tw[i]) < 2:
            far.append(i)
    if far:
        # last resort for sparsely populated corners of parameter space: the nearest ordinary stars
        # within 2.5 cell widths (250 K, 0.5 dex in logg, 0.25 dex in [Fe/H]), same vsini class
        nidx = np.where(normal & (d['snr_px_ccd2'] > 40))[0]
        P = np.stack([teff[nidx] / 100.0, logg[nidx] / 0.2, feh[nidx] / 0.1], axis=1)
        vc = np.digitize(vs[nidx], [6, 12, 20])
        for i in far:
            q = np.array([teff[i] / 100.0, logg[i] / 0.2, feh[i] / 0.1])
            dist = np.sqrt(((P - q) ** 2).sum(1)) + 10.0 * (vc != np.digitize(vs[i], [6, 12, 20]))
            dist[nidx == i] = np.inf
            o = np.argsort(dist)[:K_TWIN]
            o = o[dist[o] <= 2.5 * np.sqrt(3.0)]
            if len(o) >= 2:
                tw[i] = [int(nidx[j]) for j in o]; TWIN_MODE[int(i)] = 2
    return tw


def download_all(sids):
    sids = [int(s) for s in sids if not (os.path.exists(spectra.path(int(s))) and os.path.getsize(spectra.path(int(s))) > 10000)]
    if not sids:
        return
    print('downloading %d spectra ...' % len(sids)); sys.stdout.flush()
    with ThreadPoolExecutor(8) as ex:
        res = list(ex.map(spectra.fetch, sids))
    print('  failed:', sum(r is None for r in res))


def twin_test(S, T, masks_el):
    """S: target spectrum dict; T: list of twin dicts.  Returns per-element sums."""
    w = S['wave']; dl = np.gradient(w)
    tob = []; tmo = []; tun = []
    for t in T:
        tob.append(np.interp(w, t['wave'], t['sob'])); tmo.append(np.interp(w, t['wave'], t['smod']))
        tun.append(np.interp(w, t['wave'], t['uob']))
    tob = np.array(tob); tmo = np.array(tmo); tun = np.array(tun)
    ref_ob = np.nanmedian(tob, axis=0); ref_mo = np.nanmedian(tmo, axis=0)
    ref_un = 1.25 * np.sqrt(np.nanmean(tun ** 2, axis=0) / len(T))
    out = dict(n_lines=0, dew_obs=0.0, var=0.0, dew_mod=0.0, lines=[])
    for (w0, m0, m1, s0, s1) in masks_el:
        seg = (w >= s0) & (w <= s1); inm = (w >= m0) & (w <= m1)
        good = np.isfinite(S['sob']) & np.isfinite(S['uob']) & (S['uob'] > 0) & (S['uob'] < 0.5) & np.isfinite(ref_ob)
        if inm.sum() < 3 or (inm & good).sum() < 0.6 * inm.sum():
            continue
        side = seg & ~inm & good
        if side.sum() < 8:
            continue
        # local continuum match target -> reference (observed and model separately)
        c_ob = np.nanmedian(S['sob'][side] / ref_ob[side]); c_mo = np.nanmedian(S['smod'][side] / ref_mo[side])
        if not (np.isfinite(c_ob) and 0.8 < c_ob < 1.25):
            continue
        k = inm & good
        d_ob = np.sum((ref_ob[k] - S['sob'][k] / c_ob) * dl[k])
        d_mo = np.sum((ref_mo[k] - S['smod'][k] / c_mo) * dl[k])
        var = np.sum((S['uob'][k] ** 2 + ref_un[k] ** 2) * dl[k] ** 2)
        out['n_lines'] += 1; out['dew_obs'] += d_ob; out['dew_mod'] += d_mo; out['var'] += var
        out['lines'].append((w0, d_ob, np.sqrt(var), d_mo))
    return out


def run(d=None, c=None, verbose=True):
    if d is None:
        d = s3.load(); c = s3.compute(d)
    masks = line_masks()
    n = len(d['teff'])
    rng = np.random.default_rng(SEED)
    cand = {e: s3.candidates(d, c, e, T_VET) for e in MAIN}
    sub = {}
    for e in MAIN:
        band = np.where(s3.candidates(d, c, e, T_SUB) & ~cand[e])[0]
        sub[e] = np.sort(rng.choice(band, N_SUB, replace=False)) if len(band) > N_SUB else band
    allc = np.zeros(n, bool)
    for e in MAIN:
        allc |= cand[e]; allc[sub[e]] = True
    anyc = np.zeros(n, bool)
    for e in MAIN:
        anyc |= s3.candidates(d, c, e, T_SUB)
    # positive controls: secure s-process-rich dwarfs (high Ba and high Nd) -> real enhancements
    good = c['base'] & (d['flag_sp'] == 0)
    pos = good & c['sec']['nd'] & c['sec']['ba'] & (d['nd_fe'] >= 0.75) & (d['ba_fe'] >= 0.75) & ~anyc
    pos_idx = np.where(pos)[0]
    if len(pos_idx) > 80:
        pos_idx = np.sort(rng.choice(pos_idx, 80, replace=False))
    # null controls: ordinary stars with a secure, ordinary Nd detection
    nul = good & c['sec']['nd'] & c['sec']['ba'] & (np.abs(c['ex']['nd']) < 0.2) & (d['snr_px_ccd2'] > 30) & ~anyc
    nul_idx = np.sort(rng.choice(np.where(nul)[0], 150, replace=False))
    # completeness controls from the independent GALAH DR3 pipeline (same spectra, SME line-by-line
    # analysis): dwarfs that DR3 reports as strongly enhanced (flag 0) and that are detectable in DR4.
    #   dr3_only: DR4 does NOT report a detection  -> is real absorption being missed by DR4?
    #   dr3_both: DR4 agrees within 0.3 dex        -> reference response of a real enhancement
    d3 = dict(np.load(CACHE + 'dr3.npz'))
    o3 = np.argsort(d3['sobject_id']); p3 = np.clip(np.searchsorted(d3['sobject_id'][o3], d['sobject_id']), 0, len(o3) - 1)
    j3 = o3[p3]; common = (d3['sobject_id'][j3] == d['sobject_id']) & (d3['flag_sp'][j3] == 0)
    dr3 = {}; dr3_off = {}
    for e in DR3_CTRL:
        x4 = d[e + '_fe'].astype(float); x3 = d3[e + '_fe'][j3].astype(float); f3 = d3['flag_' + e + '_fe'][j3]
        m = good & common & c['sec'][e] & (f3 == 0) & np.isfinite(x3)
        off = float(np.median((x4 - x3)[m])); dr3_off[e] = off
        t = good & common & (f3 == 0) & np.isfinite(x3) & (x3 + off >= DR3_MIN) & (d['x90_' + e] <= x3 + off - 0.25)
        only = np.where(t & ~c['sec'][e])[0]; both = np.where(t & c['sec'][e] & (x4 >= x3 + off - 0.3))[0]
        n_only, n_both = len(only), len(both)
        if len(only) > N_DR3:
            only = np.sort(rng.choice(only, N_DR3, replace=False))
        if len(both) > N_DR3:
            both = np.sort(rng.choice(both, N_DR3, replace=False))
        dr3[e] = (only, both, n_only, n_both, int(t.sum()), {int(i): float(x3[i] + off) for i in np.concatenate([only, both])})
    targets = np.unique(np.concatenate([np.where(allc)[0], pos_idx, nul_idx] + [dr3[e][k] for e in DR3_CTRL for k in (0, 1)]))
    tw = pick_twins(d, c, targets)
    need = set(d['sobject_id'][targets].tolist())
    for i in targets:
        need |= set(d['sobject_id'][tw[i]].tolist())
    print('targets to vet: %d unique candidate stars; positive controls %d; null controls %d; spectra needed %d'
          % (allc.sum(), len(pos_idx), len(nul_idx), len(need)))
    download_all(sorted(need))
    from collections import OrderedDict
    cache = OrderedDict()      # small LRU cache: spectra are read on demand, never all held in memory

    def get(sid):
        sid = int(sid)
        if sid in cache:
            cache.move_to_end(sid)
            return cache[sid]
        try:
            sp = spectra.read(sid)
            if sp is not None:
                sp = {k: (v.astype('f4') if k != 'wave' and v.dtype != bool else v) for k, v in sp.items()}
        except Exception:
            sp = None
        cache[sid] = sp
        if len(cache) > 300:
            cache.popitem(last=False)
        return sp

    def do(i, e, kind):
        S = get(d['sobject_id'][i])
        T = [get(d['sobject_id'][j]) for j in tw[i]]
        ok = [t is not None for t in T]
        xt = np.array([d[e + '_fe'][j] for j, o in zip(tw[i], ok) if o], float)
        T = [t for t in T if t is not None]
        x_tw = float(np.nanmedian(xt)) if np.isfinite(xt).any() else 0.0
        row = dict(kind=kind, element=e, sobject_id=int(d['sobject_id'][i]), idx=int(i), ex_cat=float(c['ex'][e][i]),
                   x_cat=float(d[e + '_fe'][i]), pred=float(c['pred'][e][i]), x_twin=x_tw, n_twins=len(T),
                   twin_mode=TWIN_MODE.get(int(i), -1),
                   n_lines=0, lines='', status='ok')
        if S is None or len(T) < 2:
            row['status'] = 'no_spectrum' if S is None else 'no_twins'
            return row
        r = twin_test(S, T, masks.get(e, []))
        if r['n_lines'] == 0:
            row['status'] = 'no_lines'; return row
        row.update(n_lines=r['n_lines'],
                   lines=';'.join('%.1f:%.2f:%.2f:%.2f' % (l[0], 1e3 * l[1], 1e3 * l[2], 1e3 * l[3]) for l in r['lines']))
        return row

    rows = []
    todo = [(cell_key(d, i), int(i), e, 'candidate') for e in MAIN for i in np.where(cand[e])[0]]
    todo += [(cell_key(d, i), int(i), e, 'candidate_sub') for e in MAIN for i in sub[e]]
    todo += [(cell_key(d, i), int(i), e, 'positive_control') for i in pos_idx for e in MAIN
             if c['sec'][e][i] and d[e + '_fe'][i] >= 0.5]
    todo += [(cell_key(d, i), int(i), e, 'null_control') for i in nul_idx for e in MAIN]
    for e in DR3_CTRL:
        todo += [(cell_key(d, i), int(i), e, 'dr3_only') for i in dr3[e][0]]
        todo += [(cell_key(d, i), int(i), e, 'dr3_both') for i in dr3[e][1]]
    for _, i, e, kind in sorted(todo):
        r = do(i, e, kind)
        r['x_dr3'] = dr3[e][5].get(int(i), np.nan) if e in dr3 else np.nan
        rows.append(r)
    cols = ['kind', 'element', 'sobject_id', 'idx', 'x_cat', 'ex_cat', 'pred', 'x_twin', 'x_dr3', 'n_twins', 'twin_mode',
            'n_lines', 'lines', 'status']
    with open(DATA + 'vetting_twin_test_raw.csv', 'w', newline='') as fh:
        w = csv.DictWriter(fh, fieldnames=cols); w.writeheader()
        for r in rows:
            w.writerow({k: (f'{r[k]:.4f}' if isinstance(r[k], float) else r[k]) for k in cols})
    with open(DATA + 'vetting_subsample_sizes.csv', 'w', newline='') as fh:
        w = csv.writer(fh); w.writerow(['element', 'n_candidates_T>=%.2f' % T_VET, 'n_band_%.2f<=T<%.2f' % (T_SUB, T_VET), 'n_band_vetted'])
        for e in MAIN:
            w.writerow([e, int(cand[e].sum()), int((s3.candidates(d, c, e, T_SUB) & ~cand[e]).sum()), len(sub[e])])
    with open(DATA + 'dr3_control_counts.csv', 'w', newline='') as fh:
        w = csv.writer(fh)
        w.writerow(['element', 'dr4_minus_dr3_offset', 'n_dr3_enhanced_and_dr4_detectable', 'n_dr4_no_detection',
                    'n_dr4_agrees', 'n_no_detection_vetted', 'n_agrees_vetted'])
        for e in DR3_CTRL:
            w.writerow([e, '%.3f' % dr3_off[e], dr3[e][4], dr3[e][2], dr3[e][3], len(dr3[e][0]), len(dr3[e][1])])
    return rows


RAW_COLS = ['kind', 'element', 'sobject_id', 'idx', 'x_cat', 'ex_cat', 'pred', 'x_twin', 'x_dr3', 'n_twins', 'twin_mode',
            'n_lines', 'lines', 'status']


def shifted_masks(masks, shift):
    """Control windows: every line mask moved by `shift` A inside its own segment (no line of the
    element there).  Used to test whether extra absorption is specific to the element's lines."""
    return {el: [(m[0] + shift, m[1] + shift, m[2] + shift, m[3], m[4]) for m in masks[el]] for el in masks}


def measure_compact(d, c, todo, out_csv, masks=None, threads=8, verbose=True, cache_masks=None):
    """Twin test for a list of (target index, element, kind) on spectra that are NOT kept on disk
    in full: each spectrum is downloaded into memory and only the pixels of the line segments are
    cached (spectra.read_compact).  Same measurement as run(); writes a raw table of the same
    format to data/<out_csv> and returns the rows.  Used by the revision steps."""
    cache_masks = line_masks() if cache_masks is None else cache_masks     # defines which pixels the compact cache keeps
    masks = cache_masks if masks is None else masks
    targets = np.unique(np.array([t[0] for t in todo], dtype=int))
    tw = pick_twins(d, c, targets)
    need = set(int(d['sobject_id'][i]) for i in targets)
    for i in targets:
        need |= set(int(d['sobject_id'][j]) for j in tw[i])
    miss = [s_ for s_ in sorted(need) if not (os.path.exists(spectra.path(s_)) or os.path.exists(spectra.compact_path(s_)))]
    if verbose:
        print('  %d targets, %d spectra needed, %d to download (compact cache)' % (len(targets), len(need), len(miss))); sys.stdout.flush()
    if miss:
        with ThreadPoolExecutor(threads) as ex:
            res = list(ex.map(lambda q: spectra.fetch_compact(q, cache_masks), miss))
        if verbose:
            print('  failed downloads:', sum(r is None for r in res)); sys.stdout.flush()
    from collections import OrderedDict
    cache = OrderedDict()

    def get(sid):
        sid = int(sid)
        if sid in cache:
            cache.move_to_end(sid)
            return cache[sid]
        try:
            sp = spectra.read_compact(sid, cache_masks)
        except Exception:
            sp = None
        cache[sid] = sp
        if len(cache) > 300:
            cache.popitem(last=False)
        return sp

    rows = []
    for i, e, kind in sorted(todo, key=lambda t: (cell_key(d, t[0]), int(t[0]), t[1], t[2])):
        i = int(i)
        S = get(d['sobject_id'][i])
        T = [get(d['sobject_id'][j]) for j in tw[i]]
        ok = [t is not None for t in T]
        xt = np.array([d[e + '_fe'][j] for j, o in zip(tw[i], ok) if o], float)
        T = [t for t in T if t is not None]
        x_tw = float(np.nanmedian(xt)) if np.isfinite(xt).any() else 0.0
        row = dict(kind=kind, element=e, sobject_id=int(d['sobject_id'][i]), idx=i, ex_cat=float(c['ex'][e][i]),
                   x_cat=float(d[e + '_fe'][i]), pred=float(c['pred'][e][i]), x_twin=x_tw, n_twins=len(T),
                   twin_mode=TWIN_MODE.get(i, -1), n_lines=0, lines='', status='ok', x_dr3=np.nan)
        if S is None or len(T) < 2:
            row['status'] = 'no_spectrum' if S is None else 'no_twins'
        else:
            r = twin_test(S, T, masks.get(e, []))
            if r['n_lines'] == 0:
                row['status'] = 'no_lines'
            else:
                row.update(n_lines=r['n_lines'],
                           lines=';'.join('%.1f:%.2f:%.2f:%.2f' % (l[0], 1e3 * l[1], 1e3 * l[2], 1e3 * l[3]) for l in r['lines']))
        rows.append(row)
    with open(DATA + out_csv, 'w', newline='') as fh:
        w = csv.DictWriter(fh, fieldnames=RAW_COLS); w.writeheader()
        for r in rows:
            w.writerow({k: (f'{r[k]:.4f}' if isinstance(r[k], float) else r[k]) for k in RAW_COLS})
    return rows


def parse_lines(s):
    if not isinstance(s, str) or not s:
        return np.zeros((0, 4))
    return np.array([[float(q) for q in t.split(':')] for t in s.split(';')])


def x_spec(rho, x_cat, x_twin):
    """Abundance supported by the spectrum (weak-line limit) for a matched-filter amplitude rho."""
    return np.log10(np.maximum(10 ** x_twin + rho * (10 ** x_cat - 10 ** x_twin), 10 ** -1.5))


EDGE_TOL = 0.03      # dex: tolerance on 'at the upper edge of the model grid'


def x_spec_sat(rho, x_cat, x_twin):
    """As x_spec, but for rho > 1 the abundance grows as rho^2 (square-root part of the curve of growth)."""
    lin = x_spec(rho, x_cat, x_twin)
    with np.errstate(invalid='ignore', divide='ignore'):
        sat = x_cat + 2.0 * np.log10(np.maximum(rho, 1.0))
    return np.where(np.isfinite(rho) & (rho > 1.0), np.maximum(lin, sat), lin)


def evaluate(raw_file='vetting_twin_test_raw.csv', infl=None):
    """Read the raw twin-test table, calibrate the per-line noise on the null controls and compute
    the matched-filter amplitude, its leave-one-line-out variants and the supported excess."""
    import pandas as pd
    v = pd.read_csv(DATA + raw_file, keep_default_na=False, na_values=['nan'])
    L = [parse_lines(s) for s in v.lines]
    # noise inflation per (element, line) from the null controls: robust sigma of d_obs / s
    # (a table without null controls, e.g. the flagged-star extension, is given the inflation
    # factors of the main table through `infl`)
    given = infl is not None
    infl = dict(infl) if given else {}; nullstat = []
    isnull = (v.kind == 'null_control').values
    for e in ([] if given else MAIN):
        zs = {}
        for k in np.where(isnull & (v.element == e).values)[0]:
            for w0, dob, sg, dmo in L[k]:
                zs.setdefault(w0, []).append(dob / sg)
        for w0, z in sorted(zs.items()):
            z = np.array(z); mad = 1.4826 * np.median(np.abs(z - np.median(z)))
            infl[(e, w0)] = max(1.0, float(mad))
            nullstat.append(dict(element=e, line=w0, n=len(z), z_median=float(np.median(z)), z_robust_sigma=float(mad),
                                 z_std=float(z.std()), frac_abs_z_gt3=float((np.abs(z) > 3 * max(1.0, mad)).mean()),
                                 inflation_used=infl[(e, w0)]))
    out = {k: np.full(len(v), np.nan) for k in ['rho', 'rho_err', 'snr_mod', 'rho_dp', 'rho_dp_err', 'rho_dn', 'rho_dn_err',
                                                 'n_used', 'z_null', 'dew_obs', 'dew_sig', 'dew_mod', 'rho_all', 'rho_all_err',
                                                 'clipped_line', 'clip_pull', 'chi2_red']}
    for k in range(len(v)):
        a = L[k]
        if len(a) == 0:
            continue
        e = v.element.values[k]
        sg = a[:, 2] * np.array([infl.get((e, w0), 1.0) for w0 in a[:, 0]])
        out['z_null'][k] = a[:, 1].sum() / np.sqrt((sg ** 2).sum())
        out['dew_obs'][k] = a[:, 1].sum(); out['dew_sig'][k] = np.sqrt((sg ** 2).sum()); out['dew_mod'][k] = a[:, 3].sum()
        u = a[:, 3] > 0
        if u.sum() == 0:
            continue
        dob = a[u, 1]; dmo = a[u, 3]; sg = sg[u]
        den = np.sum(dmo ** 2 / sg ** 2)
        rho_all = np.sum(dob * dmo / sg ** 2) / den; err_all = den ** -0.5
        out['rho_all'][k] = rho_all; out['rho_all_err'][k] = err_all
        out['rho'][k] = rho_all; out['rho_err'][k] = err_all
        out['snr_mod'][k] = np.sqrt(den); out['n_used'][k] = u.sum()
        if u.sum() >= 3:
            wl = a[u, 0]
            contrib = dob * dmo / sg ** 2
            for nm, j in (('dp', int(np.argmax(contrib))), ('dn', int(np.argmin(contrib)))):
                kp = np.arange(len(dob)) != j
                den2 = np.sum(dmo[kp] ** 2 / sg[kp] ** 2)
                out['rho_' + nm][k] = np.sum(dob[kp] * dmo[kp] / sg[kp] ** 2) / den2; out['rho_' + nm + '_err'][k] = den2 ** -0.5
            # one-line outlier clipping: pull of each line against the fit to the OTHER lines
            pulls = np.zeros(len(dob)); fits = []
            for j in range(len(dob)):
                kp = np.arange(len(dob)) != j
                den2 = np.sum(dmo[kp] ** 2 / sg[kp] ** 2)
                r2 = np.sum(dob[kp] * dmo[kp] / sg[kp] ** 2) / den2; e2 = den2 ** -0.5
                pulls[j] = (dob[j] - r2 * dmo[j]) / np.sqrt(sg[j] ** 2 + (dmo[j] * e2) ** 2)
                fits.append((r2, e2))
            j = int(np.argmax(np.abs(pulls)))
            out['clip_pull'][k] = pulls[j]
            keep = np.ones(len(dob), bool)
            if abs(pulls[j]) > CLIP:
                out['rho'][k], out['rho_err'][k] = fits[j]; out['clipped_line'][k] = wl[j]; keep[j] = False
            if keep.sum() >= 2:
                out['chi2_red'][k] = np.sum(((dob[keep] - out['rho'][k] * dmo[keep]) / sg[keep]) ** 2) / (keep.sum() - 1)
    for k_, a_ in out.items():
        v[k_] = a_
    if not given:
        global NOISE_INFLATION
        NOISE_INFLATION = dict(infl)
    xc = v.x_cat.values; xt = v.x_twin.values; pr = v.pred.values
    v['ex_spec'] = x_spec(v.rho.values, xc, xt) - pr
    v['ex_spec_hi'] = x_spec(v.rho.values + 2 * v.rho_err.values, xc, xt) - pr
    v['ex_spec_lo'] = x_spec(v.rho.values - 2 * v.rho_err.values, xc, xt) - pr
    v['ex_spec_dp_hi'] = x_spec(v.rho_dp.values + 2 * v.rho_dp_err.values, xc, xt) - pr   # best line dropped
    v['ex_spec_dn_hi'] = x_spec(v.rho_dn.values + 2 * v.rho_dn_err.values, xc, xt) - pr   # worst line dropped
    v['ex_spec_all_hi'] = x_spec(v.rho_all.values + 2 * v.rho_all_err.values, xc, xt) - pr     # no clipping
    v['xs_hi'] = x_spec(v.rho.values + 2 * v.rho_err.values, xc, xt)
    v['xs_dp_hi'] = x_spec(v.rho_dp.values + 2 * v.rho_dp_err.values, xc, xt)
    v['xs_dn_hi'] = x_spec(v.rho_dn.values + 2 * v.rho_dn_err.values, xc, xt)
    # --- revision (referee point 6): the weak-line extrapolation is not valid beyond the model grid.
    # xs = abundance the spectrum supports (central value); beyond_grid = that abundance lies at or
    # above the upper edge of the DR4 label range, so the spectrum cannot bound it from above.
    edge = np.array([DR4_UPPER[e] for e in v.element.values])
    v['xs'] = x_spec(v.rho.values, xc, xt)
    v['beyond_grid'] = (np.isfinite(v.rho.values) & (v['xs'].values >= edge - EDGE_TOL)).astype(int)
    # saturation allowance (robustness variant): where the spectrum shows MORE absorption than the
    # catalogue model (rho > 1) the abundance is allowed to grow as on the square-root part of the
    # curve of growth, Delta log N = 2 Delta log W, instead of linearly.
    v['ex_spec_sat_hi'] = x_spec_sat(v.rho.values + 2 * v.rho_err.values, xc, xt) - pr
    # is the spectrum consistent with the star's own catalogue value?  (rho = 1 within 2 sigma)
    v['inconsistent_with_catalogue'] = (np.isfinite(v.rho.values) & (v.rho.values + 2 * v.rho_err.values < 1.0)).astype(int)
    return v, pd.DataFrame(nullstat)


def classify(row, T, margin=s3.MARGIN, rule='clip'):
    """Verdict for one (star, element) as a member of the population 'true excess >= T + margin'
    (the salt level that threshold T constrains).  rho is the matched-filter amplitude over the
    element's lines after dropping at most one line that deviates by more than CLIP sigma from the
    fit to the others (a blend / cosmic ray / bad window in one line is not an abundance, and one
    bad window must not hide a real signal either).
    refuted   : the spectrum excludes at 2 sigma a supported excess >= T + margin, and the abundance
                the spectrum supports lies inside the model grid (see 'beyond_grid').
    supported : not refuted and rho >= 3 sigma (extra absorption is really there).
    ambiguous : neither (spectrum too noisy to decide).   unverifiable: no usable spectrum/line.
    Positive controls measure how often a real enhancement is wrongly refuted by this rule."""
    if row['status'] != 'ok' or not np.isfinite(row['rho']):
        return 'unverifiable'
    lim = T + margin
    if rule in ('clip', 'firstpass'):
        refuted = row['ex_spec_hi'] < lim
    elif rule == 'sat':
        refuted = row['ex_spec_sat_hi'] < lim
    else:
        # alternative, more conservative rule used as a robustness variant: no clipping; the 2-sigma
        # bound must exclude the excess both with all lines and with the most discrepant-low line removed
        refuted = (row['ex_spec_all_hi'] < lim) and (row['n_used'] < 3 or row['ex_spec_dn_hi'] < lim)
    # revision: a star whose spectrum supports an abundance at or above the edge of the model grid
    # (flag 8, or rho > 1 near the edge) cannot be refuted by linear extrapolation: its lines are
    # stronger than any model the pipeline has, so the true abundance is not bounded from above.
    # (rule 'firstpass' reproduces the first-pass classification without this protection.)
    if rule != 'firstpass' and row['beyond_grid'] == 1:
        refuted = False
    if refuted:
        return 'refuted'
    if row['rho'] >= 3 * row['rho_err']:
        return 'supported'
    return 'ambiguous'


if __name__ == '__main__':
    rows = run()
    import collections
    print(collections.Counter((r['kind'], r['status']) for r in rows))
    v, ns = evaluate()
    print(ns.to_string(index=False))