"""The three renderings, from minimal to everything we learned.

Same 24 subs, same alignment, same field, same crop. The ONLY variable is how
much is done to the pixels, so the three are directly comparable.

  NGC5128-final-1-stacked.png   stack + a standard stretch, nothing else
  NGC5128-final-2-processed.png conventional processing: gradient, colour, denoise
  NGC5128-final-3-best.png      the above, corrected by what the analyses established

What image 3 does differently, and why each change is justified by a measured
result rather than by taste:

1. **Background fit that does not eat the halo.** The surface photometry
   measured the far field sitting at -17.9 ADU/px instead of zero: the plane
   fitted for image 2 absorbed real halo light, because Centaurus A's halo
   fills this field and there is no genuinely empty corner to fit to. Image 3
   fits `channel = a * galaxy_model + plane` simultaneously, using the isophote
   model from the surface photometry, so the plane can only take the part that
   is actually a gradient. The halo survives.

2. **Photometric colour calibration.** Image 2 assumed the average field star
   is grey. Image 3 uses Gaia BP-RP to pick the 1313 stars that are genuinely
   solar-coloured and neutralises on those alone, which does not care what mix
   of spectral types this particular field contains.

3. **The core, recovered.** The nucleus was never saturated - it peaks at 1944
   ADU against a ~63000 clip. A second tone curve scaled to the galaxy rather
   than to field stars restores the bulge gradient the single curve flattened.

4. **Deconvolution that does not ring.** PSF measured from the frame's own
   isolated stars, applied only where there is signal and never on a star.

5. **Noise reduction that cannot eat clusters.** The globular cluster survey
   showed this field contains 289 cluster candidates that look exactly like
   faint stars. Smoothing is therefore driven by a Gaia star mask plus a
   signal mask, so every compact source - foreground star or cluster - is
   excluded from it.
"""
import os
import sys

import numpy as np
import sep
import tifffile
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, median_filter
from skimage.restoration import denoise_tv_chambolle, richardson_lucy

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

import layout

SRC = layout.SESSION
OUT = layout.SESSION
# The finished renderings live in their own directory: they are the
# deliverables, and keeping them apart from the masters, the
# intermediates and the analysis figures makes it obvious which
# files are meant to be looked at.
FINAL = layout.path("final")
ORIG = layout.path("original")
CROP = 48
FILTERS = ["Luminance", "Red", "Green", "Blue"]


def save(arr, stem, title):
    os.makedirs(FINAL, exist_ok=True)
    u8 = (np.clip(arr, 0, 1) * 255 + 0.5).astype(np.uint8)
    Image.fromarray(u8).save(layout.path(f"{stem}.png"))
    tifffile.imwrite(layout.path(f"{stem}.tif"),
                     (np.clip(arr, 0, 1) * 65535 + 0.5).astype(np.uint16),
                     photometric="rgb")
    prev = Image.fromarray(u8)
    prev.thumbnail((2400, 2400), Image.LANCZOS)
    prev.save(layout.path(f"{stem}-preview.jpg"), quality=93)
    print(f"  wrote {stem}.png / .tif / -preview.jpg   [{title}]")


def lrgb(lum, rgb):
    """Take hue from the colour channels, brightness from the luminance."""
    ratio = lum / np.maximum(rgb.mean(axis=2), 1e-5)
    return np.clip(rgb * ratio[:, :, None], 0.0, 1.0)


# --------------------------------------------------------------- image 1
def image1():
    """Stack plus a standard stretch. No corrections of any kind.

    Uses the plain-mean alignment-only stack, so there is not even outlier
    rejection: satellite trails and the moon's gradient are all present. The
    stretch is the conventional one - a black point just below each channel's
    own sky, a white point at its 99.995th percentile, and a single shared
    midtone taken from the luminance so no colour balancing sneaks in through
    the tone curve.
    """
    print("image 1: stack + standard stretch")
    data = {}
    for f in FILTERS:
        d = fits.getdata(layout.path(f"original-{f}.fit"))
        data[f] = np.array(d[CROP:-CROP, CROP:-CROP], dtype=np.float32)
    lum_lin = data["Luminance"]
    lum, params = C.autostretch(lum_lin)
    print(f"  shared midtone {params['midtone']:.4f}")
    rgb = np.empty(lum.shape + (3,), np.float32)
    for i, f in enumerate(("Red", "Green", "Blue")):
        ch = data[f]
        sky = np.median(ch)
        mad = 1.4826 * np.median(np.abs(ch - sky))
        black, white = sky - 2.8 * mad, np.percentile(ch, 99.995)
        rgb[:, :, i] = C.mtf(np.clip((ch - black) / (white - black), 0, 1),
                             params["midtone"])
        print(f"  {f:6s} black {black:7.1f} white {white:8.1f} ADU")
    del data
    out = lrgb(lum, rgb)
    save(out, "NGC5128-final-1-stacked", "align + mean + stretch")
    return out.shape


