#!/usr/bin/env python3
"""Image-level check of the profile-fit response to waste heat that follows the starlight (revision; referee
points 3, 4 and 8).

The limits use AllWISE profile-fit (point-source) fluxes of galaxies that are marginally resolved by WISE.  The
catalogue-level injection of s05_limits.py cannot tell whether profile fitting returns the modelled W3/W1 and
W4/W1 for an extended source.  This step tests exactly that on images:

 1. Empirical effective PSFs (radial profiles) for W1-W4 are built from unsaturated, isolated AllWISE stars on
    unWISE coadd cut-outs (native WISE resolution, 2.75 arcsec pixels), normalised with the stars' AllWISE
    profile-fit magnitudes.  A linear PSF fit with these PSFs is therefore on the AllWISE profile-fit flux scale
    for point sources by construction.
 2. An emulated profile fit (position from W1; amplitude = least squares inside 1.25 FWHM; background = median
    of a 50-70 arcsec annulus) is run on unWISE cut-outs of sample galaxies and compared with the AllWISE
    catalogue fluxes of the same galaxies: this validates the emulator on resolved sources, band by band.
 3. Injection: a de Vaucouleurs source with each galaxy's own SDSS effective radius and axis ratio (waste heat
    that follows the starlight), convolved with the band PSF, is added to the real W1-W4 images and the
    emulated profile fit is repeated.  eta_b = recovered / injected flux in band b.  Two versions of
    "measured band ratio / modelled band ratio" are recorded per galaxy:
      emulator only   eta_b / eta_1
      anchored        (eta_b / eta_1) * (emulated W1 / AllWISE catalogue W1)
    The second uses the galaxy's real AllWISE W1 flux, which for well-resolved galaxies is higher than the
    emulated one (step 2); it is the lower bound if AllWISE treated a host's W3/W4 like a faint,
    background-limited source while returning more of the bright W1 light.  The limits are conservative if the
    ratio is >= 1 - EPS.  The host test (testable / compatible) is then rerun with the galaxy's catalogue
    fluxes plus the image-level measurement of the injected waste heat.
 4. eta_b on a grid of effective radius and axis ratio -> data/psf_response_grid.csv, and the W1 anchor as a
    function of size from the WHOLE sample (AllWISE profile-fit W1 / Legacy Surveys total W1, against the
    emulator's eta_1) -> data/psf_w1_anchor.csv.  Both feed the 'psf_response_*' variants of s05_limits.py.
 5. Aperture check of the profile shape: flux inside 22 arcsec (33 arcsec in W4) over profile-fit flux, relative
    to a point source, against the same quantity for the star-following model.  In W1 this tests whether the
    starlight follows the de Vaucouleurs model; in W3 (bright galaxies only) whether the NATURAL mid-infrared
    emission is more extended than the starlight.

This is an emulation of the AllWISE profile fit on coadds, not a rerun of the AllWISE pipeline (which fits
single exposures, all bands at once, with a noise model and active deblending).  Step 2 quantifies how well the
emulation reproduces the catalogue.

Outputs (data/): psf_profiles.csv, psf_stars.csv, image_injection_galaxies.csv, psf_response_grid.csv,
psf_w1_anchor.csv, image_injection_summary.json.   Raw cut-outs are cached in $MIRQ_RAW/unwise_cut.
"""
import io
import json
import os
import urllib.parse
import urllib.request
from concurrent.futures import ThreadPoolExecutor

import numpy as np
import pandas as pd
from scipy.signal import fftconvolve

import mirq_common as M

