"""Diagnostic: how good is the coordinate registration, and how many
detections survive removal of the static sky?

astroalign only reports a few dozen matched stars, which is enough to define a
similarity transform but says nothing about how well it holds across a
4788 x 3194 field. This refines each frame's transform by nearest-neighbour
matching every detection against the reference frame's detections and refitting
a full affine, then quotes the residual. It also counts, per frame, how many
detections are left once everything coincident with a master-stack source is
removed - that residual population is the input to the tracklet search.
"""
import os

import numpy as np
from astropy.io import fits
from scipy.spatial import cKDTree
from skimage.transform import AffineTransform

import layout

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


def load():
    z = np.load(CACHE, allow_pickle=True)
    meta = {r[0]: r for r in z["meta"]}
    dets = {k: z[k] for k in meta}
    return meta, dets


def refine(dets, ref_key="Luminance_002", tol=3.0):
    """Refit each frame's transform on all cross-matched detections."""
    ref = dets[ref_key]
    tree = cKDTree(ref[:, :2])
    out = {}
    for key, d in dets.items():
        src = d[:, 2:4].astype(float)          # native coords
        cur = d[:, :2].astype(float)           # astroalign-transformed
        for t in (tol, 1.5):
            dist, idx = tree.query(cur, distance_upper_bound=t)
            ok = np.isfinite(dist)
            if ok.sum() < 50:
                break
            tf = AffineTransform()
            tf.estimate(src[ok], ref[idx[ok], :2].astype(float))
            cur = tf(src)
        dist, idx = tree.query(cur, distance_upper_bound=1.5)
        ok = np.isfinite(dist)
        out[key] = (cur, np.median(dist[ok]), np.percentile(dist[ok], 90),
                    ok.sum())
    return out


def main():
    meta, dets = load()
    lum = [k for k in dets if k.startswith("Luminance")]
    ref = refine(dets)

    print("registration residual against reference frame (px):")
    for k in sorted(dets):
        _, med, p90, n = ref[k]
        print(f"  {k:16s} n={n:5d}  median={med:.3f}  p90={p90:.3f}")

    # Static sky = every source in the deep luminance master.
    with fits.open(layout.path("master-Luminance.fit")) as hd:
        master = hd[0].data.astype(np.float32)
    import sep
    bkg = sep.Background(master, bw=64, bh=64, fw=3, fh=3)
    m = sep.extract(master - bkg.back(), 2.5, err=bkg.globalrms, minarea=4,
                    deblend_cont=0.005)
    print(f"\nmaster detections (static sky): {len(m)}")
    mtree = cKDTree(np.column_stack([m["x"], m["y"]]))

    print("\nresiduals after removing anything within 4 px of a master source:")
    for k in sorted(lum):
        cur = ref[k][0]
        d, _ = mtree.query(cur, distance_upper_bound=4.0)
        left = ~np.isfinite(d)
        print(f"  {k:16s} {left.sum():5d} / {len(cur)}")


if __name__ == "__main__":
    main()