# --------------------------------------------------------------- image 3
def solar_white_balance(lum, channels, wcs, shape):
    """Neutralise on stars that are genuinely solar-coloured, per Gaia BP-RP."""
    z = np.load(layout.path("_gaia_colours.npz"))
    solar = np.abs(z["bprp"] - 0.82) < 0.15
    sky = SkyCoord(z["ra"][solar] * u.deg, z["dec"][solar] * u.deg)
    x, y = wcs.world_to_pixel(sky)
    g = z["g"][solar]
    ny, nx = shape
    # Keep them away from the edges, off the galaxy's bright core, and out of
    # saturation; faint ones carry too little signal in 20 minutes of colour.
    ok = ((x > 40) & (x < nx - 40) & (y > 40) & (y < ny - 40) &
          (g > 11.5) & (g < 16.5))
    x, y = x[ok], y[ok]
    flux = {}
    for f in ("Red", "Green", "Blue"):
        img = np.ascontiguousarray(channels[f])
        fl, _, _ = sep.sum_circle(img, x, y, 6.0, subpix=5)
        flux[f] = fl
    good = (flux["Red"] > 0) & (flux["Green"] > 0) & (flux["Blue"] > 0)
    gr = float(np.median(flux["Green"][good] / flux["Red"][good]))
    gb = float(np.median(flux["Green"][good] / flux["Blue"][good]))
    print(f"  solar-analogue white balance on {int(good.sum())} stars: "
          f"R x {gr:.4f}, B x {gb:.4f}")
    return gr, gb


def fit_background(ch, model, star_mask):
    """Solve ch = a*model + (plane) and subtract ONLY the plane.

    Fitting a plane on its own to a field this full of galaxy makes the plane
    absorb halo light - measured at -17.9 ADU/px in the far field of image 2.
    Including the galaxy model as a free component in the same least-squares
    problem gives the fit something else to attribute that light to.
    """
    ny, nx = ch.shape
    ys, xs = np.mgrid[0:ny:8, 0:nx:8]
    m = model[::8, ::8]
    v = ch[::8, ::8]
    keep = (star_mask[::8, ::8] == 0) & np.isfinite(v)
    A = np.column_stack([m[keep], xs[keep] / nx, ys[keep] / ny,
                         np.ones(keep.sum())])
    coef, *_ = np.linalg.lstsq(A, v[keep], rcond=None)
    for _ in range(3):                       # clip and refit
        pred = A @ coef
        r = v[keep] - pred
        s = 1.4826 * np.median(np.abs(r - np.median(r)))
        m2 = np.abs(r - np.median(r)) < 2.5 * s
        coef, *_ = np.linalg.lstsq(A[m2], v[keep][m2], rcond=None)
    yy, xx = np.mgrid[0:ny, 0:nx]
    plane = (coef[1] * xx / nx + coef[2] * yy / ny + coef[3]).astype(np.float32)
    print(f"    model amplitude {coef[0]:.4f}, plane offset {coef[3]:8.2f} ADU")
    return ch - plane


