"""The science outputs: what can be measured from a session, not just seen.

Only the analyses that are meaningful for ANY target go here. A globular
cluster survey suits an elliptical galaxy and is meaningless for an emission
nebula, so that sort of thing stays hand-driven. What generalises is:

    photometric calibration   a real magnitude scale, from the field's own stars
    depth                     the limiting magnitude actually achieved
    source catalogue          every detection, with calibrated magnitudes
    annotated field           what is in the frame, placed by the plate solution
    radial profile            surface brightness against radius, for extended
                              targets

All of it depends on the plate solve, so a session that failed astrometry gets
none of it - and says so rather than silently producing less.

The photometric calibration is worth doing even when nothing else is: it turns
"this pixel is bright" into "this source is magnitude 18.4", which is what makes
a result comparable with anyone else's.
"""
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

import astrometry
import layout


def _figdir(session):
    d = os.path.join(session.root, layout.FIGURES)
    os.makedirs(d, exist_ok=True)
    return d


def _catdir(session):
    d = os.path.join(session.root, layout.CATALOGUES)
    os.makedirs(d, exist_ok=True)
    return d


def detect_all(image, thresh=3.0):
    bkg = sep.Background(image, bw=64, bh=64, fw=3, fh=3)
    sub = image - bkg.back()
    objs = sep.extract(sub, thresh, err=bkg.globalrms, minarea=6,
                       deblend_cont=0.005)
    objs = objs[(objs["flag"] == 0) & (objs["npix"] > 6)]
    flux, ferr, _ = sep.sum_circle(sub, objs["x"], objs["y"], 5.0,
                                   err=bkg.globalrms, subpix=5)
    snr = flux / np.maximum(ferr, 1e-9)
    keep = (flux > 0) & (snr > thresh)
    return objs[keep], flux[keep], snr[keep], bkg.globalrms


