"""Search again the light curves data/external.json points at and compare every number with the
tables the whole search wrote.

python audit.py <data dir> <run dir> <results dir>

The light curves are read from <data dir>/lc/. A light curve agrees when its row, its candidates
(peak time, energy, whether vetted and why not), and its injection counts all match.
"""
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 run_search as rs  # noqa: E402

REL = 1e-6


def close(a, b):
    if isinstance(a, float) or isinstance(b, float):
        a, b = float(a), float(b)
        if np.isnan(a) and np.isnan(b):
            return True
        return abs(a - b) <= REL * max(1.0, abs(a), abs(b))
    return a == b


def main():
    data, run, out = sys.argv[1], sys.argv[2], sys.argv[3]
    os.makedirs(out, exist_ok=True)
    stars = pd.read_csv(os.path.join(data, 'sample_stars.csv')).set_index('tic')
    from analyze import table
    lcs = table(run, 'lcs').set_index('file')
    cands = table(run, 'cands')
    inj = table(run, 'inj')
    resp = os.path.join(data, 'tess-response-function.csv')
    files = sorted(f for f in os.listdir(os.path.join(data, 'lc')) if f.endswith('_lc.fits'))
    rows = []
    for f in files:
        ref = lcs.loc[f]
        star = stars.loc[int(ref.tic)]
        row, cs, ij = rs.process(os.path.join(data, 'lc', f), {'Teff': float(star.Teff), 'Rad': float(star.Rad)},
                                 resp, True, rs.seed_of(f))
        bad = [k for k in ('sha256', 'n_good', 'sigma_ppm', 'rvar', 'p_rot', 'ls_power', 'ls_amp', 'n_cands', 'n_dips')
               if not close(row[k], ref[k])]
        mine = sorted((c['kind'], round(c['tpk'], 6)) for c in cs)
        theirs_df = cands[cands.file == f]
        theirs = sorted(zip(theirs_df.kind, theirs_df.tpk.round(6)))
        cand_ok = mine == theirs
        if cand_ok:
            by = {(k, t): r for k, t, r in zip(theirs_df.kind, theirs_df.tpk.round(6), theirs_df.to_dict('records'))}
            for c in cs:
                r = by[(c['kind'], round(c['tpk'], 6))]
                if not (close(c['energy'], r['energy']) and bool(c['vetted']) == bool(r['vetted'])
                        and (c['reject'] or '') == (r['reject'] if isinstance(r['reject'], str) else '')):
                    cand_ok = False
        a = pd.DataFrame(ij)
        b = inj[inj.file == f]
        cols = [c for c in b.columns if c.startswith('n_')]
        key = ['model', 'log_e']
        m = a.merge(b, on=key, suffixes=('', '_ref'))
        inj_ok = len(m) == len(a) == len(b) and all((m[c] == m[c + '_ref']).all() for c in cols)
        rows.append(dict(file=f, row_ok=not bad, row_mismatch='|'.join(bad), candidates_ok=cand_ok,
                         injections_ok=bool(inj_ok), candidates=len(cs)))
    df = pd.DataFrame(rows)
    df.to_csv(os.path.join(out, 'audit_by_light_curve.csv'), index=False)
    res = {'light_curves': int(len(df)), 'rows_agree': int(df.row_ok.sum()),
           'candidates_agree': int(df.candidates_ok.sum()), 'injections_agree': int(df.injections_ok.sum()),
           'all_agree': int((df.row_ok & df.candidates_ok & df.injections_ok).sum())}
    res['fraction_all_agree'] = res['all_agree'] / max(1, res['light_curves'])
    with open(os.path.join(out, 'audit.json'), 'w') as fh:
        json.dump(res, fh, indent=1, sort_keys=True)
    print(json.dumps(res, indent=1))


if __name__ == '__main__':
    main()