VIZ = "https://tapvizier.cds.unistra.fr/TAPVizieR/tap/sync"
CUT = "https://www.legacysurvey.org/viewer/cutout.fits"
CUTDIR = os.path.join(M.RAW, "unwise_cut")
PSFDIR = os.path.join(M.RAW, "psf")
LAYER = {1: "unwise-neo7", 2: "unwise-neo7", 3: "unwise-w3w4", 4: "unwise-w3w4"}    # W1/W2: full-depth coadd
PLANE = {1: 0, 2: 1, 3: 0, 4: 1}
PIX = 2.75                      # arcsec
NPIX = 61                       # cut-out size (168 arcsec)
FINE = 0.25                     # arcsec; PIX / FINE = 11 exactly
FWHM = {1: 6.1, 2: 6.4, 3: 6.5, 4: 12.0}       # Wright et al. (2010)
RFIT = {b: 1.25 * FWHM[b] for b in FWHM}       # AllWISE profile-fit radius
BG_IN, BG_OUT = 50.0, 70.0                     # arcsec, background annulus
SEED = 20261012
# high-latitude fields inside the SDSS footprint used for the PSF stars
FIELDS = [(190.0, 30.0), (215.0, 25.0), (150.0, 35.0), (240.0, 30.0), (130.0, 20.0), (170.0, 10.0)]
STAR_SEL = {      # band: (cone radius deg, ADQL magnitude / quality cut, stars kept per field)
    1: (0.8, "W1mag BETWEEN 9.0 AND 10.5 AND sat1=0 AND ex=0", 8),
    2: (0.8, "W2mag BETWEEN 8.5 AND 10.5 AND sat2=0 AND ex=0", 8),
    3: (3.0, "W3mag BETWEEN 4.5 AND 7.0 AND sat3=0 AND ex<=1 AND W3mag-W4mag<0.4", 8),
    4: (6.0, "W4mag BETWEEN 1.5 AND 5.0 AND sat4=0 AND ex<=1 AND W3mag-W4mag<0.4", 8),
}
THETA_BINS = [0.5, 1.5, 2.5, 3.5, 5.0, 7.0, 40.0]
N_PER_BIN = 45
N_W3_BRIGHT, N_W4_BRIGHT = 60, 40
GRID_RE = [0.5, 0.75, 1.0, 1.5, 2.0, 2.5, 3.0, 4.0, 5.0, 7.0, 10.0, 15.0, 20.0, 30.0]
GRID_BA = [0.3, 0.5, 0.7, 1.0]
HYP = [(0.3, 150.0), (0.3, 300.0), (0.1, 150.0), (0.1, 300.0), (0.3, 100.0), (0.3, 600.0)]
R_AP = {1: 22.0, 2: 22.0, 3: 22.0, 4: 33.0}     # arcsec, aperture of the profile-shape check

_yy, _xx = np.mgrid[0:NPIX, 0:NPIX]
XC = (_xx - (NPIX - 1) / 2.0) * PIX          # arcsec offsets of pixel centres from the cut-out centre
YC = (_yy - (NPIX - 1) / 2.0) * PIX


# ----------------------------------------------------------------------------------------------------------
# downloads
# ----------------------------------------------------------------------------------------------------------
def _is_fits(blob):
    if not blob.startswith(b"SIMPLE"):
        raise RuntimeError("not FITS: %r" % blob[:80])


def cutout_path(ra, dec, layer):
    return os.path.join(CUTDIR, layer, "%.5f_%+.5f.fits" % (ra, dec))


def fetch_cutout(args):
    ra, dec, layer = args
    dest = cutout_path(ra, dec, layer)
    url = CUT + "?" + urllib.parse.urlencode(dict(ra="%.5f" % ra, dec="%.5f" % dec, layer=layer,
                                                  pixscale=PIX, size=NPIX))
    try:
        M.fetch(url, dest, timeout=90, tries=3, min_bytes=2880, check=_is_fits)
        return True
    except Exception as e:  # noqa
        print("  cut-out failed:", ra, dec, layer, str(e)[:80], flush=True)
        return False


def fetch_all(jobs, threads=4):
    jobs = [j for j in dict.fromkeys(jobs) if not os.path.exists(cutout_path(*j))]
    if jobs:
        print("  downloading %d cut-outs" % len(jobs), flush=True)
        with ThreadPoolExecutor(threads) as ex:
            list(ex.map(fetch_cutout, jobs))


def load_cutout(ra, dec, band):
    from astropy.io import fits
    p = cutout_path(ra, dec, LAYER[band])
    if not os.path.exists(p):
        return None
    with fits.open(p) as h:
        d = np.array(h[0].data[PLANE[band]], float)
    if d.shape != (NPIX, NPIX) or not np.isfinite(d).all() or np.all(d == 0):
        return None
    return d


