"""Pass 5: three further renderings from the same masters.

1. NGC5128-img-deconvolved - the luminance sharpened by Richardson-Lucy against a PSF
   measured from the frame's own stars, then recombined into LRGB. The seeing
   was 2.69 arcsec, so there is real detail to recover in the dust lane; the
   deconvolution is deliberately stopped early and applied only where the
   signal is strong, because RL amplifies noise and rings around bright stars
   if it is let run.
2. NGC5128-img-core-print - a full-resolution crop of the galaxy for printing.

The star/starless separation lives in starless.py, which works from this
script's output.

The stretch, colour calibration and gradient handling are imported from
compose.py rather than re-implemented, so these renderings and the main image
cannot drift apart.
"""
import os
import sys

import numpy as np
import sep
import tifffile
from astropy.io import fits
from PIL import Image
from scipy.ndimage import gaussian_filter, median_filter
from skimage.restoration import richardson_lucy

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import compose as C                                          # noqa: E402

import layout

OUT = C.OUT
RL_ITERS = 12
PSF_BOX = 25


def measure_psf(img):
    """Median-stack cutouts of isolated stars to get the frame's own PSF."""
    bkg = sep.Background(img, bw=64, bh=64, fw=3, fh=3)
    sub = img - bkg.back()
    o = sep.extract(sub, 25.0, err=bkg.globalrms, minarea=9,
                    deblend_cont=0.005)
    o = o[(o["flag"] == 0) & (o["npix"] > 15) & (o["npix"] < 400)]
    ny, nx = img.shape
    h = PSF_BOX // 2
    # Isolated stars only: a neighbour inside the cutout would drag the wings.
    keep = []
    xs, ys = o["x"], o["y"]
    for i in range(len(o)):
        if not (h + 2 < xs[i] < nx - h - 2 and h + 2 < ys[i] < ny - h - 2):
            continue
        d = np.hypot(xs - xs[i], ys - ys[i])
        if np.sort(d)[1] < 3 * PSF_BOX:
            continue
        keep.append(i)
    keep = keep[:120]
    stack = []
    for i in keep:
        cx, cy = int(round(xs[i])), int(round(ys[i]))
        cut = sub[cy - h:cy + h + 1, cx - h:cx + h + 1].astype(np.float64)
        peak = cut.max()
        if peak > 0 and peak < 60000:            # skip anything near saturation
            stack.append(cut / cut.sum())
    psf = np.median(np.stack(stack), axis=0)
    psf[psf < 0] = 0
    psf /= psf.sum()
    print(f"PSF from {len(stack)} isolated stars, "
          f"peak fraction {psf.max():.4f}")
    return psf.astype(np.float32)


def star_mask(img, rms):
    """Feathered mask over stars, used to keep deconvolution off them.

    Richardson-Lucy rings around any source whose profile is steeper than the
    PSF model can account for, and on a star field that means a dark annulus
    round every bright star. Excluding stars from the deconvolution entirely
    is the standard cure: the galaxy is what needed sharpening anyway.
    """
    bkg = sep.Background(img, bw=64, bh=64, fw=3, fh=3)
    sub = img - bkg.back()
    o = sep.extract(sub, 6.0, err=rms, minarea=6, deblend_cont=0.005)
    o = o[o["npix"] < 3000]
    ny, nx = img.shape
    m = np.zeros((ny, nx), np.float32)
    yy, xx = np.mgrid[-20:21, -20:21]
    rr = np.hypot(xx, yy)
    for x, y, npix, flux in zip(o["x"], o["y"], o["npix"], o["flux"]):
        # Bright stars ring over a wider radius than faint ones.
        r = float(np.clip(2.5 * np.sqrt(npix / np.pi) + 3.0, 5, 18))
        cx, cy = int(round(x)), int(round(y))
        x0, x1 = max(0, cx - 20), min(nx, cx + 21)
        y0, y1 = max(0, cy - 20), min(ny, cy + 21)
        patch = (rr <= r).astype(np.float32)[
            (y0 - cy + 20):(y1 - cy + 20), (x0 - cx + 20):(x1 - cx + 20)]
        np.maximum(m[y0:y1, x0:x1], patch, out=m[y0:y1, x0:x1])
    print(f"  star mask over {len(o)} sources, {m.mean():.2%} of the frame")
    return np.clip(gaussian_filter(m, 2.5), 0, 1)


