"""Revision step (referee major point 4): what the flag_sp == 0 cut hides.

The search sample requires flag_sp == 0, which removes 18% of the FGK dwarfs/subgiants.  This step
  (1) counts single-element catalogue candidates (the actual search target, same definition as in
      s03 but with flag_sp != 0) among the flagged stars, per element and per flag bit;
  (2) runs the same twin test on the spectra of every flagged Nd and La candidate (twins are clean
      stars in the candidate's catalogue Teff/logg/[Fe/H]/vsini cell);
  (3) lists what is not refuted, with the flag bits, a field-clustering count (many candidates on
      one plate in one night point to a reduction problem) and binarity indicators.
The effective sample and the limits are NOT extended to flagged stars: stellar parameters of
flagged stars are less reliable, so twins may be poorly matched and the detection model of s02 does
not apply.  This is a census of what could be hiding there, not a limit.
"""
import csv, json
import numpy as np
import pandas as pd
from common import *
import s03_candidates as s3
import s04_vetting as s4

ELEMS_VET = ['nd', 'la']
FLAG_BITS = {1: 'emission in Halpha/Hbeta', 2: 'not all 4 CCDs', 4: 'SB1', 8: 'SB2', 16: 'chi2 unusual',
             32: 'vsini warning', 64: 'vmic warning', 128: 'double-line splitting', 256: 'Teff warning',
             512: 'logg warning', 1024: '[Fe/H] warning', 2048: 'low S/N', 4096: 'not converged',
             8192: 'extrapolating model'}
PARAM_BITS = 8 | 16 | 32 | 128 | 256 | 512 | 1024 | 4096 | 8192    # bits that make the stellar parameters (and twins) suspect


def flagged_candidates(d, c, e, T):
    ok = c['base'] & (d['flag_sp'] != 0) & c['sec'][e] & c['anchor_ok'][e]
    ex = c['ex'][e]
    return ok & ((ex >= T) | (s3.at_edge(d, e) & (ex >= min(T, 0.5))))