def star_table(network=True):
    """Isolated, unsaturated AllWISE stars per band (VizieR II/328), cached."""
    os.makedirs(PSFDIR, exist_ok=True)
    cols = ("AllWISE,RAJ2000,DEJ2000,W1mag,W2mag,W3mag,W4mag,snr1,snr2,snr3,snr4,nb,na,sat1,sat2,sat3,sat4,"
            "ccf,ex")
    out = []
    for b, (rad, cut, nkeep) in STAR_SEL.items():
        for k, (ra, dec) in enumerate(FIELDS):
            dest = os.path.join(PSFDIR, "stars_w%d_field%d.csv" % (b, k))
            q = ("SELECT %s FROM \"II/328/allwise\" WHERE 1=CONTAINS(POINT('ICRS',RAJ2000,DEJ2000),"
                 "CIRCLE('ICRS',%.3f,%.3f,%.2f)) AND %s AND nb=1 AND na=0 AND ccf='0000'" % (cols, ra, dec, rad, cut))

            def chk(blob):
                if not blob.startswith(b"AllWISE"):
                    raise RuntimeError("bad payload %r" % blob[:100])
            if not os.path.exists(dest):
                if not network:
                    continue
                M.fetch(VIZ + "?" + urllib.parse.urlencode(dict(REQUEST="doQuery", LANG="ADQL", FORMAT="csv", QUERY=q)),
                        dest, timeout=300, tries=3, min_bytes=20, check=chk)
            t = pd.read_csv(dest, dtype={"ccf": str})
            t = t.sort_values("W%dmag" % b).head(nkeep)        # the brightest unsaturated stars of the field
            t["band"], t["field"] = b, k
            out.append(t)
    return pd.concat(out, ignore_index=True)


# ----------------------------------------------------------------------------------------------------------
# photometry primitives
# ----------------------------------------------------------------------------------------------------------
def background(D, x0, y0):
    r = np.hypot(XC - x0, YC - y0)
    v = D[(r >= BG_IN) & (r < BG_OUT)]
    med = np.median(v)
    s = 1.4826 * np.median(np.abs(v - med))
    v = v[np.abs(v - med) < 3 * s] if s > 0 else v
    return float(np.median(v)), float(1.4826 * np.median(np.abs(v - np.median(v))))


class EPSF:
    """Circular effective PSF: fraction of a point source's (AllWISE profile-fit) flux in a pixel whose centre
    lies r arcsec from the source."""

    def __init__(self, r, p):
        self.r, self.p = np.asarray(r, float), np.asarray(p, float)

    def __call__(self, r):
        return np.interp(r, self.r, self.p, right=0.0)

    def fwhm(self):
        half = self.p[0] / 2.0
        i = int(np.argmax(self.p < half))
        return 2 * float(np.interp(half, [self.p[i], self.p[i - 1]], [self.r[i], self.r[i - 1]]))


def pf_amp(D, psf, x0, y0, rfit, bg=None):
    """Linear least-squares PSF amplitude at a fixed position (uniform weights inside rfit)."""
    if bg is None:
        bg, _ = background(D, x0, y0)
    r = np.hypot(XC - x0, YC - y0)
    m = r <= rfit
    P = psf(r[m])
    return float(np.sum((D[m] - bg) * P) / np.sum(P * P)), bg


def aperture(D, x0, y0, rap, bg=None):
    if bg is None:
        bg, _ = background(D, x0, y0)
    r = np.hypot(XC - x0, YC - y0)
    return float(np.sum(D[r <= rap] - bg))


def pf_position(D, psf, rfit, span=5.5):
    """Position that maximises the chi^2 reduction of a single PSF fit (coarse-to-fine grid)."""
    bg, _ = background(D, 0.0, 0.0)
    best = (0.0, 0.0)
    for step, half in ((1.0, span), (0.25, 1.0)):
        cand = np.arange(-half, half + 1e-9, step)
        bx, by, bv = best[0], best[1], -np.inf
        for dx in cand:
            for dy in cand:
                x0, y0 = best[0] + dx, best[1] + dy
                r = np.hypot(XC - x0, YC - y0)
                m = r <= rfit
                P = psf(r[m])
                v = np.sum((D[m] - bg) * P) ** 2 / np.sum(P * P)
                if v > bv:
                    bx, by, bv = x0, y0, v
        best = (bx, by)
    return best


