"""Annotated version of the finished LRGB: coordinate grid and catalogued objects.

The overlay is driven entirely by the local plate solution, so every label sits
where the astrometry says it should. Objects come from SIMBAD, restricted to a
cone matching the field and to types worth marking (galaxies, clusters, radio
sources), then filtered again to those that actually fall inside the frame.

The annotation is drawn on a downsampled copy: at full 15 Mpx the labels would
be microscopic relative to the image, and nobody views a 4692 px wide frame at
1:1 to read a caption.
"""
import os

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from astropy import units as u
from astropy.coordinates import SkyCoord
from astropy.io import fits
from astropy.wcs import WCS
from PIL import Image

import layout

OUT = layout.SESSION
CACHE = layout.path("_simbad.npz")
SCALE = 3                     # downsample factor for the annotated render

with fits.open(layout.path("NGC5128-LRGB.fit")) as hd:
    hdr = hd[0].header
wcs = WCS(hdr, naxis=2)
rgb = np.asarray(Image.open(layout.path("NGC5128-LRGB.png")))
ny, nx = rgb.shape[:2]
small = np.asarray(Image.fromarray(rgb).resize((nx // SCALE, ny // SCALE),
                                               Image.LANCZOS))
# Slicing a WCS rescales it correctly whether the solution is stored as CD or
# as PC + CDELT, which hand-editing the matrix does not.
wcs_small = wcs[::SCALE, ::SCALE]
centre = wcs.pixel_to_world(nx / 2, ny / 2)
print(f"frame {nx}x{ny} -> render {small.shape[1]}x{small.shape[0]}")


def simbad_objects():
    if os.path.exists(CACHE):
        z = np.load(CACHE, allow_pickle=True)
        return z["name"], z["ra"], z["dec"], z["otype"]
    from astroquery.simbad import Simbad
    sim = Simbad()
    sim.ROW_LIMIT = 2000
    for field in ("otype", "V"):
        try:
            sim.add_votable_fields(field)
        except Exception:                                     # noqa: BLE001
            pass
    tbl = sim.query_region(centre, radius=0.42 * u.deg)
    name = np.array([str(r) for r in tbl[tbl.colnames[0]]])
    coords = SkyCoord(tbl["ra"], tbl["dec"], unit=(u.deg, u.deg))
    otype = np.array([str(t) for t in tbl["otype"]]) if "otype" in \
        tbl.colnames else np.array([""] * len(tbl))
    np.savez_compressed(CACHE, name=name, ra=coords.ra.deg,
                        dec=coords.dec.deg, otype=otype)
    return name, coords.ra.deg, coords.dec.deg, otype


name, ra, dec, otype = simbad_objects()
print(f"{len(name)} SIMBAD entries in the cone")

# SIMBAD returns 1518 rows for this field, the bulk of them anonymous entries
# from Centaurus A cluster and variable-star surveys. Marking those would bury
# the image, and the cluster system is being catalogued separately, so the
# overlay keeps only whole objects: other galaxies, planetary nebulae and
# anything carrying a mainstream catalogue designation.
GALAXY_TYPES = ("G", "GiG", "GiP", "GiC", "AGN", "SyG", "rG", "LSB")
# Confirmed nebulae only. SIMBAD lists 93 "PN?" candidates from a single
# survey of this field; they are unconfirmed, they are not visible at this
# depth, and marking them makes the image unreadable.
NEBULA_TYPES = ("PN", "HII", "SNR")
is_gal = np.isin(otype, GALAXY_TYPES)
is_neb = np.isin(otype, NEBULA_TYPES)
mainstream = np.array([n.startswith(("NGC", "IC ", "ESO", "PGC", "AM ", "SN "))
                       and not n.startswith("SNR")
                       for n in name])
sel = is_gal | is_neb | mainstream
sky = SkyCoord(ra[sel] * u.deg, dec[sel] * u.deg)
x, y = wcs_small.world_to_pixel(sky)
inside = (x > 40) & (x < small.shape[1] - 40) & (y > 40) & \
         (y < small.shape[0] - 40)
labels, kinds = name[sel][inside], otype[sel][inside]
x, y = x[inside], y[inside]
print(f"{len(labels)} catalogued objects inside the frame "
      f"({is_gal[sel][inside].sum()} galaxies)")

# Plain axes, not WCSAxes: WCSAxes insists on origin='lower', which would
# publish this image as a vertical mirror of every other deliverable. Drawing
# the graticule by hand keeps all the outputs in one orientation, and the
# lines still come from the plate solution rather than from assumption.
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()

corners = wcs_small.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()


def draw_line(coord_ra, coord_dec, label, at_ra):
    px, py = wcs_small.world_to_pixel(SkyCoord(coord_ra * u.deg,
                                               coord_dec * u.deg))
    ok = (px > 0) & (px < small.shape[1]) & (py > 0) & (py < small.shape[0])
    if ok.sum() < 2:
        return
    ax.plot(px[ok], py[ok], color="#5fa8ff", alpha=0.30, linestyle=":",
            linewidth=0.9)
    i = np.where(ok)[0][len(np.where(ok)[0]) // 2]
    ax.text(px[i], py[i], label, color="#8fc4ff", fontsize=8,
            family="monospace", rotation=0 if at_ra else 90,
            ha="center", va="center",
            bbox=dict(boxstyle="round,pad=0.12", fc="black", ec="none",
                      alpha=0.45))


RA_STEP = 15.0 / 60.0          # one minute of right ascension, in degrees
DEC_STEP = 10.0 / 60.0         # ten arcminutes
t = np.linspace(dec_lo, dec_hi, 400)
for r in np.arange(np.ceil(ra_lo / RA_STEP) * RA_STEP, ra_hi, RA_STEP):
    c = SkyCoord(r * u.deg, 0 * u.deg)
    draw_line(np.full_like(t, r), t,
              f"{int(c.ra.hms.h):02d}h{int(c.ra.hms.m):02d}m", False)
s_ = np.linspace(ra_lo, ra_hi, 400)
for d in np.arange(np.ceil(dec_lo / DEC_STEP) * DEC_STEP, dec_hi, DEC_STEP):
    dm = abs(d - int(d)) * 60
    draw_line(s_, np.full_like(s_, d), f"{int(d):+03d}d{dm:02.0f}m", True)

for lx, ly, lab, kind in zip(x, y, labels, kinds):
    colour = "#7ee08a" if kind in GALAXY_TYPES else "#ffd166"
    ax.add_patch(plt.Circle((lx, ly), 15, fill=False, color=colour,
                            linewidth=1.2, alpha=0.95))
    ax.text(lx + 19, ly - 11, f"{lab}  [{kind}]", color=colour, fontsize=7.5,
            family="monospace",
            bbox=dict(boxstyle="round,pad=0.12", fc="black", ec="none",
                      alpha=0.45))

# Scale bar: one arcminute, measured through the plate solution rather than
# assumed, plus the physical scale at Centaurus A's distance.
pix_per_arcmin = 60.0 / (0.5376 * SCALE)
bx, by = 60, small.shape[0] - 60
ax.plot([bx, bx + pix_per_arcmin], [by, by], color="white", linewidth=2.5)
ax.text(bx, by - 12, "1' = 1.1 kpc at 3.8 Mpc", color="white", fontsize=9)

# Orientation: north and east taken from the WCS, so a flipped or rotated
# solution cannot silently produce a wrong compass.
cx, cy = small.shape[1] - 120, small.shape[0] - 120
c0 = wcs_small.pixel_to_world(cx, cy)
for dlab, offset in (("N", (0 * u.arcmin, 2 * u.arcmin)),
                     ("E", (2 * u.arcmin, 0 * u.arcmin))):
    p = c0.spherical_offsets_by(*offset)
    px, py = wcs_small.world_to_pixel(p)
    ax.annotate("", xy=(px, py), xytext=(cx, cy),
                arrowprops=dict(color="white", width=1.0, headwidth=6))
    ax.text(px, py, dlab, color="white", fontsize=11, ha="center",
            va="center")

ax.text(20, 26, "NGC 5128 (Centaurus A)   iTelescope T32, Siding Spring   "
                "2026-07-21   L 12x300s  RGB 4x300s each",
        color="white", fontsize=10)
ax.text(20, 44, f"plate solved against Gaia DR3: {hdr.get('ASTRSOLV', '')}",
        color="#9fb8d0", fontsize=8)
ax.set_xlim(0, small.shape[1])
ax.set_ylim(small.shape[0], 0)

path = layout.path("NGC5128-img-annotated.jpg")
fig.savefig(path, dpi=100, pil_kwargs={"quality": 92})
print("wrote", path)