"""Transient search: sources in the luminance master with no catalogue match.

The chain is deliberately conservative, because on this field the sky is very
crowded with things that are real but not transient.

1. Detect on the deep luminance master (60 min) and measure a 5 px aperture
   magnitude on the master's own zero point.
2. Throw away everything that is not solidly detected. A transient claim rests
   on a single night's data, so the bar is SNR > 10, not the SNR 5 the frame
   nominally reaches.
3. Throw away everything that is not point-like. This is done against the
   image's own PSF, measured from Gaia stars in the frame, using sep's
   half-light radius. It removes background galaxies and the outer structure
   of NGC 5128 itself, but note carefully that it does NOT remove Centaurus A's
   globular clusters: at 3.8 Mpc a cluster with a 3 pc half-light radius spans
   ~0.2 arcsec against 2.7 arcsec seeing, so clusters are unresolved here and
   look exactly like stars. Only catalogues can separate them.
4. Cross-match against Gaia DR3 (the cached 0.42 deg, G < 20.5 catalogue).
5. Cross-match whatever is left against SkyMapper DR2 through VizieR, which
   covers this declination and goes deeper than Gaia for non-stellar objects.
6. Demand that survivors are real by requiring an independent detection in at
   least six of the twelve individual luminance subs. A cosmic ray or a
   stacking artefact cannot do that; anything astrophysical will.
7. What is left is examined one by one.

Writes mo-transient-candidates.csv and _mo_transient.npz.
"""
import csv
import os

import numpy as np
import sep
from astropy import units as u
from astropy.coordinates import SkyCoord
from astropy.io import fits
from astropy.wcs import WCS
from scipy.spatial import cKDTree

import mo_common as C

SNR_MIN = 10.0
MATCH_RADIUS = 2.0        # arcsec, Gaia
SM_RADIUS = 3.0           # arcsec, SkyMapper (worse astrometry, wider PSF)
SUB_MIN = 6               # subs a source must independently appear in
EDGE = 40                 # px


def detect_master():
    with fits.open(os.path.join(C.OUT, "master-Luminance.fit")) as hd:
        img = hd[0].data.astype(np.float32)
        hdr = hd[0].header
    bkg = sep.Background(img, bw=64, bh=64, fw=3, fh=3)
    sub = img - bkg.back()
    objs = sep.extract(sub, 3.0, err=bkg.globalrms, minarea=5,
                       deblend_cont=0.005)
    flux, ferr, _ = sep.sum_circle(sub, objs["x"], objs["y"], C.APRAD,
                                   err=bkg.globalrms, gain=1.0)
    r50, _ = sep.flux_radius(sub, objs["x"], objs["y"],
                             6.0 * np.ones(len(objs)), 0.5, normflux=flux,
                             subpix=5)
    return objs, np.asarray(flux), np.asarray(ferr), np.asarray(r50), \
        hdr, sub, bkg.globalrms


def gaia_cached():
    z = np.load(os.path.join(C.OUT, "_gaia_deep.npz"))
    return z["ra"], z["dec"], z["g"]


def skymapper(ra0, dec0, radius_deg):
    """SkyMapper DR2 over the field, cached because it is a slow query."""
    cache = os.path.join(C.OUT, "_skymapper.npz")
    if os.path.exists(cache):
        z = np.load(cache)
        print(f"SkyMapper: {len(z['ra'])} cached sources")
        return z["ra"], z["dec"], z["g"], z["r"], z["cls"]
    from astroquery.vizier import Vizier
    v = Vizier(columns=["RAICRS", "DEICRS", "gPSF", "rPSF", "ClassStar"],
               row_limit=-1)
    tbl = v.query_region(SkyCoord(ra0 * u.deg, dec0 * u.deg),
                         radius=radius_deg * u.deg,
                         catalog="II/358/smss")[0]
    ra = np.asarray(tbl["RAICRS"], dtype=float)
    dec = np.asarray(tbl["DEICRS"], dtype=float)
    g = np.asarray(tbl["gPSF"], dtype=float)
    r = np.asarray(tbl["rPSF"], dtype=float)
    cls = np.asarray(tbl["ClassStar"], dtype=float)
    np.savez_compressed(cache, ra=ra, dec=dec, g=g, r=r, cls=cls)
    print(f"SkyMapper DR2: {len(ra)} sources retrieved")
    return ra, dec, g, r, cls


