"""Inject flares of known energy into a light curve and see which the search recovers."""
import numpy as np
import sflare as sf

GRID = np.round(np.arange(33.0, 35.61, 0.2), 1)   # log10 bolometric energy, erg
PER_ENERGY = 6
RUNS = 6
MIN_SEP_D = 1.0
# FWHM of the injected flares: T_HALF_34 at 1e34 erg, scaling as E^0.39 (Maehara et al. 2015),
# with a log-normal scatter of SCATTER_DEX
T_HALF_34_D = 11.5 / 1440
DURATION_SLOPE = 0.39
SCATTER_DEX = 0.27
# Recovered with measured energy at or above each of these (log10 erg), named by 10 times the log
THRESHOLDS = {f'{int(round(x * 10))}': float(x) for x in np.round(np.arange(33.0, 35.51, 0.1), 1)}
THRESHOLD_NAMES = list(THRESHOLDS)


def inject_recover(lc, k, rng, t_half_34=T_HALF_34_D, crowdsap=1.0):
    t = lc['t']
    n = len(t)
    energies = np.repeat(GRID, PER_ENERGY)
    rng.shuffle(energies)
    tally = {e: dict(n_inj=0, n_rec_any=0, ratios=[], **{f'n_rec_{k_}': 0 for k_ in THRESHOLDS}) for e in GRID}
    for run in np.array_split(energies, RUNS):
        # Times: random good cadences at least MIN_SEP_D apart
        times = []
        for i in rng.permutation(n):
            ti = t[i]
            if all(abs(ti - x) >= MIN_SEP_D for x in times):
                times.append(ti)
            if len(times) == len(run):
                break
        model = np.zeros(n)
        inj = []
        for log_e, ti in zip(run, times):
            e = 10 ** log_e
            ed = e / k
            fwhm = t_half_34 * (e / 1e34) ** DURATION_SLOPE * 10 ** (SCATTER_DEX * rng.normal())
            amp = ed / (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)
            inj.append((log_e, ti))
        lci = dict(lc)
        lci['f'] = lc['f'] * (1 + model)
        lci['sap'] = lc['sap'] * (1 + crowdsap * model)
        _, sigma, r, cands = sf.search(lci)
        found = []
        for c in cands:
            v = sf.vet(lci, r, sigma, c, crowdsap)
            if v['vetted']:
                found.append((c['t0'] - 2 * sf.CADENCE_S / sf.DAY_S, c['t1'], c['ed'] * k))
        for log_e, ti in inj:
            d = tally[log_e]
            d['n_inj'] += 1
            hit = [e_obs for a, b, e_obs in found if a <= ti <= b]
            if hit:
                e_obs = max(hit)
                d['n_rec_any'] += 1
                d['ratios'].append(np.log10(e_obs) - log_e)
                for name, thr in THRESHOLDS.items():
                    if np.log10(e_obs) >= thr:
                        d[f'n_rec_{name}'] += 1
    out = []
    for log_e, d in tally.items():
        ratios = d.pop('ratios')
        out.append(dict(log_e=float(log_e), med_ratio=float(np.median(ratios)) if ratios else np.nan, **d))
    return out