"""Post-registration correction: the SAP test must see the event with the same sign.

The registered vetting (code/sflare.py, side_signals) compares the size of the uncorrected (SAP)
flux change with the PDCSAP amplitude but takes its absolute value, so an SAP decrease could pass
for a flare. The plan says the SAP flux must show the same event. This script measures the signed
SAP change of every vetted event of 1e34 erg or more on the stars not set aside as fast rotators,
from the light curves under data/lc (a rise for a flare, a fall for a dip), and repeats the
injections in those light curves with the sign required, to see whether completeness changes.

python sap_sign.py <data dir> <run dir> <results dir>
Writes sap_sign.csv (each event and its signed SAP change) and sap_sign.json.
"""
import json
import os
import sys

import numpy as np
import pandas as pd

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import analyze as an  # noqa: E402
import inject as ij  # noqa: E402
import run_search as rs  # noqa: E402
import sflare as sf  # noqa: E402


def signed_change(lc, r, c):
    """The SAP flux change over the candidate's cadences at half its peak or more, against the
    cadences around it, exactly as side_signals measures it, but keeping its sign."""
    n = len(lc['t'])
    s, e, ipk = c['i0'], c['i1'], c['ipk']
    base = np.r_[max(0, s - 40):max(0, s - 5), min(n, e + 10):min(n, e + 45)]
    top = np.arange(s, e + 1)
    top = top[r[top] >= 0.5 * r[ipk]]
    return float(np.mean(lc['sap'][top]) - np.median(lc['sap'][base]))


_side = sf.side_signals


def side_signals_signed(lc, r, c, crowdsap=1.0):
    """side_signals with the SAP ratio signed: an injected flare must raise the SAP flux."""
    out = _side(lc, r, c, crowdsap)
    if np.isfinite(out['sap_ratio']) and signed_change(lc, r, c) <= 0:
        out['sap_ratio'] = -out['sap_ratio']
    return out


def main():
    data, run, out = sys.argv[1:4]
    os.makedirs(out, exist_ok=True)
    st = an.Study(data, run, 'main')
    nonfast = set(st.stars.tic[~st.stars.tess_fast])
    cands = an.coincidences(an.table(run, 'cands'))
    ok = cands.vetted.astype(bool) & (cands.energy >= 1e34) & (cands.n_coincident <= an.COINCIDENCE_MAX) & cands.tic.isin(nonfast)
    rows = []
    for f, g in cands[ok].groupby('file'):
        lc, meta = sf.read_lc(os.path.join(data, 'lc', f))
        _, sigma, r, found = sf.search(lc)
        dips = sf.candidates(lc['t'], -r, sigma)
        for _, e in g.iterrows():
            pool, sign = (found, 1.0) if e.kind == 'flare' else (dips, -1.0)
            c = min(pool, key=lambda x: abs(x['tpk'] - e.tpk))
            d = signed_change(lc, sign * r, c)
            rows.append({'file': f, 'tic': int(e.tic), 'kind': e.kind, 'tpk': e.tpk, 'energy': e.energy,
                         'sap_change': d, 'same_sign': bool(d * sign > 0), 'eclipser_star': int(e.tic) not in st.keep})
    ev = pd.DataFrame(rows)
    ev.to_csv(os.path.join(out, 'sap_sign.csv'), index=False)

    # Injections again, in every light curve of the confirmatory stars searched here, with the sign required
    sf.side_signals = side_signals_signed
    stars = st.stars.set_index('tic')
    used = set(st.use.file)
    files = sorted(f for f in os.listdir(os.path.join(data, 'lc')) if f in used)
    lcs = st.use.set_index('file')
    changed, recovered = 0, 0
    for f in files:
        lc, meta = sf.read_lc(os.path.join(data, 'lc', f))
        star = stars.loc[int(lcs.loc[f].tic)]
        k = sf.energy_per_ed(float(star.Teff), float(star.Rad), os.path.join(data, 'tess-response-function.csv'))
        mine = ij.inject_recover(lc, k, np.random.default_rng(rs.seed_of(f)), crowdsap=meta['crowdsap'])
        ref = st.inj[(st.inj.file == f) & (st.inj.model == 'base')].set_index('log_e')
        for m in mine:
            a, b = m['n_rec_340'], int(ref.loc[m['log_e'], 'n_rec_340'])
            changed += abs(a - b)
            recovered += b
    sf.side_signals = _side
    flares = ev[ev.kind == 'flare']
    dips = ev[ev.kind == 'dip']
    res = {
        'flares_checked': int(len(flares)), 'flares_opposite_sign': int((~flares.same_sign).sum()),
        'counted_flares_opposite_sign': int((~flares.same_sign & ~flares.eclipser_star).sum()),
        'dips_checked': int(len(dips)), 'dips_opposite_sign': int((~dips.same_sign).sum()),
        'injection_light_curves': len(files), 'injections_recovered_above_1e34': int(recovered),
        'injections_changed_by_sign': int(changed),
    }
    with open(os.path.join(out, 'sap_sign.json'), 'w') as fh:
        json.dump(res, fh, indent=1, sort_keys=True)
    print(json.dumps(res, indent=1))


if __name__ == '__main__':
    main()