def build_epsf(stars, band, n_iter=3):
    """Radial effective PSF from star cut-outs, normalised by the AllWISE profile-fit flux of each star."""
    ims = []
    for _, s in stars.iterrows():
        D = load_cutout(s.RAJ2000, s.DEJ2000, band)
        if D is None:
            continue
        ims.append((D, 10 ** (-0.4 * (s["W%dmag" % band] - 22.5)), s.AllWISE))
    rb = np.arange(0.0, BG_IN + 0.01, 0.5)
    psf = EPSF(rb, np.exp(-0.5 * (rb / (FWHM[band] / 2.3548)) ** 2))      # start: Gaussian (shape only)
    pos = [(0.0, 0.0)] * len(ims)
    for _ in range(n_iter):
        pos = [pf_position(D, psf, RFIT[band], span=4.0) for D, _, _ in ims]
        rr, vv = [], []
        for (D, fcat, _), (x0, y0) in zip(ims, pos):
            bg, _ = background(D, x0, y0)
            r = np.hypot(XC - x0, YC - y0)
            m = r < BG_IN
            rr.append(r[m])
            vv.append((D[m] - bg) / fcat)
        rr, vv = np.concatenate(rr), np.concatenate(vv)
        edges = np.concatenate([np.arange(0, 12, 0.5), np.arange(12, 24, 1.0), np.arange(24, BG_IN + 0.1, 2.0)])
        cen, val = [], []
        for lo, hi in zip(edges[:-1], edges[1:]):
            m = (rr >= lo) & (rr < hi)
            if m.sum() >= 5:
                cen.append(np.mean(rr[m]))
                val.append(np.median(vv[m]))
        cen, val = np.array(cen), np.array(val)
        # extrapolate to r = 0 with a parabola through the first three bins (the profile is flat-topped there)
        c = np.polyfit(cen[:3] ** 2, val[:3], 1)
        psf = EPSF(np.concatenate([[0.0], cen, [BG_IN]]), np.concatenate([[c[1]], val, [0.0]]))
    # per-star check: PSF-fit flux / catalogue flux (should scatter around 1) and the PSF volume
    ratio = []
    for (D, fcat, name), (x0, y0) in zip(ims, pos):
        a, _ = pf_amp(D, psf, x0, y0, RFIT[band])
        ratio.append(dict(band=band, AllWISE=name, x0=x0, y0=y0, pf_over_cat=a / fcat))
    # renormalise so that the median star has pf/cat = 1 exactly
    med = float(np.median([q["pf_over_cat"] for q in ratio]))
    psf = EPSF(psf.r, psf.p * med)
    for q in ratio:
        q["pf_over_cat"] /= med
    vol = float(np.sum(psf(np.hypot(XC, YC))))
    return psf, pd.DataFrame(ratio), vol


# ----------------------------------------------------------------------------------------------------------
# star-following source model
# ----------------------------------------------------------------------------------------------------------
_fx = (np.arange(NPIX * 11) - (NPIX * 11 - 1) / 2.0) * FINE       # fine grid covering the cut-out; coarse pixel
_FX, _FY = np.meshgrid(_fx, _fx)                                  # centres are fine nodes 5, 16, 27, ...
_kx = np.arange(-BG_IN, BG_IN + 1e-9, FINE)
_KR = np.hypot(*np.meshgrid(_kx, _kx))


def dev_image(re_maj, ba, pa_deg, x0, y0):
    """de Vaucouleurs profile (semi-major effective radius re_maj [arcsec], axis ratio ba), truncated like the
    SDSS model (linear taper to zero between 7 and 8 r_e), unit total flux, on the fine grid centred at (x0, y0)."""
    def prof(x, y):
        c, s = np.cos(np.radians(pa_deg)), np.sin(np.radians(pa_deg))
        xr, yr = x * c + y * s, -x * s + y * c
        q = np.hypot(xr, yr / ba) / re_maj
        v = np.exp(-7.669 * (np.maximum(q, 1e-9) ** 0.25 - 1.0))
        return v * np.clip(8.0 - q, 0.0, 1.0)
    X, Y = _FX - x0, _FY - y0
    img = prof(X, Y)
    # the central cusp: supersample the fine pixels within 2 arcsec (15 x 15)
    core = np.hypot(X, Y) < 2.0
    sub = (np.arange(15) + 0.5) / 15.0 - 0.5
    acc = np.zeros(core.sum())
    for sx in sub:
        for sy in sub:
            acc += prof(X[core] + sx * FINE, Y[core] + sy * FINE)
    img[core] = acc / 225.0
    return img / img.sum()


def convolved_model(re_maj, ba, pa_deg, x0, y0, psf):
    """The star-following source seen through the band PSF, sampled at the coarse pixel centres (unit flux)."""
    x0, y0 = np.round(np.array([x0, y0]) / FINE) * FINE
    g = dev_image(re_maj, ba, pa_deg, x0, y0)
    K = psf(_KR)
    full = fftconvolve(g, K, mode="same")
    return full[5::11, 5::11]


