"""Pass 6: re-render with the core recovered.

The first composite blew out the galaxy's centre, and the earlier explanation
for that - a saturated nucleus - was wrong. The surface-photometry analysis
found, and a direct check confirmed, that the bright clipped pixels near the
middle of the frame belong to a foreground star 128 px (69 arcsec) from the
nucleus. The galaxy's own light never comes close to the clip level: inside the
inner 800 x 800 px box there is not one star-free pixel above 30000 ADU.

So the core was lost to the STRETCH, not to the sensor. Setting the white point
at the 99.995th percentile put it at 64283 ADU, a level set by field stars,
while the galaxy peaks around a twentieth of that. The midtone transfer needed
to lift a sky at 7 ADU into visibility then pushed everything above a few
thousand ADU to white.

The fix is the standard high-dynamic-range one: stretch the same data twice and
blend. A faint-biased curve for the sky and halo, a bright-biased curve that
keeps the inner galaxy on the shoulder rather than the ceiling, and a mask that
chooses between them by brightness. No extra data is needed, and none of this
invents anything: both curves are monotonic functions of the same pixels.
"""
import os
import sys

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

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

import layout

OUT = C.OUT
FAINT_TARGET = 0.10        # where the sky sits in the faint-biased curve
CORE_HEADROOM_SCALE = 1.15  # bright-curve white point, as a multiple of the
                            # galaxy's own peak: slightly above it, so the very
                            # centre keeps a little headroom
CORE_MIDTONE = 0.35         # gentle: the core needs tonal separation, not lift


def bright_curve(lum_lin, galaxy_peak, black, white):
    """A second stretch whose white point is the galaxy, not the field stars.

    This is the whole trick. The faint curve normalises against a white point
    of 64283 ADU, set by field stars, so the galaxy's entire tonal range - sky
    at 7 ADU up to a peak near 1944 - is squeezed into the top few percent of
    the curve and comes out as a featureless white blob. Rescaling so that the
    galaxy's own peak IS the white point spreads that same range across the
    full output, and the bulge's smooth gradient and the dust lane silhouetted
    against it become visible. Stars clip in this curve, which does not matter:
    it is only ever used where the faint curve has already run out of room.
    """
    lo = black
    hi = galaxy_peak * CORE_HEADROOM_SCALE
    norm = np.clip((lum_lin - lo) / (hi - lo), 0.0, 1.0)
    print(f"  bright curve: black {lo:.1f} ADU, white {hi:.0f} ADU "
          f"(galaxy peak {galaxy_peak:.0f}), midtone {CORE_MIDTONE}")
    return C.mtf(norm, CORE_MIDTONE)


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)
    data["Red"] *= np.median(flux["Green"][good] / flux["Red"][good])
    data["Blue"] *= np.median(flux["Green"][good] / flux["Blue"][good])

    lum_lin = data["Luminance"]
    rgb_lin = np.dstack([data["Red"], data["Green"], data["Blue"]])
    del data

    # The galaxy's true peak, with stars filtered out. A median filter wide
    # enough to swallow a stellar profile leaves the smooth galaxy alone.
    ny, nx = shape
    h = 500
    core = lum_lin[ny // 2 - h:ny // 2 + h, nx // 2 - h:nx // 2 + h]
    galaxy_peak = float(median_filter(core, size=41).max())
    print(f"galaxy peak (star-free) {galaxy_peak:.0f} ADU vs frame max "
          f"{lum_lin.max():.0f} ADU")

    lum_faint, params = C.autostretch(lum_lin, target=FAINT_TARGET)
    lum_bright = bright_curve(lum_lin, galaxy_peak, params["black"],
                              params["white"])

    # Blend on the FAINT curve's brightness: where it has run out of headroom,
    # hand over to the bright curve. Feathered so the transition is invisible.
    w = np.clip((lum_faint - 0.55) / 0.35, 0.0, 1.0)
    w = gaussian_filter(w.astype(np.float32), 8.0)
    lum = np.clip(lum_faint * (1.0 - w) + lum_bright * w, 0.0, 1.0)
    print(f"  HDR blend covers {float((w > 0.05).mean()):.2%} of the 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 = sky - 2.8 * mad
        white = np.percentile(ch, 99.995)
        norm = np.clip((ch - black) / (white - black), 0, 1)
        cf = C.mtf(norm, params["midtone"])
        # The colour channels get the same two-curve treatment, so the core
        # keeps its colour instead of going white while the luminance holds
        # detail.
        peak_c = float(median_filter(
            ch[ny // 2 - h:ny // 2 + h, nx // 2 - h:nx // 2 + h],
            size=41).max())
        nb = np.clip((ch - black) / (peak_c * CORE_HEADROOM_SCALE - black),
                     0.0, 1.0)
        cb = C.mtf(nb, CORE_MIDTONE)
        rgb[:, :, i] = cf * (1.0 - w) + cb * w
    del rgb_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

    detail = lum - gaussian_filter(lum, 2.0)
    weight = np.clip((lum - FAINT_TARGET) * 4.0, 0.0, 1.0)
    lum = np.clip(lum + 0.35 * detail * weight, 0.0, 1.0)
    del detail, weight

    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)

    frac = float((out.max(axis=2) > 0.995).mean())
    print(f"  pixels at full white: {frac:.3%}")

    Image.fromarray((out * 255 + 0.5).astype(np.uint8)).save(
        layout.path("NGC5128-LRGB-hdr.png"))
    tifffile.imwrite(layout.path("NGC5128-LRGB-hdr.tif"),
                     (out * 65535 + 0.5).astype(np.uint16), photometric="rgb")
    prev = Image.fromarray((out * 255 + 0.5).astype(np.uint8))
    prev.thumbnail((2400, 2400), Image.LANCZOS)
    prev.save(layout.path("NGC5128-LRGB-hdr-preview.jpg"), quality=93)
    print("wrote NGC5128-LRGB-hdr.png / .tif / -preview.jpg")


if __name__ == "__main__":
    main()