def deconvolve(img, psf):
    """Richardson-Lucy on the bright signal, feathered back into the noise."""
    rms = 1.4826 * np.median(np.abs(img - np.median(img)))
    pedestal = 5.0 * rms
    positive = np.clip(img + pedestal, 1e-3, None).astype(np.float32)
    scale = float(positive.max())
    out = richardson_lucy(positive / scale, psf, num_iter=RL_ITERS,
                          clip=False) * scale - pedestal
    # Two weights multiply together: apply the result only where there is
    # signal to sharpen, and only where there is no star to ring.
    signal = np.clip((img - 3.0 * rms) / (20.0 * rms), 0.0, 1.0)
    w = (signal * (1.0 - star_mask(img, rms))).astype(np.float32)
    return (out * w + img * (1.0 - w)).astype(np.float32)


def build_rgb(lum_lin, rgb_lin):
    """The colour half of compose.main(), reused verbatim in spirit."""
    lum, params = C.autostretch(lum_lin)
    rgb = np.empty_like(rgb_lin)
    for i in range(3):
        ch = rgb_lin[:, :, i]
        sky = np.median(ch)
        mad = 1.4826 * np.median(np.abs(ch - sky))
        black = sky - 2.8 * mad
        white = np.percentile(ch, 99.995)
        rgb[:, :, i] = C.mtf(np.clip((ch - black) / (white - black), 0, 1),
                             params["midtone"])
    return lum, rgb


def main():
    data, hdr = C.load()
    shape = data["Luminance"].shape
    mask = C.galaxy_mask(shape)
    for name in C.CHANNELS:
        data[name], *_ = C.remove_gradient(data[name], mask)
    o, flux = C.star_photometry(data["Luminance"], data)
    good = (flux["Red"] > 0) & (flux["Green"] > 0) & (flux["Blue"] > 0)
    gr = np.median(flux["Green"][good] / flux["Red"][good])
    gb = np.median(flux["Green"][good] / flux["Blue"][good])
    data["Red"] *= gr
    data["Blue"] *= gb

    psf = measure_psf(data["Luminance"])
    print(f"deconvolving luminance, {RL_ITERS} Richardson-Lucy iterations")
    lum_lin = deconvolve(data["Luminance"], psf)
    rgb_lin = np.dstack([data["Red"], data["Green"], data["Blue"]])
    del data

    lum, rgb = build_rgb(lum_lin, rgb_lin)
    del rgb_lin, lum_lin

    sky_med = [float(np.median(rgb[:, :, i][~mask])) for i in range(3)]
    target = float(np.mean(sky_med))
    for i in range(3):
        rgb[:, :, i] = np.clip(rgb[:, :, i] - (sky_med[i] - target), 0.0, 1.0)
    rgb_lum = rgb.mean(axis=2, keepdims=True)
    chroma = rgb - rgb_lum
    for i in range(3):
        chroma[:, :, i] = gaussian_filter(median_filter(chroma[:, :, i], 3),
                                          1.5)
    rgb = np.clip(rgb_lum + chroma * C.SATURATION, 0.0, 1.0)
    del chroma, rgb_lum

    ratio = lum / np.maximum(rgb.mean(axis=2), 1e-5)
    out = np.clip(rgb * ratio[:, :, None], 0.0, 1.0)
    del ratio, rgb

    neutral = 0.5 * (out[:, :, 0] + out[:, :, 2])
    green = out[:, :, 1]
    out[:, :, 1] = np.where(green > neutral, green * 0.15 + neutral * 0.85,
                            green)

    u8 = (out * 255 + 0.5).astype(np.uint8)
    Image.fromarray(u8).save(layout.path("NGC5128-img-deconvolved.png"))
    tifffile.imwrite(layout.path("NGC5128-img-deconvolved.tif"),
                     (out * 65535 + 0.5).astype(np.uint16), photometric="rgb")
    print("wrote NGC5128-img-deconvolved.png / .tif")

    # A crop for print, taken from the deconvolved version at full resolution.
    ny, nx = out.shape[:2]
    cw, ch = 2600, 1950
    crop = u8[ny // 2 - ch // 2:ny // 2 + ch // 2,
              nx // 2 - cw // 2:nx // 2 + cw // 2]
    Image.fromarray(crop).save(layout.path("NGC5128-img-core-print.jpg"),
                               quality=95)
    print(f"wrote NGC5128-img-core-print.jpg ({cw}x{ch}, "
          f"{cw * 0.5376 / 60:.1f}' x {ch * 0.5376 / 60:.1f}')")
    del crop, u8


if __name__ == "__main__":
    main()