"""Split the finished image into a starless view and a star-only layer.

The first attempt used a greyscale morphological opening, the classical trick,
and it failed in a way worth recording: an opening replaces each pixel with a
local minimum, so on a noisy sky it does not merely delete the star, it digs a
hole slightly BELOW sky level and leaves a black dot behind. It also left the
brightest stars standing, because their wings are wider than any structuring
element small enough to spare the galaxy.

What works instead is catalogue-and-fill. The stars to remove are taken from
Gaia DR3 rather than from a blind detection: Gaia is essentially complete for
stars down to G = 20, which is fainter than anything visible here, so the
catalogue identifies the foreground almost perfectly. That matters because a
blind detector cannot tell a Milky Way star from one of Centaurus A's own
globular clusters or from the blue knots in its dust lane, and removing those
guts the very structure the image is about. Each Gaia star is painted into a
mask whose radius follows its magnitude, and the masked pixels are
replaced by a normalised convolution - a Gaussian blur of the unmasked pixels
divided by the same blur of the mask itself, which interpolates across each
star using only real neighbouring pixels. That follows the galaxy's gradient
instead of flattening it, and leaves no holes.

The star layer is then simply what was removed. Two honest limitations: the
extended halos and diffraction spikes of the very brightest stars reach beyond
any sensible mask radius and leave soft residual glows behind, and anything
Gaia does not list (background galaxies, the cluster system) stays in the
starless frame by design. This is a presentation tool, not a measurement.
"""
import os

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
from scipy.ndimage import gaussian_filter

import layout

OUT = layout.SESSION
SOURCE = "NGC5128-img-deconvolved.png"
FILL_SIGMA = 9.0            # interpolation scale for the normalised convolution

rgb = np.asarray(Image.open(layout.path(SOURCE))).astype(np.float32) / 255
ny, nx = rgb.shape[:2]
print(f"{SOURCE}: {nx} x {ny}")

z = np.load(layout.path("_gaia_deep.npz"))
with fits.open(layout.path("NGC5128-LRGB.fit")) as hd:
    wcs = WCS(hd[0].header, naxis=2)
gx, gy = wcs.world_to_pixel(SkyCoord(z["ra"] * u.deg, z["dec"] * u.deg))
gmag = z["g"]
on_frame = (gx > -30) & (gx < nx + 30) & (gy > -30) & (gy < ny + 30)
gx, gy, gmag = gx[on_frame], gy[on_frame], gmag[on_frame]
print(f"{len(gx)} Gaia stars on the frame, G {gmag.min():.1f}-{gmag.max():.1f}")

mask = np.zeros((ny, nx), np.float32)
R = 30
yy, xx = np.mgrid[-R:R + 1, -R:R + 1]
rr = np.hypot(xx, yy)
for x, y, g in zip(gx, gy, gmag):
    # Radius from magnitude: a bright star's visible disc and wings are far
    # wider than a faint one's, and a fixed radius would either miss the
    # bright ones or scour the field around the faint ones.
    r = float(np.clip(30.0 - 1.9 * (g - 8.0), 4, R - 2))
    cx, cy = int(round(x)), int(round(y))
    x0, x1 = max(0, cx - R), min(nx, cx + R + 1)
    y0, y1 = max(0, cy - R), min(ny, cy + R + 1)
    patch = (rr <= r).astype(np.float32)[(y0 - cy + R):(y1 - cy + R),
                                         (x0 - cx + R):(x1 - cx + R)]
    np.maximum(mask[y0:y1, x0:x1], patch, out=mask[y0:y1, x0:x1])
mask = np.clip(gaussian_filter(mask, 1.5), 0, 1)
print(f"mask covers {mask.mean():.2%} of the frame")

starless = np.empty_like(rgb)
keep = 1.0 - mask
denom = gaussian_filter(keep, FILL_SIGMA)
denom[denom < 1e-3] = 1e-3
for i in range(3):
    filled = gaussian_filter(rgb[:, :, i] * keep, FILL_SIGMA) / denom
    starless[:, :, i] = rgb[:, :, i] * keep + filled * mask
stars = np.clip(rgb - starless, 0.0, 1.0)

Image.fromarray((np.clip(starless, 0, 1) * 255 + 0.5).astype(np.uint8)).save(
    layout.path("NGC5128-img-starless.png"))
# The star layer is faint on its own; a square-root stretch makes it legible
# without adding or removing anything.
Image.fromarray((np.sqrt(stars) * 255 + 0.5).astype(np.uint8)).save(
    layout.path("NGC5128-img-stars.png"))
print("wrote NGC5128-img-starless.png and NGC5128-img-stars.png")