"""Moving-object search, pass 2: strip the static sky and link tracklets.

The static sky is removed twice over. First, anything coincident with a source
in the deep luminance master is dropped. Second, anything that sits at the same
reference-frame position in four or more of the twelve subs is dropped - that
catches stars the master's sigma clipping or deblending missed, and by
construction cannot remove a real mover, which is never in the same place
twice. The price of the second cut is a floor on detectable motion: an object
slower than about 2 px per three exposures looks static and is removed with the
stars. mo_sensitivity.py measures where that floor actually falls.

Third, and on this data much the most important, anything that sits at the same
DETECTOR-frame position in four or more subs is dropped. The subs are dithered
and the field rotates ~0.2 deg through the sequence, so registration - which
holds the sky still - drags anything fixed to the detector across the
registered frame on a perfectly straight, perfectly constant-rate, perfectly
constant-brightness track. Hot pixels are therefore ideal fake asteroids, and
they pass every cut a linker would normally apply. Skipping this one cut turns
a clean null result into 141 confident false detections.

What is left is a few dozen detections per frame: cosmic rays, hot pixels that
survived the bad-pixel map, deblending artefacts around bright stars, noise
peaks at the 3.5 sigma threshold, and - if there is one - a minor planet.

Linking is brute force over detection pairs (mo_common.link). Each pair of
residual detections from two frames separated by at least three exposures
defines a candidate velocity; velocities outside a plausible sky-motion range
are discarded, and the rest are propagated to every frame to see how many other
residuals fall on the predicted track. A candidate must be recovered in at
least five frames, lie on a straight constant-rate line to better than 1.2 px,
and hold its brightness to better than 0.5 mag.

The point of demanding a straight line through many frames is that noise and
cosmic rays are independent between frames: the chance that five unrelated
false detections are collinear in space AND linear in time to sub-pixel
precision is tiny, and is measured directly by rerunning the linker on
time-permuted data (mo_sensitivity.py).

Writes _mo_residuals.npz (the residual detection lists, reused by the
false-alarm control and the figures) and mo-candidates rows for the CSV.
"""
import csv
import os

import numpy as np
from astropy.io import fits
from astropy.wcs import WCS
from scipy.spatial import cKDTree

import mo_common as C


def main():
    times = C.lum_times()
    keys = [k for k, _ in C.lum_frames()]
    print(f"{len(keys)} luminance subs spanning {times[-1] * 60:.1f} min")

    tforms = C.frame_transforms()
    mobjs, map_, mrms, mhdr, mimg = C.master_sources()
    del mimg
    mtree = cKDTree(np.column_stack([mobjs["x"], mobjs["y"]]))
    print(f"static sky: {len(mobjs)} sources in the luminance master")

    pts, nats, mags, shapes = [], [], [], []
    for key, path in C.lum_frames():
        with fits.open(path, memmap=False) as hd:
            img = hd[0].data.astype(np.float32)
        objs, ap, rms = C.detect(img)
        del img
        keep = (objs["npix"] > 3) & (objs["npix"] < 3000) & (ap > 0)
        objs, ap = objs[keep], ap[keep]
        zp, nz = C.frame_zeropoint(objs["x"], objs["y"], ap, mobjs, map_,
                                   tforms[key])
        nat = np.column_stack([objs["x"], objs["y"]]).astype(float)
        nats.append(nat)
        pts.append(tforms[key](nat))
        mags.append(-2.5 * np.log10(ap) + zp)
        shapes.append(np.column_stack([objs["a"], objs["b"],
                                       objs["npix"]]).astype(float))
        print(f"  {key}: {len(objs)} det, rms {rms:.2f}, zp {zp:.3f}")

    mask = C.residual_mask(pts, nats, mtree)
    rpts = [p[m] for p, m in zip(pts, mask)]
    rmags = [g[m] for g, m in zip(mags, mask)]
    rnat = [n[m] for n, m in zip(nats, mask)]
    rshape = [s[m] for s, m in zip(shapes, mask)]
    print("\nresiduals per frame:", " ".join(str(len(p)) for p in rpts))
    print(f"total residual detections: {sum(len(p) for p in rpts)}")

    cands = C.link(times, rpts, rmags, nat=rnat, shape=rshape)
    print(f"\n{len(cands)} tracklet candidates after all cuts")

    wcs = WCS(mhdr)
    tmid = times.mean()
    rows = []
    for c in cands:
        xm = c["cx"][0] + c["cx"][1] * tmid
        ym = c["cy"][0] + c["cy"][1] * tmid
        sky = wcs.pixel_to_world(xm, ym)
        rows.append(dict(
            ra_deg=round(sky.ra.deg, 6), dec_deg=round(sky.dec.deg, 6),
            x_ref=round(xm, 2), y_ref=round(ym, 2),
            rate_arcsec_per_hr=round(c["rate"] * C.SCALE, 2),
            pa_deg=round(c["ang"], 1), g_mag=round(c["mag"], 2),
            g_scatter=round(c["magsig"], 2), n_frames=c["nhit"],
            line_rms_px=round(c["rms"], 2),
            elongation=round(c["a"] / max(c["b"], 1e-6), 2)))
        print("  " + str(rows[-1]))

    np.savez(os.path.join(C.OUT, "_mo_residuals.npz"),
             keys=np.array(keys), times=times,
             **{f"p{i}": p for i, p in enumerate(rpts)},
             **{f"m{i}": g for i, g in enumerate(rmags)},
             **{f"n{i}": n for i, n in enumerate(rnat)})
    np.save(os.path.join(C.OUT, "_mo_tracklets.npy"),
            np.array(cands, dtype=object), allow_pickle=True)

    out = os.path.join(C.OUT, "mo-moving-candidates.csv")
    fields = ["ra_deg", "dec_deg", "x_ref", "y_ref", "rate_arcsec_per_hr",
              "pa_deg", "g_mag", "g_scatter", "n_frames", "line_rms_px",
              "elongation"]
    with open(out, "w", newline="") as fh:
        w = csv.DictWriter(fh, fieldnames=fields)
        w.writeheader()
        w.writerows(rows)
    print(f"\nwrote {out} ({len(rows)} rows)")


if __name__ == "__main__":
    main()