"""Stage 2: measure every frame and cache its star list.

Two jobs. It produces the quality numbers used to choose a registration
reference and to weight the stack - sky level, noise, star sharpness - and it
caches each frame's detected stars, because both registration and the plate
solve need them and detection costs far more than reading a small array back.

Frames are opened one at a time. A 4096x4096 float32 frame is 67 MB and a
session can hold ninety of them; there is no reason to have more than one in
memory at once.
"""
import os

import numpy as np
import sep

from astropy.io import fits

import layout

MAX_STARS = 500          # cached per frame, brightest first
DETECT_SIGMA = 5.0


def measure_frame(path, max_stars=MAX_STARS):
    """Sky, noise, seeing and a star list for one frame."""
    with fits.open(path, memmap=False) as hd:
        data = hd[0].data.astype(np.float32)

    bkg = sep.Background(data, bw=64, bh=64, fw=3, fh=3)
    sky = float(np.median(bkg.back()))
    rms = float(bkg.globalrms)
    sub = data - bkg.back()

    objs = sep.extract(sub, DETECT_SIGMA, err=rms, minarea=5,
                       deblend_cont=0.005)
    # Saturated cores and cosmic-ray hits make poor registration anchors and
    # poor seeing estimates, so they go first.
    objs = objs[(objs["flag"] == 0) & (objs["npix"] < 3000)]

    # The minimum size has to follow the SAMPLING, not be a fixed number of
    # pixels. A fixed floor of 12 px assumes a well sampled star: at 0.53
    # arcsec/px with 5.9 px seeing that is right, but at 3.5 arcsec/px the
    # stars are undersampled at 1.75 px FWHM and cover only a handful of
    # pixels each, so the same floor discards nearly every real star and keeps
    # blends and galaxies instead. That is what made the wide-field sessions
    # impossible to plate solve: the surviving detections were not the objects
    # the catalogue lists.
    if len(objs) > 10:
        rad0, _ = sep.flux_radius(sub, objs["x"], objs["y"], 6.0 * objs["a"],
                                  0.5, normflux=objs["flux"])
        fwhm0 = float(np.median(rad0) * 2.0)
        min_npix = max(4, int(0.5 * np.pi * (fwhm0 / 2.0) ** 2))
        objs = objs[objs["npix"] >= min_npix]
    if len(objs) == 0:
        del data, sub, bkg
        return dict(sky=sky, rms=rms, fwhm=np.nan, nstars=0,
                    xy=np.zeros((0, 2)), flux=np.zeros(0))

    order = np.argsort(objs["flux"])[::-1][:max_stars]
    objs = objs[order]
    radius, _ = sep.flux_radius(sub, objs["x"], objs["y"], 6.0 * objs["a"],
                                0.5, normflux=objs["flux"])
    fwhm = float(np.median(radius) * 2.0)

    result = dict(
        sky=sky, rms=rms, fwhm=fwhm, nstars=int(len(objs)),
        xy=np.column_stack([objs["x"], objs["y"]]).astype(np.float64),
        flux=objs["flux"].astype(np.float64),
    )
    del data, sub, bkg, objs
    return result


def run(session, force=False, verbose=True):
    """Measure every frame in a session, caching to intermediates/."""
    cache_path = os.path.join(session.root, layout.INTERMEDIATES,
                              "_measure.npz")
    os.makedirs(os.path.dirname(cache_path), exist_ok=True)

    cached = {}
    if os.path.exists(cache_path) and not force:
        z = np.load(cache_path, allow_pickle=True)
        cached = {k: z[k] for k in z.files}

    store, rows = {}, []
    for frame in session.frames:
        key = frame.name
        if f"{key}|xy" in cached and not force:
            xy = cached[f"{key}|xy"]
            stats = cached[f"{key}|stats"].item()
        else:
            stats = measure_frame(frame.path)
            xy = stats.pop("xy")
            flux = stats.pop("flux")
            store[f"{key}|flux"] = flux
        store[f"{key}|xy"] = xy
        store[f"{key}|stats"] = np.array(stats, dtype=object)
        rows.append((frame, stats))
        if verbose:
            print(f"      {frame.filter:10s} {frame.name[-28:]:28s} "
                  f"sky={stats['sky']:8.1f} rms={stats['rms']:7.2f} "
                  f"fwhm={stats['fwhm']:5.2f}px stars={stats['nstars']:4d}")

    # Carry forward anything already cached that was not re-measured.
    for k, v in cached.items():
        store.setdefault(k, v)
    np.savez_compressed(cache_path, **store)
    return rows


def load(session):
    """Star lists and stats for a session, as measured earlier."""
    path = os.path.join(session.root, layout.INTERMEDIATES, "_measure.npz")
    if not os.path.exists(path):
        raise FileNotFoundError(
            f"{path} not found - run the measure stage first")
    z = np.load(path, allow_pickle=True)
    xy = {k.split("|")[0]: z[k] for k in z.files if k.endswith("|xy")}
    stats = {k.split("|")[0]: z[k].item() for k in z.files
             if k.endswith("|stats")}
    return xy, stats


def choose_reference(session, stats):
    """Pick the frame everything else is aligned to.

    The reference sets the output grid, so it should be the sharpest frame of
    the filter with the most signal - which is also usually the filter with the
    finest sampling. Choosing a poor reference costs resolution in every other
    frame, permanently, because registration can only resample onto it.
    """
    best_filter = session.filters[0]
    candidates = [f for f in session.by_filter(best_filter)
                  if stats.get(f.name, {}).get("nstars", 0) >= 8]
    if not candidates:
        candidates = session.by_filter(best_filter)
    ref = min(candidates,
              key=lambda f: stats.get(f.name, {}).get("fwhm", np.inf))
    return ref