# ----------------------------------------------------------------------------------------------------------
def pick_galaxies(d):
    import s05_limits as L
    P = L.phot(d, "psf")
    cc = d.cc_flags
    good = P["ok"] & (d.nb.values == 1) & (d.n_allwise_6as.values == 1) & (cc == "0000").values
    rng = np.random.default_rng(SEED)
    rows = []
    for lo, hi in zip(THETA_BINS[:-1], THETA_BINS[1:]):
        idx = np.where(good & (d.theta_e.values >= lo) & (d.theta_e.values < hi))[0]
        for i in rng.choice(idx, size=min(N_PER_BIN, len(idx)), replace=False):
            rows.append((i, "theta_e_bin"))
    with np.errstate(invalid="ignore"):
        s3 = np.where(good & (P["F3"] / P["e3"] > 15))[0]
        s4 = np.where(good & (P["F4"] / P["e4"] > 10))[0]
    for i in rng.choice(s3, size=min(N_W3_BRIGHT, len(s3)), replace=False):
        rows.append((i, "w3_bright"))
    for i in rng.choice(s4, size=min(N_W4_BRIGHT, len(s4)), replace=False):
        rows.append((i, "w4_bright"))
    t = pd.DataFrame(rows, columns=["row", "set"]).drop_duplicates("row")
    return t


def to_mjy(counts, band):
    return M.WISE_F0[band] * 10 ** (-0.4 * 22.5) * counts * 1e3


def response_grid(psfs):
    rows = []
    for re_ in GRID_RE:
        for ba in GRID_BA:
            row = dict(devrad=re_, ba=ba)
            for b in (1, 2, 3, 4):
                eta = []
                for (x0, y0) in ((0.0, 0.0), (1.25, 0.75)):          # two pixel phases
                    Mb = convolved_model(re_, ba, 30.0, x0, y0, psfs[b])
                    a, _ = pf_amp(Mb, psfs[b], x0, y0, RFIT[b])
                    eta.append(a)
                row["eta%d" % b] = float(np.mean(eta))
            rows.append(row)
    g = pd.DataFrame(rows)
    g.to_csv(os.path.join(M.DATA, "psf_response_grid.csv"), index=False, float_format="%.5g")
    return g


def eta_of(grid, devrad, ba, band):
    """Interpolate eta_b on the (log devrad, ba) grid."""
    from scipy.interpolate import RegularGridInterpolator
    res = np.array(sorted(grid.devrad.unique()))
    bas = np.array(sorted(grid.ba.unique()))
    tab = grid.pivot(index="devrad", columns="ba", values="eta%d" % band).loc[res, bas].values
    f = RegularGridInterpolator((np.log(res), bas), tab, bounds_error=False, fill_value=None)
    x = np.column_stack([np.log(np.clip(devrad, res[0], res[-1])), np.clip(ba, bas[0], bas[-1])])
    return f(x)


