"""Superflare search in TESS 2-minute SPOC light curves.

One light curve (one star, one sector) at a time: detrend, find flare candidates, measure them,
vet them, give them bolometric energies, and measure how often flares of known energy injected
into the same light curve are recovered. Everything here is deterministic given its inputs and
the seed it is handed.
"""
import numpy as np
from scipy import special
from scipy.optimize import least_squares

CADENCE_S = 120.0
DAY_S = 86400.0
# Cadences SPOC flags for reasons that make the flux unusable. Impulsive outliers (512) are kept:
# SPOC flags many flare peaks that way, and their fluxes stay in PDCSAP_FLUX.
BAD_QUALITY = 1 | 2 | 4 | 8 | 16 | 32 | 128
DESAT = 32  # reaction wheel desaturation (momentum dump)

SEGMENT_GAP_D = 0.25   # a gap longer than this starts a new segment
BIN_D = 0.05           # trend bins (72 minutes)
TREND_BINS = 5         # running median over this many bins (6 hours)
NOISE_BINS = 11        # running median of per-bin scatter over this many bins
CORE_Z = 3.0           # every cadence of a candidate's core is at least this many sigma high
CORE_N = 3             # ... for at least this many consecutive cadences
EDGE_Z = 1.0           # a candidate extends while the flux stays this many sigma high
FIT_PAD = 20           # cadences either side of a candidate used in shape fits
SIGMA_SB = 5.670374419e-5  # erg cm-2 s-1 K-4
R_SUN_CM = 6.957e10
H = 6.62607015e-27
C = 2.99792458e10
K_B = 1.380649e-16


# --------------------------------------------------------------------------------------------
# Reading
# --------------------------------------------------------------------------------------------
def read_lc(path):
    """Columns and header values of a SPOC light-curve file, good cadences only."""
    from astropy.io import fits
    with fits.open(path, memmap=False) as h:
        h0, h1, d = h[0].header, h[1].header, h[1].data
        t = np.asarray(d['TIME'], float)
        f = np.asarray(d['PDCSAP_FLUX'], float)
        fe = np.asarray(d['PDCSAP_FLUX_ERR'], float)
        q = np.asarray(d['QUALITY'], np.int64)
        bkg = np.asarray(d['SAP_BKG'], float)
        sap = np.asarray(d['SAP_FLUX'], float)
        c1 = np.asarray(d['MOM_CENTR1'], float)
        c2 = np.asarray(d['MOM_CENTR2'], float)
        meta = {
            'tic': int(h0['TICID']), 'sector': int(h0['SECTOR']), 'camera': int(h0['CAMERA']),
            'ccd': int(h0['CCD']), 'tessmag': float(h0.get('TESSMAG') or np.nan),
            'crowdsap': float(h1.get('CROWDSAP') or np.nan), 'flfrcsap': float(h1.get('FLFRCSAP') or np.nan),
            'cdpp1': float(h1.get('CDPP1_0') or np.nan),
        }
    # Every cadence the file holds, before any is dropped, so desaturation events still count
    desat_t = t[(q & DESAT) > 0]
    good = np.isfinite(t) & np.isfinite(f) & np.isfinite(fe) & ((q & BAD_QUALITY) == 0)
    lc = {k: v[good] for k, v in dict(t=t, f=f, fe=fe, q=q, bkg=bkg, sap=sap, c1=c1, c2=c2).items()}
    lc['desat_t'] = desat_t[np.isfinite(desat_t)]
    return lc, meta


# --------------------------------------------------------------------------------------------
# Detrending
# --------------------------------------------------------------------------------------------
def segments(t):
    br = np.where(np.diff(t) > SEGMENT_GAP_D)[0]
    return list(zip(np.r_[0, br + 1], np.r_[br + 1, len(t)]))


def _running_nanmedian(x, k):
    h = k // 2
    xp = np.concatenate([np.full(h, np.nan), x, np.full(h, np.nan)])
    w = np.lib.stride_tricks.sliding_window_view(xp, k)
    with np.errstate(all='ignore'):
        import warnings
        with warnings.catch_warnings():
            warnings.simplefilter('ignore', RuntimeWarning)
            return np.nanmedian(w, axis=1)


def _binned(t, y, use, t0):
    idx = np.floor((t - t0) / BIN_D).astype(np.int64)
    nb = int(idx.max()) + 1
    med = np.full(nb, np.nan)
    tc = (np.arange(nb) + 0.5) * BIN_D + t0
    mad = np.full(nb, np.nan)
    order = np.argsort(idx, kind='stable')
    bounds = np.searchsorted(idx[order], np.arange(nb + 1))
    for b in range(nb):
        sel = order[bounds[b]:bounds[b + 1]]
        sel = sel[use[sel]]
        if len(sel) >= 8:
            v = y[sel]
            m = np.median(v)
            med[b] = m
            mad[b] = np.median(np.abs(v - m))
    return tc, med, mad


