"""Stage 3: register every frame onto one grid and combine per filter.

Registration matches asterisms between the cached star lists rather than
cross-correlating pixels: matching a few hundred coordinates is far cheaper
than comparing 15 megapixel frames, and it copes with the field rotating
between the east and west sides of the meridian.

**One reference for all filters, not one per filter.** That is what makes the
masters pixel-aligned, so the colour composite needs no further registration.

The awkward case this has to survive: NGC 6744 was shot with luminance at bin1
(4096x4096) and colour at bin2 (2048x2048). Frames do not share a shape or a
pixel scale. Because astroalign returns a similarity transform - rotation,
translation AND scale - the maths already handles it; what does not handle it
is assuming the output is the same shape as the input. Every warp is therefore
given the reference's shape explicitly, and the bin2 frames are resampled up
onto the bin1 grid.
"""
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
import measure as measure_mod

aa.MIN_MATCHES_FRACTION = 0.6
aa.NUM_NEAREST_NEIGHBORS = 8


def combine(cube, weights, sigma=3.0):
    """Weighted sigma-clipped mean, in row blocks to bound peak memory.

    sigma_clip allocates a mask and float64 intermediates; on a full cube that
    can triple peak usage, which matters on a machine with a few GB free.
    """
    n, ny, nx = cube.shape
    out = np.zeros((ny, nx), np.float32)
    w = np.asarray(weights, np.float32)[:, None, None]
    step = 256
    for y0 in range(0, ny, step):
        y1 = min(y0 + step, ny)
        block = cube[:, y0:y1, :]
        if n >= 3:
            clipped = sigma_clip(block, sigma=sigma, maxiters=1, axis=0,
                                 masked=True, copy=True)
            good = ~clipped.mask
        else:
            # With one or two frames there is nothing to reject against, and
            # clipping would just throw away signal.
            good = np.ones(block.shape, bool)
        finite = np.isfinite(block)
        good &= 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, finite, good, vals, wb, denom
    return np.nan_to_num(out, nan=0.0)


def run(session, verbose=True):
    """Register and stack every filter; write masters to stacks/masters/."""
    xy, stats = measure_mod.load(session)
    ref = measure_mod.choose_reference(session, stats)
    ref_xy = xy[ref.name]
    ref_shape = ref.shape
    if verbose:
        print(f"      reference {ref.name[-32:]} ({ref.filter}, "
              f"fwhm {stats[ref.name]['fwhm']:.2f} px, "
              f"{ref_shape[1]}x{ref_shape[0]})")

    out_dir = os.path.join(session.root, layout.MASTERS)
    os.makedirs(out_dir, exist_ok=True)
    written = {}

    for filt in session.filters:
        frames = session.by_filter(filt)
        planes, weights, used = [], [], []
        header0 = None
        for frame in frames:
            with fits.open(frame.path, memmap=False) as hd:
                data = hd[0].data.astype(np.float32)
                if header0 is None:
                    header0 = hd[0].header.copy()
            if frame.name == ref.name:
                reg = data
            else:
                src = xy.get(frame.name)
                if src is None or len(src) < 8:
                    print(f"        {frame.name[-28:]}: too few stars - "
                          f"dropped")
                    del data
                    continue
                try:
                    tform, _ = aa.find_transform(src, ref_xy)
                except Exception as exc:                      # noqa: BLE001
                    print(f"        {frame.name[-28:]}: registration failed "
                          f"({type(exc).__name__}) - dropped")
                    del data
                    continue
                # output_shape is the whole point: without it a bin2 frame
                # would be written onto a bin2-sized grid and silently fail to
                # line up with a bin1 reference.
                reg = warp(data, inverse_map=tform.inverse,
                           output_shape=ref_shape, order=3, mode="constant",
                           cval=np.nan, preserve_range=True).astype(np.float32)
                del data
            sky = float(np.nanmedian(reg))
            reg -= sky
            planes.append(reg)
            rms = stats.get(frame.name, {}).get("rms", 1.0) or 1.0
            weights.append(1.0 / (rms * rms))
            used.append(frame)

        if not planes:
            print(f"      {filt}: no frames survived registration - skipped")
            continue

        cube = np.stack(planes, axis=0)
        del planes
        # Normalise for transparency: scale each frame so its bright signal
        # matches the group, so a frame 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)

        master = combine(cube, weights)
        nframes = cube.shape[0]
        total_exp = sum(f.exptime for f in used)
        del cube

        hdr = header0
        for key in ("CBLACK", "CWHITE", "PEDESTAL", "HISTORY"):
            hdr.remove(key, ignore_missing=True, remove_all=True)
        hdr["FILTER"] = filt
        hdr["NCOMBINE"] = (nframes, "frames in this master")
        hdr["EXPTOTAL"] = (total_exp, "[s] total integration")
        hdr["STACKREF"] = (ref.name[:60], "registration reference frame")
        hdr["STACKALG"] = ("sigma-clipped weighted mean" if nframes >= 3
                           else "weighted mean (too few frames to clip)",
                           "combine method")
        hdr["IMAGETYP"] = "Master Light"
        path = os.path.join(out_dir, f"master-{filt}.fit")
        fits.PrimaryHDU(master.astype(np.float32), hdr).writeto(
            path, overwrite=True)
        written[filt] = path
        if verbose:
            print(f"      {filt:10s} {nframes:3d} frames  "
                  f"{total_exp/60:6.1f} min  -> master-{filt}.fit")
        del master
    return written, ref