def image3():
    print("image 3: science-informed")
    model_full = fits.getdata(layout.path("sb-model.fits")).astype(
        np.float32)
    star_full = fits.getdata(layout.path("sb-mask-stars.fits"))
    with fits.open(layout.path("master-Luminance.fit")) as hd:
        wcs_full = WCS(hd[0].header, naxis=2)

    data = {}
    for f in FILTERS:
        d = fits.getdata(layout.path(f"master-{f}.fit"))
        data[f] = np.array(d, dtype=np.float32)
    model = model_full
    smask = star_full
    print("  background fit with the galaxy model as a free component:")
    for f in FILTERS:
        print(f"   {f}")
        data[f] = fit_background(data[f], model, smask)
    del model_full, star_full, model, smask

    shape_full = data["Luminance"].shape
    gr, gb = solar_white_balance(data["Luminance"], data, wcs_full, shape_full)
    data["Red"] *= gr
    data["Blue"] *= gb

    for f in FILTERS:
        data[f] = np.ascontiguousarray(data[f][CROP:-CROP, CROP:-CROP])
    shape = data["Luminance"].shape

    psf = E.measure_psf(data["Luminance"])
    print("  deconvolving luminance (star-protected)")
    lum_lin = E.deconvolve(data["Luminance"], psf)
    rgb_lin = np.dstack([data["Red"], data["Green"], data["Blue"]])
    del data

    # Two tone curves, blended: the faint one for sky and halo, one scaled to
    # the galaxy for the core the single curve flattened.
    ny, nx = shape
    h = 500
    core = lum_lin[ny // 2 - h:ny // 2 + h, nx // 2 - h:nx // 2 + h]
    peak = float(median_filter(core, size=41).max())
    lum_faint, params = C.autostretch(lum_lin)
    hi = peak * 1.15
    lum_bright = C.mtf(np.clip((lum_lin - params["black"]) /
                               (hi - params["black"]), 0, 1), 0.35)
    w = gaussian_filter(np.clip((lum_faint - 0.55) / 0.35, 0, 1).astype(
        np.float32), 8.0)
    lum = np.clip(lum_faint * (1 - w) + lum_bright * w, 0, 1)
    print(f"  galaxy peak {peak:.0f} ADU, HDR blend over "
          f"{float((w > 0.05).mean()):.2%} of frame")

    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, white = sky - 2.8 * mad, np.percentile(ch, 99.995)
        cf = C.mtf(np.clip((ch - black) / (white - black), 0, 1),
                   params["midtone"])
        pc = float(median_filter(
            ch[ny // 2 - h:ny // 2 + h, nx // 2 - h:nx // 2 + h], 41).max())
        cb = C.mtf(np.clip((ch - black) / (pc * 1.15 - black), 0, 1), 0.35)
        rgb[:, :, i] = cf * (1 - w) + cb * w
    del rgb_lin

    # Sky neutralisation, measured where the galaxy model says there is no
    # galaxy rather than outside an arbitrary ellipse.
    mask = C.galaxy_mask(shape)
    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, 1)

    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 * 1.4, 0, 1)
    del chroma, rgb_lum

    detail = lum - gaussian_filter(lum, 2.0)
    protect = np.clip((lum - 0.10) * 4.0, 0.0, 1.0)
    lum = np.clip(lum + 0.35 * detail * protect, 0, 1)

    # Denoise the sky only. The cluster survey found 289 cluster candidates
    # that look like faint stars, so compact sources are excluded from the
    # smoothing along with the galaxy itself.
    z = np.load(layout.path("_gaia_deep.npz"))
    gx, gy = wcs_full.world_to_pixel(SkyCoord(z["ra"] * u.deg,
                                              z["dec"] * u.deg))
    gx, gy = gx - CROP, gy - CROP
    point = np.zeros(shape, np.float32)
    R = 14
    yy, xx = np.mgrid[-R:R + 1, -R:R + 1]
    rr = np.hypot(xx, yy)
    for x, y, g in zip(gx, gy, z["g"]):
        if not (R < x < nx - R and R < y < ny - R):
            continue
        r = float(np.clip(16.0 - 0.7 * (g - 8.0), 4, R - 1))
        cx, cy = int(x), int(y)
        patch = (rr <= r).astype(np.float32)
        sl = (slice(cy - R, cy + R + 1), slice(cx - R, cx + R + 1))
        np.maximum(point[sl], patch, out=point[sl])
    keep_sharp = np.clip(protect + gaussian_filter(point, 2.0), 0, 1)
    smooth = denoise_tv_chambolle(lum, weight=0.012)
    lum = np.clip(lum * keep_sharp + smooth * (1 - keep_sharp), 0, 1)
    print(f"  sky denoise applied to {float((keep_sharp < 0.5).mean()):.1%} "
          f"of the frame; stars and clusters excluded")
    del detail, protect, smooth, point, keep_sharp

    out = lrgb(lum, rgb)
    neutral = 0.5 * (out[:, :, 0] + out[:, :, 2])
    green = out[:, :, 1]
    out[:, :, 1] = np.where(green > neutral, green * 0.15 + neutral * 0.85,
                            green)
    save(out, "NGC5128-final-3-best", "science-informed")


if __name__ == "__main__":
    image1()
    print("image 2: the conventional pipeline output (compose.py)")
    os.makedirs(FINAL, exist_ok=True)
    for ext in ("png", "tif"):
        src = layout.path(f"NGC5128-LRGB.{ext}")
        dst = layout.path(f"NGC5128-final-2-processed.{ext}")
        with open(src, "rb") as a, open(dst, "wb") as b:
            b.write(a.read())
    im = Image.open(layout.path("NGC5128-final-2-processed.png"))
    im.thumbnail((2400, 2400), Image.LANCZOS)
    im.save(layout.path("NGC5128-final-2-processed-preview.jpg"), quality=93)
    print("  wrote NGC5128-final-2-processed.png / .tif / -preview.jpg")
    image3()