def detrend(t, f, mask):
    """Trend and per-cadence noise (relative flux), ignoring cadences in mask."""
    trend = np.full_like(f, np.nan)
    sigma = np.full_like(f, np.nan)
    for a, b in segments(t):
        ts, fs, ms = t[a:b], f[a:b], mask[a:b]
        if b - a < 30:
            continue
        tc, med, _ = _binned(ts, fs, ~ms, ts[0])
        rm = _running_nanmedian(med, TREND_BINS)
        ok = np.isfinite(rm)
        if ok.sum() < 2:
            continue
        tr = np.interp(ts, tc[ok], rm[ok])
        trend[a:b] = tr
        r = fs / tr - 1.0
        _, _, mad = _binned(ts, r, ~ms, ts[0])
        sm = _running_nanmedian(mad, NOISE_BINS)
        ok = np.isfinite(sm) & (sm > 0)
        if ok.sum() < 2:
            continue
        sigma[a:b] = 1.4826 * np.interp(ts, tc[ok], sm[ok])
    return trend, sigma


# --------------------------------------------------------------------------------------------
# Candidates
# --------------------------------------------------------------------------------------------
def find_runs(z, t):
    """Candidates: CORE_N or more consecutive cadences at CORE_Z sigma or more, with no gap."""
    hi = np.isfinite(z) & (z >= CORE_Z)
    out = []
    i, n = 0, len(z)
    while i < n:
        if not hi[i]:
            i += 1
            continue
        j = i
        while j + 1 < n and hi[j + 1] and (t[j + 1] - t[j]) < 1.5 * CADENCE_S / DAY_S:
            j += 1
        if j - i + 1 >= CORE_N:
            out.append((i, j))
        i = j + 1
    # Merge cores separated by two cadences or fewer
    merged = []
    for a, b in out:
        if merged and a - merged[-1][1] <= 3:
            merged[-1] = (merged[-1][0], b)
        else:
            merged.append((a, b))
    return merged


def extent(z, t, a, b):
    """Grow a core while the flux stays EDGE_Z sigma high: back to the rise, on to the decay."""
    n = len(z)
    s = a
    while s - 1 >= 0 and np.isfinite(z[s - 1]) and z[s - 1] > EDGE_Z and (t[s] - t[s - 1]) < 1.5 * CADENCE_S / DAY_S and a - s < 15:
        s -= 1
    e = b
    while e + 1 < n and (t[e + 1] - t[e]) < 1.5 * CADENCE_S / DAY_S:
        nxt = z[e + 1]
        nxt2 = z[e + 2] if e + 2 < n else -np.inf
        if np.isfinite(nxt) and (nxt > EDGE_Z or (np.isfinite(nxt2) and nxt2 > EDGE_Z and e + 2 < n)):
            e += 1
        else:
            break
        if e - b > 360:
            break
    return s, e


def candidates(t, r, sigma):
    z = r / sigma
    out = []
    for a, b in find_runs(z, t):
        s, e = extent(z, t, a, b)
        seg = r[s:e + 1]
        ipk = s + int(np.argmax(seg))
        out.append(dict(i0=s, i1=e, ipk=ipk, t0=t[s], t1=t[e], tpk=t[ipk],
                        n=e - s + 1, n_core=b - a + 1, amp=float(r[ipk]), z_pk=float(z[ipk]),
                        ed=float(np.sum(seg) * CADENCE_S),
                        snr=float(np.sum(z[s:e + 1]) / np.sqrt(e - s + 1))))
    # Overlapping extents become one candidate
    out.sort(key=lambda c: c['i0'])
    merged = []
    for c in out:
        if merged and c['i0'] <= merged[-1]['i1']:
            m = merged[-1]
            if c['amp'] > m['amp']:
                c['i0'] = min(c['i0'], m['i0']); c['i1'] = max(c['i1'], m['i1'])
                merged[-1] = c
            else:
                m['i1'] = max(m['i1'], c['i1'])
        else:
            merged.append(c)
    for c in merged:
        s, e = c['i0'], c['i1']
        c['t0'], c['t1'], c['n'] = t[s], t[e], e - s + 1
        c['ed'] = float(np.sum(r[s:e + 1]) * CADENCE_S)
        c['snr'] = float(np.sum(z[s:e + 1]) / np.sqrt(e - s + 1))
    return merged


