"""Search again the light curves under data/lc and compare every column of their rows in the
search's tables: the light curve's row, each candidate's row, and each injection row.

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

Numbers agree within a relative 1e-6 (absolute 1e-9 near zero), so the last bits platforms'
math libraries round differently don't count; text and flags must match exactly. The template
fits (code/sflare.py, shape) are iterative, and on another platform can stop a little apart, so
their parameters are reported apart from everything else: the decisions, energies, counts and
every other field, and the largest relative difference in a fitted parameter.
"""
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
from analyze import table  # noqa: E402


def same(a, b):
    if isinstance(a, (bool, np.bool_)) or isinstance(b, (bool, np.bool_)):
        return bool(a) == bool(b)
    try:
        fa, fb = float(a), float(b)
    except (TypeError, ValueError):
        sa = '' if a is None or (isinstance(a, float) and np.isnan(a)) else str(a)
        sb = '' if b is None or (isinstance(b, float) and np.isnan(b)) else str(b)
        return sa == sb
    if np.isnan(fa) and np.isnan(fb):
        return True
    if np.isinf(fa) or np.isinf(fb):
        return fa == fb
    return abs(fa - fb) <= max(1e-9, 1e-6 * max(abs(fa), abs(fb)))


FIT = {'fwhm', 'amp_fit', 'tpk_fit', 'ed_fit', 'chi2_flare', 'chi2_gauss', 'chi2_null', 'dchi2'}


def rel(a, b):
    try:
        fa, fb = float(a), float(b)
    except (TypeError, ValueError):
        return np.inf
    if np.isnan(fa) or np.isnan(fb):
        return 0.0 if np.isnan(fa) and np.isnan(fb) else np.inf
    return abs(fa - fb) / max(1e-12, abs(fa), abs(fb))


def main():
    data, run, out = sys.argv[1:4]
    os.makedirs(out, exist_ok=True)
    stars = pd.read_csv(os.path.join(data, 'sample_stars.csv')).set_index('tic')
    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 in lcs.index)
    fields = mismatched = 0
    other = 0          # mismatches outside the fitted parameters
    max_fit = 0.0      # the largest relative difference in a fitted parameter
    fit_vetted = 0     # fitted parameters that differ on a vetted candidate
    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 = []
        for k in rs.LC_COLS[1:]:
            fields += 1
            if not same(row[k], ref[k]):
                bad.append('lc.' + k)
        mine = {(c['kind'], round(c['tpk'], 6)): c for c in cs}
        theirs = {(k, round(t, 6)): r for k, t, r in zip(cands[cands.file == f].kind, cands[cands.file == f].tpk,
                                                           cands[cands.file == f].to_dict('records'))}
        if set(mine) != set(theirs):
            bad.append('candidates')
        else:
            for key, c in mine.items():
                for k in rs.CAND_COLS:
                    fields += 1
                    if not same(c.get(k, np.nan), theirs[key][k]):
                        bad.append('cand.' + k)
                        if k in FIT:
                            max_fit = max(max_fit, rel(c.get(k, np.nan), theirs[key][k]))
                            fit_vetted += int(bool(c['vetted']) or bool(theirs[key]['vetted']))
                        else:
                            other += 1
        a = {(r['model'], r['log_e']): r for r in ij}
        b = {(r['model'], r['log_e']): r for r in inj[inj.file == f].to_dict('records')}
        if set(a) != set(b):
            bad.append('injections')
        else:
            for key, r in a.items():
                for k in rs.INJ_COLS:
                    fields += 1
                    if not same(r.get(k, np.nan), b[key][k]):
                        bad.append('inj.' + k)
        other += sum(1 for b in bad if not b.startswith('cand.') or b in ('candidates',))
        mismatched += len(bad)
        rows.append({'file': f, 'agree': not bad, 'mismatches': '|'.join(sorted(set(bad)))})
    df = pd.DataFrame(rows)
    df.to_csv(os.path.join(out, 'audit_all_by_light_curve.csv'), index=False)
    res = {'light_curves': int(len(df)), 'light_curves_agree': int(df.agree.sum()),
           'fields_compared': int(fields), 'fields_mismatched': int(mismatched),
           'all_but_fits_agree': bool(other == 0), 'fit_differences_on_vetted': int(fit_vetted),
           'max_fit_relative_difference': float(max_fit), 'fits_within_one_percent': bool(max_fit <= 0.01)}
    with open(os.path.join(out, 'audit_all.json'), 'w') as fh:
        json.dump(res, fh, indent=1, sort_keys=True)
    print(json.dumps(res, indent=1))


if __name__ == '__main__':
    main()