def sky_match(sc, ra, dec, radius_arcsec):
    """Nearest-neighbour match; returns index and separation in arcsec."""
    cat = SkyCoord(ra * u.deg, dec * u.deg)
    idx, d2d, _ = sc.match_to_catalog_sky(cat)
    return idx, d2d.arcsec, d2d.arcsec < radius_arcsec


def sub_support(xy_ref):
    """How many of the twelve luminance subs independently show each source."""
    _, dets = C.load_dets()
    tforms = C.frame_transforms()
    n = np.zeros(len(xy_ref), dtype=int)
    for key, _ in C.lum_frames():
        d = dets[key]
        p = tforms[key](d[:, 2:4].astype(float))
        t = cKDTree(p)
        dd, _ = t.query(xy_ref, distance_upper_bound=2.5)
        n += np.isfinite(dd)
    return n


def rgb_support(xy_ref, rms_scale=3.0):
    """Peak significance of each position in the R, G and B masters.

    A real object on the sky is in every filter. A luminance-only artefact is
    not. The masters are pixel-aligned by construction, so this is a direct
    look-up rather than another registration.
    """
    out = {}
    for filt in ("Red", "Green", "Blue"):
        with fits.open(os.path.join(C.OUT, f"master-{filt}.fit")) as hd:
            img = hd[0].data.astype(np.float32)
        bkg = sep.Background(img, bw=64, bh=64, fw=3, fh=3)
        s = img - bkg.back()
        fl, _, _ = sep.sum_circle(s, xy_ref[:, 0], xy_ref[:, 1], C.APRAD,
                                  err=bkg.globalrms, gain=1.0)
        noise = bkg.globalrms * np.sqrt(np.pi * C.APRAD ** 2)
        out[filt] = np.asarray(fl) / noise
        del img, s
    return out


