"""Exploratory: how wide are the flares the search recovers, against the widths the injections were
given? The counted flares are selected by the search, so their widths are compared with those of
injected flares recovered the same way, under each of the three duration models, rather than with
the injected law itself.

python width_check.py <data dir> <run dir> <results dir> <counted flares json>
Writes width_check.json: for each model, the median and central range of
log10(fitted FWHM / (11.5 min x (measured energy / 1e34 erg)^0.39)) over injections recovered at
1e34 erg or more in the confirmatory light curves under data/lc, and the same for the counted flares.
"""
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 law(e):
    return 11.5 * (e / 1e34) ** ij.DURATION_SLOPE


def recovered_widths(lc, k, rng, t_half_34, crowdsap):
    """inject.inject_recover's runs, keeping each recovery's fitted FWHM and measured energy."""
    t = lc['t']
    n = len(t)
    energies = np.repeat(ij.GRID, ij.PER_ENERGY)
    rng.shuffle(energies)
    out = []
    for run in np.array_split(energies, ij.RUNS):
        times = []
        for i in rng.permutation(n):
            ti = t[i]
            if all(abs(ti - x) >= ij.MIN_SEP_D for x in times):
                times.append(ti)
            if len(times) == len(run):
                break
        model = np.zeros(n)
        for log_e, ti in zip(run, times):
            e = 10 ** log_e
            fwhm = t_half_34 * (e / 1e34) ** ij.DURATION_SLOPE * 10 ** (ij.SCATTER_DEX * rng.normal())
            amp = e / k / (fwhm * sf.DAY_S * sf.DAVENPORT_AREA)
            w = (t > ti - 2 * fwhm - 0.01) & (t < ti + 30 * fwhm + 0.01)
            model[w] += sf.flare_model(t[w], ti, fwhm, amp)
        lci = dict(lc)
        lci['f'] = lc['f'] * (1 + model)
        lci['sap'] = lc['sap'] * (1 + crowdsap * model)
        _, sigma, r, cands = sf.search(lci)
        for c in cands:
            v = sf.vet(lci, r, sigma, c, crowdsap)
            if v['vetted'] and any(c['t0'] - 2 * sf.CADENCE_S / sf.DAY_S <= ti <= c['t1'] for ti in times):
                out.append((c['ed'] * k, v['fwhm'] * 1440))
    return out


def summary(e, w):
    x = np.log10(np.asarray(w) / law(np.asarray(e)))
    return {'n': int(len(x)), 'median_log_ratio': round(float(np.median(x)), 2),
            'p16': round(float(np.percentile(x, 16)), 2), 'p84': round(float(np.percentile(x, 84)), 2)}


def main():
    data, run, out, counted = sys.argv[1:5]
    st = an.Study(data, run, 'main')
    stars = st.stars.set_index('tic')
    lcs = st.use.set_index('file')
    files = sorted(f for f in os.listdir(os.path.join(data, 'lc')) if f in lcs.index)
    res = {'light_curves': len(files)}
    for name, scale in (('short', 0.5), ('base', 1.0), ('long', 2.5)):
        es, ws = [], []
        for j, f in enumerate(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'))
            rng = np.random.default_rng(rs.seed_of(f) + 1000 + j)
            for e, w in recovered_widths(lc, k, rng, ij.T_HALF_34_D * scale, meta['crowdsap']):
                if e >= 1e34:
                    es.append(e)
                    ws.append(w)
        res[name] = summary(es, ws)
    fl = pd.DataFrame(json.load(open(counted))['flares'])
    res['counted'] = summary(fl.energy.values, fl.fwhm.values * 1440)
    res['counted_median_fwhm_min'] = round(float(np.median(fl.fwhm.values * 1440)), 1)
    with open(os.path.join(out, 'width_check.json'), 'w') as fh:
        json.dump(res, fh, indent=1, sort_keys=True)
    print(json.dumps(res, indent=1))


if __name__ == '__main__':
    main()