"""Where on the detector did each counted flare happen? The difference image of the flare's
brightest cadences against the cadences around it, from the target pixel file, located against
the target's own image.

python localize.py <flares.csv> <tpf dir> <out.csv>

flares.csv has file, tpk, t0, t1 (as cands.csv). The target pixel file for a light curve
tess...-s_lc.fits is tess...-s_tp.fits in <tpf dir>.
"""
import os
import sys

import numpy as np
import pandas as pd

CAD_D = 120.0 / 86400.0
BOX = 2              # centroids over the 5x5 pixels around the target's brightest pixel
OFF_PX = 0.5         # off target: farther than this ...
OFF_SIGMA = 3.0      # ... and more than this many standard errors
DRAWS = 300


def centroid(img, r0, c0):
    sub = img[r0 - BOX:r0 + BOX + 1, c0 - BOX:c0 + BOX + 1]
    rr, cc = np.mgrid[r0 - BOX:r0 + BOX + 1, c0 - BOX:c0 + BOX + 1]
    w = np.clip(sub, 0, None)
    s = w.sum()
    if s <= 0:
        return np.nan, np.nan
    return float((w * rr).sum() / s), float((w * cc).sum() / s)


def locate(tpf, tpk, t0, t1, seed=0):
    from astropy.io import fits
    with fits.open(tpf, memmap=False) as h:
        d = h[1].data
        t = np.asarray(d['TIME'], float)
        flux = np.asarray(d['FLUX'], float)
        q = np.asarray(d['QUALITY'], np.int64)
        ap = np.asarray(h[2].data, np.int64)
    good = np.isfinite(t) & ((q & (1 | 2 | 4 | 8 | 16 | 32 | 128)) == 0) & np.all(np.isfinite(flux), axis=(1, 2))
    t, flux = t[good], flux[good]
    inap = (ap & 2) > 0
    phot = flux[:, inap].sum(axis=1)
    near = (t > t0 - 45 * CAD_D) & (t < t1 + 45 * CAD_D)
    base = near & ((t < t0 - 5 * CAD_D) | (t > t1 + 10 * CAD_D))
    inside = (t >= t0) & (t <= t1)
    if base.sum() < 10 or inside.sum() == 0:
        return dict(n_top=0, offset_px=np.nan, offset_sigma=np.nan, on_target=np.nan)
    excess = phot - np.median(phot[base])
    peak = excess[inside].max()
    top = inside & (excess >= 0.5 * peak)
    quiet = flux[base].mean(axis=0)
    diff = flux[top].mean(axis=0) - quiet
    r0, c0 = np.unravel_index(np.nanargmax(np.where(inap, quiet, -np.inf)), quiet.shape)
    r0 = int(np.clip(r0, BOX, quiet.shape[0] - BOX - 1))
    c0 = int(np.clip(c0, BOX, quiet.shape[1] - BOX - 1))
    cq = centroid(quiet, r0, c0)
    cd = centroid(diff, r0, c0)
    off = float(np.hypot(cd[0] - cq[0], cd[1] - cq[1]))
    # Noise of the difference image from the baseline cadences' scatter
    sig = flux[base].std(axis=0) * np.sqrt(1.0 / top.sum() + 1.0 / base.sum())
    rng = np.random.default_rng(seed)
    offs = []
    for _ in range(DRAWS):
        c = centroid(diff + rng.normal(0, 1, diff.shape) * sig, r0, c0)
        offs.append(np.hypot(c[0] - cq[0], c[1] - cq[1]))
    offs = np.array(offs)
    err = float(np.nanstd(offs))
    on = not (off > OFF_PX and off > OFF_SIGMA * err)
    return dict(n_top=int(top.sum()), peak_e_per_s=float(peak), offset_px=off, offset_err_px=err,
                offset_sigma=off / err if err > 0 else np.nan, on_target=bool(on),
                row_quiet=cq[0], col_quiet=cq[1], row_flare=cd[0], col_flare=cd[1])


def main():
    flares = pd.read_csv(sys.argv[1])
    tpf_dir, out = sys.argv[2], sys.argv[3]
    rows = []
    for k, f in flares.iterrows():
        tpf = os.path.join(tpf_dir, f.file.replace('_lc.fits', '_tp.fits'))
        res = dict(file=f.file, tpk=f.tpk)
        if os.path.exists(tpf):
            res.update(locate(tpf, f.tpk, f.t0, f.t1, seed=k))
        rows.append(res)
    pd.DataFrame(rows).to_csv(out, index=False)


if __name__ == '__main__':
    main()