def main(verbose=True):
    d = s3.load(); c = s3.compute(d)
    fl = d['flag_sp'].astype(int); base = c['base']
    clean = base & (fl == 0); flagged = base & (fl != 0)
    T = s4.T_VET
    # ---- (1) census
    rows = []
    for e in s4.MAIN:
        cf = flagged_candidates(d, c, e, T); cc = s3.candidates(d, c, e, T)
        f8 = (d['flag_' + e + '_fe'] & 8) > 0
        r = dict(element=e, n_clean_stars=int(clean.sum()), n_flagged_stars=int(flagged.sum()),
                 n_candidates_clean=int(cc.sum()), n_candidates_flagged=int(cf.sum()),
                 rate_clean=float(cc.sum() / clean.sum()), rate_flagged=float(cf.sum() / flagged.sum()),
                 rate_ratio=float((cf.sum() / flagged.sum()) / max(cc.sum() / clean.sum(), 1e-12)),
                 n_above_model_range_clean=int((clean & f8).sum()), n_above_model_range_flagged=int((flagged & f8).sum()))
        for b, name in FLAG_BITS.items():
            n = int((cf & ((fl & b) > 0)).sum())
            if (flagged & ((fl & b) > 0)).sum() > 0:
                r['cand_with_bit_%d' % b] = n
        r['cand_chi2_only_flag16'] = int((cf & (fl == 16)).sum())
        rows.append(r)
    census = pd.DataFrame(rows)
    census.to_csv(DATA + 'flagged_candidate_census.csv', index=False, float_format='%.5g')
    bits = [dict(bit=b, meaning=name, n_flagged_stars=int((flagged & ((fl & b) > 0)).sum())) for b, name in FLAG_BITS.items()
            if (flagged & ((fl & b) > 0)).sum() > 0]
    # ---- (2) twin test on every flagged Nd / La candidate
    s4.evaluate()                      # sets s4.NOISE_INFLATION from the null controls of the main table
    todo = [(int(i), e, 'flagged_candidate') for e in ELEMS_VET for i in np.where(flagged_candidates(d, c, e, T))[0]]
    s4.measure_compact(d, c, todo, 'flagged_twin_test_raw.csv', verbose=verbose)
    v, _ = s4.evaluate('flagged_twin_test_raw.csv', infl=s4.NOISE_INFLATION)
    v['verdict'] = [s4.classify(r, T) for _, r in v.iterrows()]
    idx = v.idx.values.astype(int)
    v['flag_sp'] = fl[idx]
    for col in ['ra', 'dec', 'teff', 'logg', 'fe_h', 'vsini', 'snr_px_ccd1', 'snr_px_ccd2', 'chi2_sp', 'ruwe', 'rv_comp_2', 'ba_fe', 'y_fe']:
        v[col] = np.asarray(d[col])[idx].astype(float)
    v['significance'] = v.rho / v.rho_err
    v['parameters_suspect'] = ((v.flag_sp.values & PARAM_BITS) > 0).astype(int)
    # field clustering: candidates (any element, clean or flagged, T >= 0.75) on the same plate/night
    field = d['sobject_id'] // 1000
    anyc = np.zeros(len(fl), bool)
    for e in s4.MAIN:
        anyc |= flagged_candidates(d, c, e, T) | s3.candidates(d, c, e, T)
    uf, cnt = np.unique(field[anyc], return_counts=True); fc = dict(zip(uf.tolist(), cnt.tolist()))
    nf = dict(zip(*[a.tolist() for a in np.unique(field[base], return_counts=True)]))
    v['field'] = field[idx]
    v['n_candidates_in_field'] = [fc.get(int(f), 0) for f in v.field]
    v['n_dwarfs_in_field'] = [nf.get(int(f), 0) for f in v.field]
    cols = ['element', 'sobject_id', 'ra', 'dec', 'teff', 'logg', 'fe_h', 'vsini', 'snr_px_ccd1', 'snr_px_ccd2', 'flag_sp',
            'parameters_suspect', 'chi2_sp', 'ruwe', 'rv_comp_2', 'ba_fe', 'y_fe', 'x_cat', 'pred', 'ex_cat', 'x_twin',
            'n_twins', 'twin_mode', 'status', 'n_used', 'rho', 'rho_err', 'significance', 'snr_mod', 'chi2_red',
            'clipped_line', 'ex_spec', 'ex_spec_hi', 'beyond_grid', 'inconsistent_with_catalogue', 'verdict',
            'field', 'n_candidates_in_field', 'n_dwarfs_in_field']
    v[cols].to_csv(DATA + 'flagged_candidates_vetting.csv', index=False, float_format='%.4f')
    # ---- (2b) is the extra absorption specific to the element?  A single-element salt strengthens
    # the lines of ONE element; a composite (SB2) or mis-parameterised spectrum, a barium star or a
    # chemically peculiar (Ap) star strengthens many.  For every 'supported' flagged candidate the
    # same twin statistic (extra absorption relative to twins, in sigma) is measured in the windows
    # of nine elements; null-control stars of the main sample calibrate the three elements (Y, Ba,
    # Sr) that the main table does not cover.
    ALL = ['nd', 'la', 'ce', 'sm', 'eu', 'zr', 'y', 'ba', 'sr']
    sup = v[v.verdict == 'supported']
    main_raw = pd.read_csv(DATA + 'vetting_twin_test_raw.csv', keep_default_na=False, na_values=['nan'])
    null_idx = sorted(set(main_raw[main_raw.kind == 'null_control'].idx.astype(int)))
    todo2 = [(int(i), e2, 'flagged_supported') for i in sorted(set(sup.idx.astype(int))) for e2 in ALL]
    todo2 += [(int(i), e2, 'null_control') for i in null_idx for e2 in ('y', 'ba', 'sr')]
    rows2 = s4.measure_compact(d, c, todo2, 'flagged_multi_element_raw.csv', verbose=verbose)
    infl = dict(s4.NOISE_INFLATION)
    zs = {}
    for r in rows2:
        if r['kind'] == 'null_control':
            for w0, dob, sg, dmo in s4.parse_lines(r['lines']):
                zs.setdefault((r['element'], w0), []).append(dob / sg)
    for k, z in zs.items():
        z = np.array(z); infl[k] = max(1.0, float(1.4826 * np.median(np.abs(z - np.median(z)))))
    zel = {}
    for r in rows2:
        if r['kind'] != 'flagged_supported':
            continue
        a = s4.parse_lines(r['lines'])
        if len(a) == 0:
            continue
        sg = a[:, 2] * np.array([infl.get((r['element'], w0), 1.5) for w0 in a[:, 0]])
        zel[(r['idx'], r['element'])] = float(a[:, 1].sum() / np.sqrt((sg ** 2).sum()))
    multi = []
    for _, r in sup.iterrows():
        o = dict(element=r.element, sobject_id=int(r.sobject_id), ra=r.ra, dec=r.dec, teff=r.teff, logg=r.logg, fe_h=r.fe_h,
                 vsini=r.vsini, snr_px_ccd2=r.snr_px_ccd2, flag_sp=int(r.flag_sp), ruwe=r.ruwe, rv_comp_2=r.rv_comp_2,
                 significance=r.significance, rho=r.rho, ex_spec=r.ex_spec, chi2_red=r.chi2_red,
                 n_candidates_in_field=int(r.n_candidates_in_field))
        others = [e2 for e2 in ALL if e2 != r.element]
        for e2 in ALL:
            o['z_' + e2] = zel.get((int(r.idx), e2), np.nan)
        strong = [e2 for e2 in others if np.isfinite(o['z_' + e2]) and o['z_' + e2] >= 3]
        o['n_other_elements_ge3sigma'] = len(strong); o['other_elements_ge3sigma'] = '+'.join(strong)
        o['n_other_elements_measured'] = int(sum(np.isfinite(o['z_' + e2]) for e2 in others))
        # Ap-like signature: Eu and Sr both strong together with a lanthanide (SrCrEu-type pattern)
        o['eu_and_sr_strong'] = int(o.get('z_eu', np.nan) >= 5 and o.get('z_sr', np.nan) >= 5)
        # strong lines (Ba II, Y II) clearly SHALLOWER than in the twins: line doubling or extra
        # broadening, i.e. the twins do not match and the comparison is not valid
        weak_anchor = (o.get('z_ba', np.nan) <= -4) or (o.get('z_y', np.nan) <= -4)
        if o['eu_and_sr_strong']:
            o['class'] = 'A: Ap-like line forest (Sr, Eu and several lanthanides strong)'
        elif weak_anchor:
            o['class'] = 'B: twins do not match (Ba/Y lines much shallower than in twins: composite or broadened spectrum)'
        elif len(strong) >= 2:
            o['class'] = 'C: several elements strong together (not a single-element excess)'
        else:
            o['class'] = 'D: element-specific and twins match (unresolved)'
        multi.append(o)
    multi = pd.DataFrame(multi)
    multi.to_csv(DATA + 'flagged_supported_multi_element.csv', index=False, float_format='%.3f')
    # ---- (3) summary
    summ = dict(n_flagged_dwarfs=int(flagged.sum()), n_clean_dwarfs=int(clean.sum()), flag_bits=bits,
                census=census.to_dict('records'), vetting={})
    for e in ELEMS_VET:
        g = v[v.element == e]
        o = dict(n_candidates=int(len(g)), verdicts={k: int(n) for k, n in g.verdict.value_counts().items()},
                 n_inconsistent_with_catalogue=int(g.inconsistent_with_catalogue.sum()))
        sub = {}
        for lab, m in [('chi2_only_flag16', g.flag_sp == 16), ('chi2_bit_set', (g.flag_sp & 16) > 0),
                       ('low_snr_bit_set', (g.flag_sp & 2048) > 0), ('missing_ccd_bit_set', (g.flag_sp & 2) > 0),
                       ('sb1_or_sb2_bit_set', (g.flag_sp & 12) > 0), ('vsini_or_vmic_warning', (g.flag_sp & 96) > 0),
                       ('parameters_not_suspect', g.parameters_suspect == 0)]:
            gg = g[m]
            sub[lab] = dict(n=int(len(gg)), **{k: int(n) for k, n in gg.verdict.value_counts().items()})
        o['by_flag'] = sub
        sup = g[g.verdict == 'supported']
        o['supported'] = [dict(sobject_id=int(r.sobject_id), flag_sp=int(r.flag_sp), significance=float(r.significance),
                               rho=float(r.rho), ex_spec=float(r.ex_spec), ba_fe=float(r.ba_fe), ruwe=float(r.ruwe),
                               n_candidates_in_field=int(r.n_candidates_in_field)) for _, r in sup.iterrows()]
        mm = multi[multi.element == e] if len(multi) else multi
        for lab in 'ABCD':
            o['supported_class_' + lab] = int(mm['class'].str.startswith(lab).sum()) if len(mm) else 0
        o['class_D_stars'] = [int(x) for x in mm[mm['class'].str.startswith('D')].sobject_id] if len(mm) else []
        o['class_A_stars'] = [int(x) for x in mm[mm['class'].str.startswith('A')].sobject_id] if len(mm) else []
        o['supported_with_binary_flag_or_ruwe'] = int((((mm.flag_sp & (4 | 8 | 128)) > 0) | (mm.ruwe >= 1.4) | np.isfinite(mm.rv_comp_2)).sum()) if len(mm) else 0
        summ['vetting'][e] = o
    json.dump(summ, open(DATA + 'flagged_summary.json', 'w'), indent=1)
    if verbose:
        pd.set_option('display.width', 250); pd.set_option('display.max_columns', 50); pd.set_option('display.max_rows', 300)
        print(census.to_string(index=False))
        for e in ELEMS_VET:
            print(e, json.dumps({k: w for k, w in summ['vetting'][e].items() if k != 'supported'}, indent=0))
        print(multi[['element', 'sobject_id', 'teff', 'vsini', 'flag_sp', 'ruwe', 'significance'] + ['z_' + e2 for e2 in ALL] +
                    ['n_other_elements_ge3sigma', 'class']].sort_values(['class', 'significance']).to_string(index=False))
        print(v[v.verdict == 'supported'][['element', 'sobject_id', 'teff', 'logg', 'fe_h', 'vsini', 'snr_px_ccd2', 'flag_sp', 'ruwe',
                                           'ba_fe', 'x_cat', 'pred', 'rho', 'rho_err', 'significance', 'ex_spec', 'ex_spec_hi',
                                           'chi2_red', 'n_used', 'n_candidates_in_field', 'n_dwarfs_in_field']].to_string(index=False))
    return summ, v


if __name__ == '__main__':
    main()