def main():
    objs, flux, ferr, r50, hdr, _sub, mrms = detect_master()
    wcs = WCS(hdr)
    ny, nx = 3194, 4788
    print(f"master detections: {len(objs)}")

    snr = np.where(ferr > 0, flux / np.maximum(ferr, 1e-9), 0.0)
    mag = np.where(flux > 0, -2.5 * np.log10(np.maximum(flux, 1e-9)) + C.ZP,
                   np.nan)

    # PSF reference from the frame's own bright stars.
    gra, gdec, gmag = gaia_cached()
    sc_all = wcs.pixel_to_world(objs["x"], objs["y"])
    gidx, gsep, gok = sky_match(sc_all, gra, gdec, MATCH_RADIUS)
    star = gok & (gmag[gidx] > 15) & (gmag[gidx] < 18.5) & (snr > 30)
    r50_star = float(np.median(r50[star]))
    r50_sig = float(np.std(r50[star]))
    print(f"PSF from {star.sum()} Gaia stars: r50 = {r50_star:.2f} +- "
          f"{r50_sig:.2f} px ({r50_star * C.SCALE:.2f} arcsec)")

    inframe = ((objs["x"] > EDGE) & (objs["x"] < nx - EDGE) &
               (objs["y"] > EDGE) & (objs["y"] < ny - EDGE))
    good = inframe & (snr > SNR_MIN) & np.isfinite(mag)
    pointlike = np.abs(r50 - r50_star) < 3.0 * max(r50_sig, 0.25)
    print(f"SNR > {SNR_MIN:.0f} and in frame: {good.sum()}")
    print(f"  ... of which point-like: {(good & pointlike).sum()}")
    print(f"  ... of which Gaia-matched: {(good & pointlike & gok).sum()}")

    cand = good & pointlike & ~gok
    print(f"\nno Gaia DR3 counterpart within {MATCH_RADIUS}\": {cand.sum()}")

    # SkyMapper DR2
    cen = wcs.pixel_to_world(nx / 2, ny / 2)
    sra, sdec, sg, sr, scls = skymapper(cen.ra.deg, cen.dec.deg, 0.42)
    sidx, ssep, sok = sky_match(sc_all, sra, sdec, SM_RADIUS)
    cand2 = cand & ~sok
    print(f"also no SkyMapper DR2 counterpart within {SM_RADIUS}\": "
          f"{cand2.sum()}")

    ci = np.where(cand2)[0]
    xy = np.column_stack([objs["x"][ci], objs["y"][ci]])
    nsub = sub_support(xy)
    rgb = rgb_support(xy)
    print(f"of those, detected in >= {SUB_MIN} individual subs: "
          f"{(nsub >= SUB_MIN).sum()}")

    sc = wcs.pixel_to_world(xy[:, 0], xy[:, 1])
    # Distance from the centre of NGC 5128 - a supernova is expected on or
    # near the galaxy, and this ranks the list accordingly.
    cenA = SkyCoord("13h25m27.6s", "-43d01m09s")
    dgal = sc.separation(cenA).arcmin

    rows = []
    for k, i in enumerate(ci):
        rows.append(dict(
            id=int(i), ra_deg=round(sc[k].ra.deg, 6),
            dec_deg=round(sc[k].dec.deg, 6),
            ra_hms=sc[k].ra.to_string(u.hour, sep=":", precision=2),
            dec_dms=sc[k].dec.to_string(u.deg, sep=":", precision=1,
                                        alwayssign=True),
            x=round(float(xy[k, 0]), 2), y=round(float(xy[k, 1]), 2),
            g_mag=round(float(mag[i]), 2), snr=round(float(snr[i]), 1),
            r50_px=round(float(r50[i]), 2),
            r50_over_psf=round(float(r50[i] / r50_star), 2),
            n_subs=int(nsub[k]),
            snr_R=round(float(rgb["Red"][k]), 1),
            snr_G=round(float(rgb["Green"][k]), 1),
            snr_B=round(float(rgb["Blue"][k]), 1),
            gaia_sep_arcsec=round(float(gsep[i]), 2),
            smss_sep_arcsec=round(float(ssep[i]), 2),
            dist_from_cenA_arcmin=round(float(dgal[k]), 2)))
    rows.sort(key=lambda r: -r["snr"])

    out = os.path.join(C.OUT, "mo-transient-candidates.csv")
    if rows:
        with open(out, "w", newline="") as fh:
            w = csv.DictWriter(fh, fieldnames=list(rows[0].keys()))
            w.writeheader()
            w.writerows(rows)
    print(f"\nwrote {out} ({len(rows)} rows)")
    for r in rows[:40]:
        print(f"  {r['ra_hms']} {r['dec_dms']} G={r['g_mag']:5.2f} "
              f"SNR={r['snr']:6.1f} r50/psf={r['r50_over_psf']:.2f} "
              f"subs={r['n_subs']:2d} RGB=({r['snr_R']:.0f},{r['snr_G']:.0f},"
              f"{r['snr_B']:.0f}) d={r['dist_from_cenA_arcmin']:.1f}'")

    np.savez(os.path.join(C.OUT, "_mo_transient.npz"),
             x=objs["x"], y=objs["y"], mag=mag, snr=snr, r50=r50,
             gaia_ok=gok, gaia_sep=gsep, sm_ok=sok, sm_sep=ssep,
             good=good, pointlike=pointlike, cand=cand, cand2=cand2,
             cand_idx=ci, nsub=nsub, r50_star=r50_star, r50_sig=r50_sig,
             snr_R=rgb["Red"], snr_G=rgb["Green"], snr_B=rgb["Blue"],
             dgal=dgal)
    print("saved _mo_transient.npz")


if __name__ == "__main__":
    main()