def main(network=True):
    import s05_limits as L
    from s03_model import Model
    os.makedirs(CUTDIR, exist_ok=True)
    d = L.load_master()
    summ = {}

    # ---------------- 1. effective PSFs from stars ----------------
    stars = star_table(network)
    if network:
        fetch_all([(s.RAJ2000, s.DEJ2000, LAYER[int(s.band)]) for _, s in stars.iterrows()])
    psfs, chk, prof_rows = {}, [], []
    for b in (1, 2, 3, 4):
        psfs[b], ratio, vol = build_epsf(stars[stars.band == b], b)
        chk.append(ratio)
        summ["psf_w%d" % b] = dict(n_stars=int(len(ratio)), fwhm_arcsec=psfs[b].fwhm(), volume_within_50as=vol,
                                   star_pf_over_cat_scatter=float(1.4826 * np.median(np.abs(ratio.pf_over_cat - 1))))
        for r_, p_ in zip(psfs[b].r, psfs[b].p):
            prof_rows.append(dict(band=b, r_arcsec=r_, psf_per_pixel=p_))
    pd.DataFrame(prof_rows).to_csv(os.path.join(M.DATA, "psf_profiles.csv"), index=False, float_format="%.5g")
    pd.concat(chk).to_csv(os.path.join(M.DATA, "psf_stars.csv"), index=False, float_format="%.5g")
    print(json.dumps({k: v for k, v in summ.items() if k.startswith("psf_")}, indent=1), flush=True)

    # ---------------- 4. response grid (needs only the PSFs) ----------------
    grid = response_grid(psfs)
    print(grid.to_string(index=False, float_format=lambda v: "%.3f" % v), flush=True)

    # ---------------- 2/3. galaxies: emulated profile fit, injection ----------------
    sel = pick_galaxies(d)
    g = d.loc[sel.row.values].copy()
    g["set"] = sel.set.values
    g["ra"], g["dec"] = g.ra.astype(float), g["dec"].astype(float)
    if network:
        fetch_all([(r.ra, r["dec"], lay) for _, r in g.iterrows() for lay in ("unwise-neo7", "unwise-w3w4")])
    m = Model()
    rng = np.random.default_rng(SEED + 1)
    rows = []
    for _, r in g.iterrows():
        Ds = {b: load_cutout(r.ra, r["dec"], b) for b in (1, 2, 3, 4)}
        if any(v is None for v in Ds.values()):
            continue
        x0, y0 = pf_position(Ds[1], psfs[1], RFIT[1])
        devrad = r.theta_e / np.sqrt(r.ba)
        pa = float(rng.uniform(0, 180))
        row = dict(id=r.id, set=r["set"], z=r.zhelio, theta_e=r.theta_e, devrad=devrad, ba=r.ba, ext_flg=r.ext_flg,
                   line_class=r.line_class, x0=x0, y0=y0, pos_offset=float(np.hypot(x0, y0)), sep_w=r.sep_w)
        emu, Mb = {}, {}
        for b in (1, 2, 3, 4):
            a, bg = pf_amp(Ds[b], psfs[b], x0, y0, RFIT[b])
            emu[b] = to_mjy(a, b)
            row["emu%d" % b], row["cat%d" % b], row["cat_e%d" % b] = emu[b], r["F%d" % b], r["e%d" % b]
            Mb[b] = convolved_model(devrad, r.ba, pa, x0, y0, psfs[b])
            # eta: response of the emulated fit to the star-following source added to the real image
            # (position refitted in W1 on the injected image; amplitude equal to the galaxy's own W1 flux)
            amp_counts = max(emu[1], 1e-3) / to_mjy(1.0, b)
            Dinj = Ds[b] + amp_counts * Mb[b]
            if b == 1:
                xi, yi = pf_position(Dinj, psfs[1], RFIT[1])
                row["x0_inj"], row["y0_inj"] = xi, yi
            a2, _ = pf_amp(Dinj, psfs[b], row["x0_inj"], row["y0_inj"], RFIT[b])
            a1, _ = pf_amp(Ds[b], psfs[b], row["x0_inj"], row["y0_inj"], RFIT[b])
            row["eta%d" % b] = (a2 - a1) / amp_counts
            # profile-shape check: aperture / profile-fit flux relative to a point source, data and model
            star = float(np.sum(psfs[b](np.hypot(XC - x0, YC - y0))[np.hypot(XC - x0, YC - y0) <= R_AP[b]]))
            row["ext%d_obs" % b] = aperture(Ds[b], x0, y0, R_AP[b], bg) / a / star if a > 0 else np.nan
            am, bgm = pf_amp(Mb[b], psfs[b], x0, y0, RFIT[b])
            row["ext%d_model" % b] = aperture(Mb[b], x0, y0, R_AP[b], bgm) / am / star
        anchor = emu[1] / r.F1                 # emulated / catalogue W1 of this galaxy
        row["anchor_w1"] = anchor
        for b in (2, 3, 4):
            row["ratio_emu%d" % b] = row["eta%d" % b] / row["eta1"]
            row["ratio_anch%d" % b] = row["eta%d" % b] / row["eta1"] * anchor
        # host test: catalogue fluxes of the galaxy (as in the main analysis) plus the image-level measurement of
        # the injected waste heat.  The starlight is taken to follow the same profile, so the true stellar W1
        # flux is (emulated W1) / eta_1 and the true waste-heat flux in band b is g R_b times that.
        for f, T in HYP:
            gg = f / (1.0 - f)
            R = {b: float(m.R_of(T, b, r.zhelio)) for b in (1, 2, 3, 4)}
            f1_true = emu[1] / row["eta1"]
            dF = {}
            for b in (1, 2, 3, 4):
                Dinj = Ds[b] + gg * R[b] * f1_true / to_mjy(1.0, b) * Mb[b]
                if b == 1:
                    xi, yi = pf_position(Dinj, psfs[1], RFIT[1])
                a, _ = pf_amp(Dinj, psfs[b], xi, yi, RFIT[b])
                dF[b] = to_mjy(a, b) - emu[b]
            # W1 of the host as the catalogue would list it; the stellar part is recovered as in s05_limits.py
            f1_host = r.F1 + dF[1] * (r.F1 / emu[1])
            f1star = f1_host / (1.0 + gg * R[1])
            test, comp = False, True
            for b in (2, 3, 4):
                S = gg * R[b] * f1star
                e = r["e%d" % b]
                if b in (3, 4):
                    test |= bool((1 - L.EPS) * S >= L.K_S * e)
                comp &= bool(r["F%d" % b] + dF[b] >= (1 - L.EPS) * S - L.K_C * e)
            tag = "f%.1f_T%d" % (f, T)
            row["test_" + tag], row["rec_" + tag] = test, test and comp
            bb = 4 if T <= 200 else 3            # the band that carries the test
            row["ratio_" + tag] = dF[bb] / (gg * R[bb] * f1star)
        rows.append(row)
    t = pd.DataFrame(rows)
    t.to_csv(os.path.join(M.DATA, "image_injection_galaxies.csv"), index=False, float_format="%.5g")
    summ["n_galaxies_measured"] = int(len(t))

    def med(x):
        x = np.asarray(x, float)
        x = x[np.isfinite(x)]
        return [float(np.median(x)), float(np.percentile(x, 16)), float(np.percentile(x, 84)), int(len(x))] if len(x) else None
    # validation of the emulator against the AllWISE catalogue
    val = {}
    for b, snr_min in ((1, 20), (2, 20), (3, 10), (4, 8)):
        ok = t["cat%d" % b] / t["cat_e%d" % b] > snr_min
        val["w%d_all" % b] = med((t["emu%d" % b] / t["cat%d" % b])[ok])
        for lab, mm in (("pointlike_theta_lt_2.5", t.theta_e < 2.5), ("theta_2.5_5", (t.theta_e >= 2.5) & (t.theta_e < 5)),
                        ("theta_ge_5", t.theta_e >= 5)):
            val["w%d_%s" % (b, lab)] = med((t["emu%d" % b] / t["cat%d" % b])[ok & mm])
        if b > 1:      # colours: emulated W_b/W1 over catalogue W_b/W1
            val["colour_w%d_over_w1" % b] = med(((t["emu%d" % b] / t.emu1) / (t["cat%d" % b] / t.cat1))[ok])
    summ["emulated_over_catalogue [median, p16, p84, n]"] = val
    # response to a star-following source, by size
    resp = {}
    tb = t[t.set == "theta_e_bin"]
    for lo, hi in zip(THETA_BINS[:-1], THETA_BINS[1:]):
        mm = (tb.theta_e >= lo) & (tb.theta_e < hi)
        resp["theta_e_%g_%g" % (lo, hi)] = dict(n=int(mm.sum()), eta1=med(tb.eta1[mm])[:3],
                                                 eta3_over_eta1=med((tb.eta3 / tb.eta1)[mm])[:3],
                                                 eta4_over_eta1=med((tb.eta4 / tb.eta1)[mm])[:3],
                                                 eta2_over_eta1=med((tb.eta2 / tb.eta1)[mm])[:3])
    summ["response_by_size [median, p16, p84]"] = resp
    for lo, hi in zip(THETA_BINS[:-1], THETA_BINS[1:]):
        mm = (tb.theta_e >= lo) & (tb.theta_e < hi)
        resp["theta_e_%g_%g" % (lo, hi)].update(
            anchor_w1=med(tb.anchor_w1[mm])[:3], anchored_w3=med(tb.ratio_anch3[mm])[:3],
            anchored_w4=med(tb.ratio_anch4[mm])[:3], anchored_w2=med(tb.ratio_anch2[mm])[:3],
            anchored_w3_min=float(tb.ratio_anch3[mm].min()), anchored_w4_min=float(tb.ratio_anch4[mm].min()))
    summ["response_all_365_galaxies"] = {
        "emulator_only_w%d [median, p16, p84, n]" % b: med(t["ratio_emu%d" % b]) for b in (2, 3, 4)}
    summ["response_all_365_galaxies"].update({
        "anchored_w%d [median, p16, p84, n]" % b: med(t["ratio_anch%d" % b]) for b in (2, 3, 4)})
    summ["response_all_365_galaxies"].update({
        "min_emulator_only_w3": float(t.ratio_emu3.min()), "min_emulator_only_w4": float(t.ratio_emu4.min()),
        "min_anchored_w3": float(t.ratio_anch3.min()), "min_anchored_w4": float(t.ratio_anch4.min()),
        "p05_anchored_w3": float(np.percentile(t.ratio_anch3, 5)), "p05_anchored_w4": float(np.percentile(t.ratio_anch4, 5)),
        "frac_anchored_w3_below_0.8": float((t.ratio_anch3 < 0.8).mean()),
        "frac_anchored_w4_below_0.8": float((t.ratio_anch4 < 0.8).mean()),
        "frac_anchored_w3_below_0.7": float((t.ratio_anch3 < 0.7).mean())})
    # profile-shape check
    shape = {}
    hi1 = t.cat1 / t.cat_e1 > 20
    shape["w1_obs_over_model [median, p16, p84, n]"] = med((t.ext1_obs / t.ext1_model)[hi1])
    for lab, mm in (("theta_lt_2.5", t.theta_e < 2.5), ("theta_2.5_5", (t.theta_e >= 2.5) & (t.theta_e < 5)),
                    ("theta_ge_5", t.theta_e >= 5)):
        shape["w1_obs_over_model_" + lab] = med((t.ext1_obs / t.ext1_model)[hi1 & mm])
    for b, snr in ((3, 15), (4, 10)):
        hib = t["cat%d" % b] / t["cat_e%d" % b] > snr
        shape["w%d_obs_over_model_bright [median, p16, p84, n]" % b] = med((t["ext%d_obs" % b] / t["ext%d_model" % b])[hib])
        shape["w%d_obs_extent_bright (1 = point source)" % b] = med(t["ext%d_obs" % b][hib])
        shape["w%d_model_extent_bright" % b] = med(t["ext%d_model" % b][hib])
        shape["w%d_frac_more_extended_than_stars" % b] = float(((t["ext%d_obs" % b] / t["ext%d_model" % b])[hib] > 1.1).mean())
    summ["profile_shape_check"] = shape
    rec = {}
    for f, T in HYP:
        tag = "f%.1f_T%d" % (f, T)
        nt = int(t["test_" + tag].sum())
        rr = t["ratio_" + tag][t["test_" + tag]]
        rec[tag] = dict(n_testable=nt, n_recovered=int(t["rec_" + tag].sum()),
                        recovery=float(t["rec_" + tag].sum() / nt) if nt else None,
                        measured_over_modelled_signal=med(rr), min_ratio=float(rr.min()) if nt else None,
                        n_ratio_below_0p8=int((rr < 0.8).sum()))
    summ["image_level_recovery"] = rec

    # ---------------- sample-wide check: observed W1 loss of profile fitting vs the response model ----------------
    P = L.phot(d, "psf")
    devrad_all = d.theta_e.values / np.sqrt(d.ba.values)
    eta1_all = eta_of(grid, devrad_all, d.ba.values, 1)
    with np.errstate(invalid="ignore", divide="ignore"):
        okl = P["ok"] & (d.sep_ls.values <= 1.5) & (d.ls_type.values != "DUP") & (d.L1.values > 0) \
            & (d.nb.values == 1) & (d.n_allwise_6as.values == 1)
        obs = P["F1"] / d.L1.values
    chk_rows = []
    for lo, hi in zip(THETA_BINS[:-1], THETA_BINS[1:]):
        mm = okl & (d.theta_e.values >= lo) & (d.theta_e.values < hi)
        chk_rows.append(dict(theta_lo=lo, theta_hi=hi, n=int(mm.sum()), observed_F1pf_over_L1total=float(np.median(obs[mm])),
                             predicted_eta1=float(np.median(eta1_all[mm]))))
    summ["w1_loss_check"] = chk_rows
    # W1 anchor as a function of size for the whole sample: emulator eta_1 / observed profile-fit fraction
    anc = pd.DataFrame(chk_rows)
    anc["theta_mid"] = [float(np.median(d.theta_e.values[okl & (d.theta_e.values >= lo) & (d.theta_e.values < hi)]))
                        for lo, hi in zip(THETA_BINS[:-1], THETA_BINS[1:])]
    anc["anchor"] = anc.predicted_eta1 / anc.observed_F1pf_over_L1total
    anc.to_csv(os.path.join(M.DATA, "psf_w1_anchor.csv"), index=False, float_format="%.5g")
    hb = {b: eta_of(grid, devrad_all, d.ba.values, b) / eta1_all for b in (2, 3, 4)}
    a_all = np.interp(np.log(d.theta_e.values), np.log(anc.theta_mid.values), anc.anchor.values)
    summ["sample_response_factors"] = {"emulator_only_w%d [p1, p16, p50, p84, p99]" % b:
                                       [float(v) for v in np.percentile(hb[b][P["ok"]], [1, 16, 50, 84, 99])] for b in (2, 3, 4)}
    summ["sample_response_factors"].update({"anchored_w%d [p1, p16, p50, p84, p99]" % b:
                                            [float(v) for v in np.percentile((hb[b] * a_all)[P["ok"]], [1, 16, 50, 84, 99])]
                                            for b in (2, 3, 4)})
    summ["sample_response_factors"]["frac_anchored_w3_below_0.8"] = float(((hb[3] * a_all)[P["ok"]] < 0.8).mean())
    summ["sample_response_factors"]["frac_anchored_w4_below_0.8"] = float(((hb[4] * a_all)[P["ok"]] < 0.8).mean())
    json.dump(summ, open(os.path.join(M.DATA, "image_injection_summary.json"), "w"), indent=1)
    print(json.dumps(summ, indent=1))


if __name__ == "__main__":
    main()