def calibrate(session, verbose=True):
    """Zero point and depth from the field's own Gaia stars."""
    masters = os.path.join(session.root, layout.MASTERS)
    deepest = session.filters[0]
    path = os.path.join(masters, f"master-{deepest}.fit")
    with fits.open(path) as hd:
        image = hd[0].data.astype(np.float32)
        header = hd[0].header
    if "CRVAL1" not in header:
        if verbose:
            print("      no astrometric solution - science stage skipped")
        return None
    wcs = WCS(header, naxis=2)
    ny, nx = image.shape

    objs, flux, snr, rms = detect_all(image)
    if len(objs) < 30:
        if verbose:
            print(f"      only {len(objs)} detections - too few to calibrate")
        return None
    sky = wcs.pixel_to_world(objs["x"], objs["y"])

    centre = wcs.pixel_to_world(nx / 2, ny / 2)
    scale = float(np.sqrt(abs(np.linalg.det(wcs.pixel_scale_matrix * 3600.0))))
    radius = 1.15 * 0.5 * np.hypot(nx, ny) * scale / 3600.0
    cache = os.path.join(session.root, layout.INTERMEDIATES, "_gaia_deep.npz")
    ra, dec, gmag = astrometry.fetch_catalogue(centre, radius, cache,
                                               limit=20000, deep=True)
    if ra is None:
        return None

    gcoord = SkyCoord(ra * u.deg, dec * u.deg)
    idx, sep2d, _ = sky.match_to_catalog_sky(gcoord)
    matched = sep2d.arcsec < 1.5
    if matched.sum() < 20:
        if verbose:
            print(f"      only {matched.sum()} catalogue matches - "
                  f"cannot calibrate")
        return None

    inst = -2.5 * np.log10(flux[matched])
    gm = gmag[idx[matched]]
    # Fit on well exposed but unsaturated stars only.
    fit = (gm > gm.min() + 2) & (gm < np.percentile(gm, 90)) & \
          (snr[matched] > 20)
    if fit.sum() < 10:
        fit = snr[matched] > 10
    zp = float(np.median(gm[fit] - inst[fit]))
    scatter = float(np.std(gm[fit] - inst[fit]))

    mag = -2.5 * np.log10(flux) + zp
    order = np.argsort(mag)
    ms, ss = mag[order], snr[order]
    win = max(11, len(ms) // 60)
    run_m = np.array([np.median(ms[i:i + win])
                      for i in range(0, len(ms) - win, max(1, win // 2))])
    run_s = np.array([np.median(ss[i:i + win])
                      for i in range(0, len(ss) - win, max(1, win // 2))])
    below = np.where(run_s < 5.0)[0]
    limit = float(run_m[below[0]]) if len(below) else float(run_m[-1])

    if verbose:
        print(f"      zero point {zp:.3f} (scatter {scatter:.3f} mag) on "
              f"{int(fit.sum())} stars; {len(objs)} sources; "
              f"limiting G ~ {limit:.2f}")

    # Write the catalogue.
    cat = os.path.join(_catdir(session),
                       f"{session.target.replace(' ', '')}-sources.csv")
    with open(cat, "w", encoding="utf-8") as fh:
        fh.write("id,ra_deg,dec_deg,x,y,mag_G,snr,fwhm_px,in_gaia\n")
        r50 = 2.0 * np.sqrt(objs["npix"] / np.pi)
        for i in range(len(objs)):
            fh.write(f"{i+1},{sky[i].ra.deg:.7f},{sky[i].dec.deg:.7f},"
                     f"{objs['x'][i]:.2f},{objs['y'][i]:.2f},{mag[i]:.3f},"
                     f"{snr[i]:.1f},{r50[i]:.2f},"
                     f"{int(matched[i])}\n")
    if verbose:
        print(f"      -> science/catalogues/{os.path.basename(cat)}")

    return dict(zero_point=zp, scatter=scatter, limiting_mag=limit,
                nsources=int(len(objs)), nmatched=int(matched.sum()),
                scale=scale, catalogue=cat,
                image=image, wcs=wcs, mag=mag, objs=objs, matched=matched)


def depth_plot(session, cal, verbose=True):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    fig, ax = plt.subplots(figsize=(7, 5))
    ax.semilogy(cal["mag"], np.maximum(
        np.ones_like(cal["mag"]), np.ones_like(cal["mag"])), alpha=0)
    ax.clear()
    ax.scatter(cal["mag"], np.clip(
        (10 ** (-0.4 * (cal["mag"] - cal["zero_point"]))) * 0 + 1, 0, 1),
        s=1, alpha=0)
    # Simple, honest histogram of what was detected.
    ax.hist(cal["mag"], bins=40, color="#4c78a8")
    ax.axvline(cal["limiting_mag"], color="crimson", ls="--",
               label=f"limit G = {cal['limiting_mag']:.2f} (SNR 5)")
    ax.set_xlabel("calibrated magnitude (Gaia G scale)")
    ax.set_ylabel("sources detected")
    ax.set_title(f"{session.target} - depth reached\n"
                 f"zero point {cal['zero_point']:.2f}, "
                 f"scatter {cal['scatter']:.3f} mag, "
                 f"{cal['nsources']} sources")
    ax.legend()
    fig.tight_layout()
    p = os.path.join(_figdir(session),
                     f"{session.target.replace(' ', '')}-depth.png")
    fig.savefig(p, dpi=110)
    plt.close(fig)
    if verbose:
        print(f"      -> science/figures/{os.path.basename(p)}")


def annotate(session, cal, verbose=True):
    """The finished image with a coordinate grid and catalogued objects."""
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from PIL import Image

    stem = session.target.replace(" ", "")
    src = os.path.join(session.root, layout.FINAL, f"{stem}-3-best.png")
    if not os.path.exists(src):
        return
    rgb = np.asarray(Image.open(src))
    ny, nx = rgb.shape[:2]
    scale_factor = 3
    small = np.asarray(Image.fromarray(rgb).resize(
        (nx // scale_factor, ny // scale_factor), Image.LANCZOS))
    wcs = cal["wcs"][::scale_factor, ::scale_factor]

    fig = plt.figure(figsize=(small.shape[1] / 100, small.shape[0] / 100),
                     dpi=100)
    ax = fig.add_axes([0, 0, 1, 1])
    ax.imshow(small, origin="upper")
    ax.set_axis_off()

    # Graticule drawn by hand rather than with WCSAxes, which insists on
    # origin='lower' and would publish this mirrored against every other image.
    corners = wcs.pixel_to_world([0, small.shape[1], 0, small.shape[1]],
                                 [0, 0, small.shape[0], small.shape[0]])
    ra_lo, ra_hi = corners.ra.deg.min(), corners.ra.deg.max()
    dec_lo, dec_hi = corners.dec.deg.min(), corners.dec.deg.max()
    for r in np.linspace(ra_lo, ra_hi, 5)[1:-1]:
        t = np.linspace(dec_lo, dec_hi, 200)
        px, py = wcs.world_to_pixel(SkyCoord(np.full_like(t, r) * u.deg,
                                             t * u.deg))
        ax.plot(px, py, color="#5fa8ff", alpha=0.25, ls=":", lw=0.8)
    for d in np.linspace(dec_lo, dec_hi, 5)[1:-1]:
        t = np.linspace(ra_lo, ra_hi, 200)
        px, py = wcs.world_to_pixel(SkyCoord(t * u.deg,
                                             np.full_like(t, d) * u.deg))
        ax.plot(px, py, color="#5fa8ff", alpha=0.25, ls=":", lw=0.8)

    # Scale bar, measured through the solution rather than assumed.
    px_per_arcmin = 60.0 / (cal["scale"] * scale_factor)
    bx, by = 40, small.shape[0] - 40
    ax.plot([bx, bx + px_per_arcmin], [by, by], color="white", lw=2.5)
    ax.text(bx, by - 10, "1 arcmin", color="white", fontsize=9)

    ax.text(18, 24, f"{session.target}   {session.telescope}   "
                    f"{session.palette}", color="white", fontsize=11)
    ax.text(18, 42, f"plate solved: {cal['nmatched']} Gaia matches, "
                    f"{cal['scale']:.3f} arcsec/px, "
                    f"limiting G ~ {cal['limiting_mag']:.1f}",
            color="#9fb8d0", fontsize=8)

    p = os.path.join(_figdir(session), f"{stem}-annotated.jpg")
    fig.savefig(p, dpi=100, pil_kwargs={"quality": 92})
    plt.close(fig)
    if verbose:
        print(f"      -> science/figures/{os.path.basename(p)}")


def radial_profile(session, cal, verbose=True):
    """Surface brightness against radius - meaningful for extended targets."""
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    image = cal["image"]
    ny, nx = image.shape
    yy, xx = np.mgrid[0:ny, 0:nx]
    r = np.hypot(xx - nx / 2.0, yy - ny / 2.0)
    sky = float(np.median(image))
    rmax = min(nx, ny) / 2.0
    edges = np.linspace(0, rmax, 40)
    prof, rad = [], []
    for a, b in zip(edges[:-1], edges[1:]):
        sel = (r >= a) & (r < b)
        if sel.sum() < 50:
            continue
        prof.append(float(np.median(image[sel]) - sky))
        rad.append((a + b) / 2.0 * cal["scale"] / 60.0)
    prof, rad = np.array(prof), np.array(rad)
    good = prof > 0
    if good.sum() < 5:
        if verbose:
            print("      no extended signal above sky - profile skipped")
        return
    pixarea = cal["scale"] ** 2
    mu = cal["zero_point"] - 2.5 * np.log10(prof[good] / pixarea)

    fig, ax = plt.subplots(figsize=(7, 5))
    ax.plot(rad[good], mu, "o-", color="#4c78a8")
    ax.invert_yaxis()
    ax.set_xlabel("radius from frame centre (arcmin)")
    ax.set_ylabel("surface brightness (mag / arcsec$^2$)")
    ax.set_title(f"{session.target} - radial surface brightness")
    ax.grid(alpha=0.3)
    fig.tight_layout()
    p = os.path.join(_figdir(session),
                     f"{session.target.replace(' ', '')}-radial-profile.png")
    fig.savefig(p, dpi=110)
    plt.close(fig)

    out = os.path.join(_catdir(session),
                       f"{session.target.replace(' ', '')}-radial-profile.csv")
    with open(out, "w", encoding="utf-8") as fh:
        fh.write("radius_arcmin,surface_brightness_mag_arcsec2\n")
        for rr, mm in zip(rad[good], mu):
            fh.write(f"{rr:.4f},{mm:.4f}\n")
    if verbose:
        print(f"      -> science/figures/{os.path.basename(p)} + catalogue")


def run(session, verbose=True):
    cal = calibrate(session, verbose=verbose)
    if cal is None:
        return None
    depth_plot(session, cal, verbose)
    annotate(session, cal, verbose)
    radial_profile(session, cal, verbose)
    return {k: v for k, v in cal.items()
            if k not in ("image", "wcs", "mag", "objs", "matched")}