"""
gc-complete-inner.py -- step 3b. Artificial stars concentrated on the INNER field.

The uniform injection run (gc-complete.py) puts only ~2% of its fakes inside
3 arcmin, because that annulus is a tiny fraction of the frame area. That left
the innermost completeness bins under-sampled, and the correction visibly failed
to flatten the foreground-star control inside ~4 arcmin. This run injects only
within r < 7 arcmin so the inner bins are properly measured. The two runs are
merged into _gc_complete_all.npz.
"""
import os, time, gc, numpy as np, sep, warnings
from astropy.io import fits
from astropy.wcs import WCS

import layout
warnings.filterwarnings("ignore")

S = layout.SESSION
BW, FW, THRESH, MINAREA, APR = 16, 3, 3.0, 5, 5.0
ZP_L, PIXSCALE = 27.941, 0.5376
NUC_RA, NUC_DEC = 201.365063, -43.019113
NRUN, NPER = 6, 900
MAGS = (17.5, 20.5)
RMAX = 7.0
HALF = 15
RNG = np.random.default_rng(31415)


def gk(fwhm=5.0):
    sig = fwhm / 2.3548
    n = int(2 * round(3 * sig) + 1)
    r = np.arange(n) - n // 2
    xx, yy = np.meshgrid(r, r)
    k = np.exp(-(xx ** 2 + yy ** 2) / (2 * sig ** 2))
    return (k / k.sum()).astype(np.float32)


def main():
    t0 = time.time()
    hdr = fits.getheader(layout.path("master-Luminance.fit"))
    w = WCS(hdr)
    ny, nx = hdr["NAXIS2"], hdr["NAXIS1"]
    cx, cy = [float(v) for v in w.all_world2pix(NUC_RA, NUC_DEC, 0)]
    cat = dict(np.load(layout.path("_gc_cat.npz"), allow_pickle=True))
    FW_MIN, FW_MAX = float(cat["FW_MIN"]), float(cat["FW_MAX"])
    psf = np.load(layout.path("_gc_psf.npy"))
    apfrac = float(np.load(layout.path("_gc_complete.npz"))["apfrac"])
    sep.set_extract_pixstack(1000000)
    K = gk()
    rmax_px = RMAX * 60 / PIXSCALE
    rm_, rr_, ro_ = [], [], []

    for run in range(NRUN):
        # re-read rather than keeping a pristine copy in memory: only ~3.5 GB
        # is free and each master is 61 MB
        img = np.ascontiguousarray(
            fits.getdata(layout.path('master-Luminance.fit')).astype(np.float32))
        # uniform in area within the circle, clipped to the frame
        th = RNG.uniform(0, 2 * np.pi, NPER * 3)
        rad = rmax_px * np.sqrt(RNG.uniform(0, 1, NPER * 3))
        xs = cx + rad * np.cos(th); ys = cy + rad * np.sin(th)
        ok = ((xs > HALF + 5) & (xs < nx - HALF - 5) &
              (ys > HALF + 5) & (ys < ny - HALF - 5))
        xs, ys = xs[ok][:NPER], ys[ok][:NPER]
        mags = RNG.uniform(*MAGS, len(xs))
        ftot = 10 ** (-0.4 * (mags - ZP_L)) / apfrac
        for xi, yi, f in zip(xs, ys, ftot):
            ix, iy = int(round(xi)), int(round(yi))
            img[iy - HALF:iy + HALF + 1, ix - HALF:ix + HALF + 1] += f * psf

        b = sep.Background(img, bw=BW, bh=BW, fw=FW, fh=FW)
        ds = img - b.back(); rmm = b.rms()
        o = sep.extract(ds, THRESH, err=rmm, minarea=MINAREA, filter_kernel=K,
                        filter_type="matched", deblend_nthresh=32,
                        deblend_cont=0.005, clean=True)
        fl, fe, _ = sep.sum_circle(ds, o["x"], o["y"], APR, err=rmm, subpix=5)
        fl2, _, _ = sep.sum_circle(ds, o["x"], o["y"], 2 * APR, err=rmm, subpix=5)
        rh, _ = sep.flux_radius(ds, o["x"], o["y"], np.full(len(o), 6 * APR), 0.5,
                                normflux=fl2, subpix=5)
        fet = np.sqrt(fe ** 2 + np.clip(fl, 0, None) / (0.2467 * 12.0))
        snr = fl / np.clip(fet, 1e-9, None)
        omag = -2.5 * np.log10(np.clip(fl, 1e-9, None)) + ZP_L
        ofw, oel = 2 * rh, o["a"] / np.clip(o["b"], 1e-6, None)
        sel = ((snr >= 5) & ((o["flag"].astype(int) & 0b1110) == 0)
               & (ofw > FW_MIN) & (ofw < FW_MAX) & (oel < 2.0)
               & (omag > 17.5) & (omag < 20.0) & (fl > 0))
        ox, oy = o["x"][sel], o["y"][sel]
        for xi, yi, mg in zip(xs, ys, mags):
            dd = np.hypot(ox - xi, oy - yi)
            rm_.append(mg)
            rr_.append(np.hypot(xi - cx, yi - cy) * PIXSCALE / 60.0)
            ro_.append(bool(dd.min() < 2.0) if len(dd) else False)
        print("  inner run %d/%d (%.0f s)" % (run + 1, NRUN, time.time() - t0), flush=True)
        del img, b, ds, rmm, o, fl, fe, fl2, rh, fet, snr, omag, ofw, oel, ox, oy
        gc.collect()

    rm_, rr_, ro_ = np.array(rm_), np.array(rr_), np.array(ro_)
    np.savez(layout.path("_gc_complete_inner.npz"), mag=rm_, rad=rr_, ok=ro_)

    u = np.load(layout.path("_gc_complete.npz"))
    np.savez(layout.path("_gc_complete_all.npz"),
             mag=np.concatenate([u["mag"], rm_]),
             rad=np.concatenate([u["rad"], rr_]),
             ok=np.concatenate([u["ok"], ro_]), apfrac=u["apfrac"])
    print("\ninner injections %d, recovered %.1f%%" % (len(ro_), 100 * ro_.mean()))
    print("inner completeness grid (rows mag, cols radius arcmin):")
    rb = [0, 1.5, 3, 4.5, 6, 7]
    mb = np.arange(17.5, 20.51, 0.5)
    print("      " + "".join("%8s" % ("%.1f-%.1f" % (rb[i], rb[i + 1]))
                             for i in range(len(rb) - 1)))
    for j in range(len(mb) - 1):
        row = "%4.1f  " % mb[j]
        for i in range(len(rb) - 1):
            s = ((rm_ >= mb[j]) & (rm_ < mb[j + 1]) & (rr_ >= rb[i]) & (rr_ < rb[i + 1]))
            row += "%8s" % ("%.2f" % ro_[s].mean() if s.sum() > 15 else "-")
        print(row)


if __name__ == "__main__":
    main()