"""Moving-object search, pass 1: per-sub source detection and frame registration.

Every calibrated sub is opened one at a time (they are 61 MB each and there is
only a few GB of RAM), background-subtracted with sep, and its sources
extracted. The detection list, not the pixels, is what gets carried forward.

Registration is done on the coordinates rather than the images. astroalign
matches asterisms between each sub's star list and the reference sub's star
list and returns a similarity transform; applying that transform to the
detection coordinates puts every sub's sources into one common pixel grid
without ever warping a 15 Mpx array. The reference is Luminance_002, the same
frame the master stack was registered to, so the master's WCS applies directly
to the common grid and every detection can be turned into RA/Dec.

Output: _mo_dets.npz with one record array per frame plus the transform
parameters and mid-exposure times.
"""
import os
import sys

import astroalign as aa
import numpy as np
import sep
from astropy.io import fits

import layout

SRC = layout.SESSION
OUT = layout.SESSION
CACHE = layout.path("_mo_dets.npz")
REFERENCE = "Luminance_002"

aa.MIN_MATCHES_FRACTION = 0.6
aa.NUM_NEAREST_NEIGHBORS = 8

# Detection threshold in sigma. Deliberately low: a moving object is only in
# any one sub for 300 s, so it is much fainter per-frame than in the 60 min
# master. Spurious detections are cheap here because the tracklet linker
# demands a straight line through five or more frames, which noise does not
# supply.
THRESH = 3.5
MINAREA = 5


def frames():
    """(key, filename, filter, mid-exposure MJD) for every calibrated sub."""
    rows = []
    counts = {}
    for fn in sorted(os.listdir(SRC)):
        if not (fn.startswith("calibrated-") and fn.endswith(".fit")):
            continue
        filt = fn.split("-")[6]
        counts[filt] = counts.get(filt, 0) + 1
        rows.append((f"{filt}_{counts[filt]:03d}", fn, filt))
    return rows


def detect(image, thresh=THRESH, minarea=MINAREA):
    bkg = sep.Background(image, bw=64, bh=64, fw=3, fh=3)
    sub = image - bkg.back()
    rms = bkg.globalrms
    objs = sep.extract(sub, thresh, err=rms, minarea=minarea,
                       deblend_cont=0.005, filter_kernel=None)
    return objs, sub, rms


def main():
    rows = frames()
    print(f"{len(rows)} calibrated subs")

    # Reference star list first: everything else is matched onto it.
    ref_fn = [r[1] for r in rows if r[0] == REFERENCE][0]
    with fits.open(layout.path(ref_fn), memmap=False) as hd:
        img = hd[0].data.astype(np.float32)
    objs, _, rms = detect(img)
    bright = objs[(objs["flag"] == 0) & (objs["npix"] > 12) &
                  (objs["npix"] < 3000)]
    bright = bright[np.argsort(bright["flux"])[::-1][:600]]
    ref_xy = np.column_stack([bright["x"], bright["y"]])
    del img
    print(f"reference {REFERENCE}: {len(ref_xy)} registration stars, "
          f"rms {rms:.2f}")

    store = {}
    meta = []
    for key, fn, filt in rows:
        path = layout.path(fn)
        with fits.open(path, memmap=False) as hd:
            img = hd[0].data.astype(np.float32)
            mjd = float(hd[0].header["JD"]) - 2400000.5
        mjd_mid = mjd + 150.0 / 86400.0          # mid-exposure
        objs, _, rms = detect(img)
        del img

        clean = objs[(objs["flag"] == 0) & (objs["npix"] > 12) &
                     (objs["npix"] < 3000)]
        clean = clean[np.argsort(clean["flux"])[::-1][:600]]
        xy = np.column_stack([clean["x"], clean["y"]])

        if key == REFERENCE:
            params = (1.0, 0.0, 0.0, 0.0)
            nmatch = len(xy)
            tx, ty = objs["x"].copy(), objs["y"].copy()
        else:
            try:
                tform, (src_m, _) = aa.find_transform(xy, ref_xy)
            except Exception as exc:                            # noqa: BLE001
                print(f"  {key:16s} REGISTRATION FAILED ({exc}) - dropped")
                continue
            nmatch = len(src_m)
            pts = np.column_stack([objs["x"], objs["y"]])
            warped = tform(pts)
            tx, ty = warped[:, 0], warped[:, 1]
            params = (tform.scale, np.degrees(tform.rotation),
                      tform.translation[0], tform.translation[1])

        # Everything the tracklet linker and the vetting need, per detection.
        rec = np.column_stack([
            tx, ty,                                   # reference-frame x, y
            objs["x"], objs["y"],                     # native x, y
            objs["flux"], objs["peak"],
            objs["a"], objs["b"], objs["theta"],
            objs["npix"].astype(float), objs["flag"].astype(float),
        ]).astype(np.float32)
        store[key] = rec
        meta.append((key, fn, filt, f"{mjd_mid:.8f}", f"{rms:.4f}",
                     str(len(rec)), str(nmatch),
                     *[f"{p:.6f}" for p in params]))
        print(f"  {key:16s} {len(rec):5d} det  match={nmatch:3d}  "
              f"rms={rms:6.2f}  rot={params[1]:+7.3f} deg")

    np.savez_compressed(CACHE, meta=np.array(meta, dtype=object), **store)
    print(f"\nwrote {CACHE}")


if __name__ == "__main__":
    sys.exit(main())