# --------------------------------------------------------------------------------------------
# Shapes
# --------------------------------------------------------------------------------------------
_SUB = (np.arange(15) + 0.5) / 15.0 - 0.5  # sub-exposure offsets, in cadences


def davenport(x):
    """Davenport et al. (2014) flare template; x is time from peak in units of the FWHM."""
    y = np.zeros_like(x)
    r = (x >= -1) & (x <= 0)
    y[r] = 1 + 1.941 * x[r] - 0.175 * x[r] ** 2 - 2.246 * x[r] ** 3 - 1.125 * x[r] ** 4
    d = x > 0
    y[d] = 0.6890 * np.exp(-1.600 * x[d]) + 0.3030 * np.exp(-0.2783 * x[d])
    return y


DAVENPORT_AREA = 1.0 - 1.941 / 2 - 0.175 / 3 + 2.246 / 4 - 1.125 / 5 + 0.6890 / 1.600 + 0.3030 / 0.2783


def flare_model(t, tpk, fwhm, amp):
    """The template integrated over each 2-minute exposure."""
    ts = t[:, None] + _SUB[None, :] * CADENCE_S / DAY_S
    return amp * davenport((ts - tpk) / fwhm).mean(axis=1)


def gauss_model(t, tc, w, amp):
    ts = t[:, None] + _SUB[None, :] * CADENCE_S / DAY_S
    return amp * np.exp(-0.5 * ((ts - tc) / w) ** 2).mean(axis=1)


def shape(t, r, sigma, c):
    """Fit the flare template and a symmetric bump around a candidate; compare them."""
    n = len(t)
    a, b = max(0, c['i0'] - FIT_PAD), min(n, c['i1'] + 1 + 2 * FIT_PAD)
    tt, rr, ss = t[a:b], r[a:b], sigma[a:b]
    ok = np.isfinite(rr) & np.isfinite(ss) & (ss > 0)
    tt, rr, ss = tt[ok], rr[ok], ss[ok]
    out = dict(fwhm=np.nan, amp_fit=np.nan, tpk_fit=np.nan, chi2_flare=np.nan, chi2_gauss=np.nan,
               chi2_null=np.nan, ed_fit=np.nan, dchi2=np.nan)
    if len(tt) < 8:
        return out
    tpk, amp = c['tpk'], max(c['amp'], 1e-6)
    dur = max(c['t1'] - c['t0'], 2 * CADENCE_S / DAY_S)
    f0 = [tpk, max(dur / 4, 2 * CADENCE_S / DAY_S), amp]
    lo = [tpk - 5 * CADENCE_S / DAY_S, 0.3 * CADENCE_S / DAY_S, 0.0]
    hi = [tpk + 5 * CADENCE_S / DAY_S, 1.0, 50 * amp]
    try:
        pf = least_squares(lambda p: (flare_model(tt, *p) - rr) / ss, f0, bounds=(lo, hi), x_scale=[1e-3, 1e-3, amp]).x
        g0 = [tpk, max(dur / 4, CADENCE_S / DAY_S), amp]
        pg = least_squares(lambda p: (gauss_model(tt, *p) - rr) / ss, g0,
                           bounds=([tpk - 10 * CADENCE_S / DAY_S, 0.3 * CADENCE_S / DAY_S, 0.0], [tpk + 10 * CADENCE_S / DAY_S, 1.0, 50 * amp]),
                           x_scale=[1e-3, 1e-3, amp]).x
    except Exception:
        return out
    cf = float(np.sum(((flare_model(tt, *pf) - rr) / ss) ** 2))
    cg = float(np.sum(((gauss_model(tt, *pg) - rr) / ss) ** 2))
    out.update(tpk_fit=float(pf[0]), fwhm=float(pf[1]), amp_fit=float(pf[2]), chi2_flare=cf, chi2_gauss=cg,
               chi2_null=float(np.sum((rr / ss) ** 2)), ed_fit=float(pf[2] * pf[1] * DAY_S * DAVENPORT_AREA),
               dchi2=cg - cf, npts=int(len(tt)))
    return out


