"""Pass 2: register every frame to one reference and combine per filter.

One reference frame is used for ALL four filters, not one per filter, so the
four masters come out pixel-aligned and the colour composite needs no further
registration.

Registration works on the cached star lists rather than the pixels: astroalign
matches asterisms between the two point sets and returns a similarity transform,
which is then applied to the image with a bicubic warp. Matching a few hundred
coordinates is far cheaper than cross-correlating 15 Mpx frames, and it copes
with the field rotation between the east and west sides of the meridian.

Before combining, each frame is sky-subtracted and then rescaled so that its
bright-signal level (the 99.5th percentile, which on this field is set by stars
and the galaxy core rather than by sky) matches the group median. That corrects
for transparency changes without needing per-star photometry. A
sigma-clipped mean then rejects cosmic rays and satellite trails; with only four
frames per colour the clip is deliberately gentle (3 sigma, one iteration) so it
does not start eating real signal.
"""
import os

import astroalign as aa
import numpy as np
from astropy.io import fits
from astropy.stats import sigma_clip
from skimage.transform import warp

import layout

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

# Best-seeing luminance frame, chosen from the pass-1 FWHM table.
REFERENCE = "Luminance_002"
FILTERS = ["Luminance", "Red", "Green", "Blue"]

aa.MIN_MATCHES_FRACTION = 0.6
aa.NUM_NEAREST_NEIGHBORS = 8


def load_cache():
    z = np.load(CACHE, allow_pickle=True)
    meta = {row[0]: row for row in z["meta"]}
    stars = {k: z[k + "_xy"] for k in meta}
    return meta, stars, z


def frame_path(fname):
    return layout.path(fname)


def combine(cube, weights):
    """Weighted sigma-clipped mean along axis 0, done row-block by row-block.

    The full cube is already in memory; the blocking here is only to keep the
    boolean mask and the float64 intermediates that sigma_clip allocates from
    tripling peak usage on a machine with a few GB free.
    """
    n, ny, nx = cube.shape
    out = np.zeros((ny, nx), dtype=np.float32)
    w = np.asarray(weights, dtype=np.float32)[:, None, None]
    step = 256
    for y0 in range(0, ny, step):
        y1 = min(y0 + step, ny)
        block = cube[:, y0:y1, :]
        clipped = sigma_clip(block, sigma=3.0, maxiters=1, axis=0,
                             masked=True, copy=True)
        # Warped frames carry NaN outside their footprint; those pixels must
        # drop out of both the sum and the weight total, exactly like a
        # clipped outlier.
        finite = np.isfinite(block)
        good = (~clipped.mask) & finite
        vals = np.where(finite, block, 0.0)
        wb = np.broadcast_to(w, block.shape) * good
        denom = wb.sum(axis=0)
        denom[denom == 0] = np.nan
        out[y0:y1, :] = np.nansum(vals * wb, axis=0) / denom
        del block, clipped, good, finite, vals, wb, denom
    return np.nan_to_num(out, nan=0.0)


def main():
    meta, stars, _ = load_cache()
    ref_xy = stars[REFERENCE]
    ref_row = meta[REFERENCE]
    print(f"reference {REFERENCE} ({ref_row[1]})")
    os.makedirs(OUT, exist_ok=True)

    report = []
    for filt in FILTERS:
        keys = sorted(k for k in meta if meta[k][2] == filt)
        planes, weights, headers, used = [], [], [], []
        for key in keys:
            row = meta[key]
            fname, rms = row[1], float(row[7])
            with fits.open(frame_path(fname), memmap=False) as hd:
                data = hd[0].data.astype(np.float32)
                hdr = hd[0].header
            if key == REFERENCE:
                reg = data
                nmatch = len(ref_xy)
            else:
                try:
                    tform, (src_m, tgt_m) = aa.find_transform(stars[key],
                                                              ref_xy)
                except Exception as exc:                      # noqa: BLE001
                    print(f"  {key}: registration FAILED ({exc}) - dropped")
                    del data
                    continue
                nmatch = len(src_m)
                reg = warp(data, inverse_map=tform.inverse, order=3,
                           mode="constant", cval=np.nan,
                           preserve_range=True).astype(np.float32)
                del data
            # Sky subtraction uses the frame's own median, which on a field
            # this empty outside the galaxy is a fair estimate of sky level.
            sky = float(np.nanmedian(reg))
            reg -= sky
            planes.append(reg)
            weights.append(1.0 / (rms * rms))
            headers.append(hdr)
            used.append((key, nmatch, sky, rms))
            print(f"  {key:16s} matched={nmatch:3d} sky={sky:8.1f} "
                  f"shift-corrected")

        cube = np.stack(planes, axis=0)
        del planes
        # Photometric scaling: normalise each frame to the stack's own median
        # signal so a frame taken through thin cloud cannot drag the mean down.
        levels = np.array([np.nanpercentile(p, 99.5) for p in cube])
        ref_level = np.nanmedian(levels)
        for i, lv in enumerate(levels):
            if lv > 0:
                cube[i] *= float(ref_level / lv)
        print(f"  {filt}: photometric scale factors "
              f"{np.round(ref_level / levels, 4)}")

        master = combine(cube, weights)
        nframes = cube.shape[0]
        del cube

        hdr = headers[0].copy()
        for k in ("CBLACK", "CWHITE", "PEDESTAL", "HISTORY"):
            hdr.remove(k, ignore_missing=True, remove_all=True)
        hdr["FILTER"] = filt
        hdr["NCOMBINE"] = (nframes, "frames in this master")
        hdr["EXPTOTAL"] = (300.0 * nframes, "[s] total integration")
        hdr["STACKREF"] = (REFERENCE, "registration reference frame")
        hdr["STACKALG"] = ("sigma-clipped weighted mean", "combine method")
        hdr["IMAGETYP"] = "Master Light"
        out = layout.path(f"master-{filt}.fit")
        fits.PrimaryHDU(master.astype(np.float32), hdr).writeto(out,
                                                                overwrite=True)
        print(f"  -> {out}  ({nframes} x 300 s = {nframes * 5:.0f} min)")
        report.append((filt, nframes, out))
        del master

    print("\nmasters written:")
    for filt, n, path in report:
        print(f"  {filt:10s} {n:2d} frames  {path}")


if __name__ == "__main__":
    main()