# --------------------------------------------------------------------------------------------
# Vetting signals measured from the light curve itself
# --------------------------------------------------------------------------------------------
def side_signals(lc, r, c, crowdsap=1.0):
    """Quality flags, momentum dumps, and what the background, the centroid, and the
    uncorrected (SAP) flux did during a candidate."""
    t, q = lc['t'], lc['q']
    s, e, ipk = c['i0'], c['i1'], c['ipk']
    n = len(t)
    w0, w1 = max(0, s - 5), min(n, e + 6)
    flagged = int(np.sum((q[w0:w1] & ~512) != 0))
    dt_desat = float(np.min(np.abs(lc['desat_t'] - c['tpk']))) if len(lc['desat_t']) else np.inf
    # Baseline: up to 35 cadences either side, outside the candidate
    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]]
    res = dict(flagged=flagged, dt_desat=dt_desat, bkg_ratio=np.nan, bkg_z=np.nan, sap_ratio=np.nan,
               cen_dx=np.nan, cen_dy=np.nan, cen_z=np.nan, cen_offset=np.nan, cen_offset_z=np.nan)
    if len(base) < 10 or len(top) == 0:
        return res

    def mad(v):
        return 1.4826 * np.median(np.abs(v - np.median(v)))
    sap_med = np.median(lc['sap'][base])
    dflux = np.mean(lc['sap'][top]) - sap_med
    dbkg = np.mean(lc['bkg'][top]) - np.median(lc['bkg'][base])
    if dflux != 0:
        res['bkg_ratio'] = float(abs(dbkg) / abs(dflux))
    sb = mad(lc['bkg'][base]) / np.sqrt(len(top))
    if sb > 0:
        res['bkg_z'] = float(abs(dbkg) / sb)
    rel = float(np.mean(r[top]))
    if rel > 0 and sap_med > 0 and np.isfinite(crowdsap) and crowdsap > 0:
        res['sap_ratio'] = float(abs(dflux) / sap_med / (crowdsap * rel))
    for k, name in (('c1', 'cen_dx'), ('c2', 'cen_dy')):
        v = lc[k]
        if np.all(np.isfinite(v[base])) and np.all(np.isfinite(v[top])):
            res[name] = float(np.mean(v[top]) - np.median(v[base]))
    if np.isfinite(res['cen_dx']) and np.isfinite(res['cen_dy']):
        sx = mad(lc['c1'][base]) / np.sqrt(len(top))
        sy = mad(lc['c2'][base]) / np.sqrt(len(top))
        if sx > 0 and sy > 0:
            res['cen_z'] = float(np.hypot(res['cen_dx'] / sx, res['cen_dy'] / sy))
            # A source that brightens by a fraction f of the aperture's flux moves the centroid
            # by f times its distance from the centroid: the shift over f is that distance
            f = abs(dflux) / sap_med if sap_med > 0 else np.nan
            if np.isfinite(f) and f > 0:
                res['cen_offset'] = float(np.hypot(res['cen_dx'], res['cen_dy']) / f)
                res['cen_offset_z'] = res['cen_z']
    return res


# --------------------------------------------------------------------------------------------
# Energies
# --------------------------------------------------------------------------------------------
_RESP = None


def response(path):
    global _RESP
    if _RESP is None:
        d = np.genfromtxt(path, delimiter=',', names=True)
        _RESP = (d['lambda'] * 1e-7, d['resp'])  # nm to cm
    return _RESP


def planck_photon_weighted(lam_cm, T):
    """B_lambda(T) * lambda, the photon-counting weight Vasilyev et al. (2024) integrate."""
    x = H * C / (lam_cm * K_B * T)
    with np.errstate(over='ignore'):
        b = 2 * H * C ** 2 / lam_cm ** 5 / np.expm1(x)
    return b * lam_cm


def energy_per_ed(teff, radius_rsun, resp_path, t_flare=9000.0):
    """Bolometric flare energy per second of equivalent duration (erg/s), for one star.

    E = sigma T_fl^4 pi R^2 [int R B(T*) lambda dl / int R B(T_fl) lambda dl] * ED
    (Shibayama et al. 2013, as Vasilyev et al. 2024 apply it).
    """
    lam, resp = response(resp_path)
    num = np.trapezoid(resp * planck_photon_weighted(lam, teff), lam)
    den = np.trapezoid(resp * planck_photon_weighted(lam, t_flare), lam)
    return SIGMA_SB * t_flare ** 4 * np.pi * (radius_rsun * R_SUN_CM) ** 2 * num / den


# --------------------------------------------------------------------------------------------
# One light curve
# --------------------------------------------------------------------------------------------
def search(lc, flux=None):
    """Detrend twice (masking candidates the second time) and return residuals and candidates."""
    t = lc['t']
    f = lc['f'] if flux is None else flux
    mask = np.zeros(len(t), bool)
    trend, sigma = detrend(t, f, mask)
    r = f / trend - 1.0
    cands = candidates(t, r, sigma)
    if cands:
        for c in cands:
            mask[max(0, c['i0'] - 5):min(len(t), c['i1'] + 1 + max(10, 2 * c['n']))] = True
        trend, sigma = detrend(t, f, mask)
        r = f / trend - 1.0
        cands = candidates(t, r, sigma)
    return trend, sigma, r, cands


def rotation(t, trend):
    """Lomb-Scargle of the trend, 0.2 to 12 days: best period, its power and semi-amplitude."""
    from astropy.timeseries import LombScargle
    ok = np.isfinite(trend)
    if ok.sum() < 500:
        return np.nan, np.nan, np.nan
    tt, yy = t[ok][::5], trend[ok][::5] / np.nanmedian(trend) - 1.0
    freq = np.linspace(1 / 12.0, 1 / 0.2, 3000)
    ls = LombScargle(tt, yy)
    p = ls.power(freq)
    k = int(np.argmax(p))
    model = ls.model(tt, freq[k])
    amp = 0.5 * (np.percentile(model, 99) - np.percentile(model, 1))
    return float(1 / freq[k]), float(p[k]), float(amp)


# --------------------------------------------------------------------------------------------
# Vetting rules
# --------------------------------------------------------------------------------------------
EDGE_MIN_D = 0.02      # a candidate needs this much data before and after it in its segment
DESAT_MIN_D = 0.02     # ... and this far from any momentum dump
RISE_MAX = 4           # cadences from half maximum to the peak
Z_PEAK_MIN = 5.0       # the peak cadence at least this many sigma high
OFFSET_MAX_PX = 0.5    # the source the centroid shift implies at most this far from the centroid ...
OFFSET_Z_MAX = 3.0     # ... unless the shift is not significant
FWHM_MAX_D = 1.0 / 24  # fitted FWHM at most 1 hour
BKG_RATIO_MAX = 0.25   # background change at most this fraction of the flux change ...
BKG_Z_MAX = 4.0        # ... unless it is not significant
SAP_RANGE = (0.5, 2.0)  # the uncorrected flux must show the same event


def vet(lc, r, sigma, c, crowdsap=1.0):
    """Measure a candidate and apply the vetting rules; r is signed so the event is positive."""
    t = lc['t']
    out = shape(t, r, sigma, c)
    out.update(side_signals(lc, r, c, crowdsap))
    # Momentum dumps: distance from the candidate's extent, zero if one falls inside it
    if len(lc['desat_t']):
        d = np.where((lc['desat_t'] >= c['t0']) & (lc['desat_t'] <= c['t1']), 0.0,
                     np.minimum(np.abs(lc['desat_t'] - c['t0']), np.abs(lc['desat_t'] - c['t1'])))
        out['dt_desat'] = float(d.min())
    for a, b in segments(t):
        if a <= c['i0'] < b:
            out['dt_edge0'] = float(t[c['i0']] - t[a])
            out['dt_edge1'] = float(t[b - 1] - t[c['i1']])
            break
    # Rise and decay at half maximum: cadences from the first at half the peak or more to the
    # peak, and from the peak to the last
    half = r[c['i0']:c['i1'] + 1] >= 0.5 * r[c['ipk']]
    k = c['ipk'] - c['i0']
    out['rise'] = int(k - np.argmax(half))
    out['decay'] = int(len(half) - 1 - np.argmax(half[::-1]) - k)
    why = []
    if out.get('dt_edge0', 0) < EDGE_MIN_D or out.get('dt_edge1', 0) < EDGE_MIN_D:
        why.append('edge')
    if out['dt_desat'] < DESAT_MIN_D:
        why.append('desat')
    if out['flagged'] > 0:
        why.append('flagged')
    if out['rise'] > RISE_MAX:
        why.append('slow_rise')
    if out['rise'] > out['decay']:
        why.append('rise')
    if c['z_pk'] < Z_PEAK_MIN:
        why.append('weak')
    if not np.isfinite(out['dchi2']):
        why.append('nofit')
    elif out['fwhm'] > FWHM_MAX_D:
        why.append('long')
    if (np.isfinite(out['cen_offset']) and out['cen_offset'] > OFFSET_MAX_PX
            and np.isfinite(out['cen_offset_z']) and out['cen_offset_z'] > OFFSET_Z_MAX):
        why.append('offset')
    if np.isfinite(out['bkg_ratio']) and out['bkg_ratio'] > BKG_RATIO_MAX and np.isfinite(out['bkg_z']) and out['bkg_z'] > BKG_Z_MAX:
        why.append('bkg')
    if not (np.isfinite(out['sap_ratio']) and SAP_RANGE[0] <= out['sap_ratio'] <= SAP_RANGE[1]):
        why.append('sap')
    out['reject'] = '|'.join(why)
    out['vetted'] = not why
    return out