"""Stage 5: assemble a viewable image from whatever filters the session has.

The palette is chosen from what is present, not assumed:

    LRGB   colour from R/G/B, brightness from the deep luminance
    RGB    no luminance, so a synthetic one is made from the colour channels
    SHO    the Hubble palette: SII -> red, Ha -> green, OIII -> blue
    HOO    Ha -> red, OIII -> green and blue
    MONO   one filter, rendered as greyscale

Background handling deserves a note, because the obvious approach is wrong in a
way that is invisible. Fitting a plane to a frame that a large target fills
makes the plane absorb the target's own outer light - on Centaurus A this was
measured at -17.9 ADU/px of real halo quietly subtracted away. So the fit
excludes a central region whose size is derived from where the signal actually
is, and only a plane is used, never a flexible surface.
"""
import os

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

import layout

BG_TARGET = 0.10          # where the sky sits in the stretched image
SATURATION = 1.35

PALETTE_MAP = {
    "SHO": {"Red": "SII", "Green": "Ha", "Blue": "OIII"},
    "HOO": {"Red": "Ha", "Green": "OIII", "Blue": "OIII"},
}


def mtf(x, midtone):
    """Midtone transfer function on data already scaled to [0, 1]."""
    x = np.clip(x, 0.0, 1.0)
    return ((midtone - 1.0) * x) / ((2.0 * midtone - 1.0) * x - midtone)


def autostretch(img, target=BG_TARGET, shadow_sigma=2.8):
    """Black point just below sky, then an MTF putting sky at `target`."""
    sky = float(np.median(img))
    mad = 1.4826 * float(np.median(np.abs(img - sky)))
    if mad <= 0:
        mad = max(float(np.std(img)), 1e-6)
    black = sky - shadow_sigma * mad
    white = float(np.percentile(img, 99.995))
    if white <= black:
        white = black + 1.0
    norm = np.clip((img - black) / (white - black), 0.0, 1.0)
    sky_norm = (sky - black) / (white - black)
    m = ((target - 1.0) * sky_norm) / (2.0 * target * sky_norm - target -
                                       sky_norm)
    m = float(np.clip(m, 1e-4, 0.9))
    return mtf(norm, m), dict(black=black, white=white, midtone=m, sky=sky,
                              mad=mad)


def target_mask(img, frac=0.45):
    """Ellipse covering the bright central object, sized from the data.

    A fixed radius cannot suit both a galaxy filling the frame and a small
    nebula. This grows the ellipse until it contains most of the flux above
    sky, then stops - so the background fit is excluded from wherever the
    target actually is, whatever its size.
    """
    ny, nx = img.shape
    yy, xx = np.mgrid[0:ny, 0:nx]
    cy, cx = ny / 2.0, nx / 2.0
    r = np.hypot((xx - cx) / (nx / 2.0), (yy - cy) / (ny / 2.0))
    sky = float(np.median(img))
    signal = np.clip(img - sky, 0, None)
    total = signal.sum()
    if total <= 0:
        return r < 0.5
    for cut in np.arange(0.25, 0.96, 0.05):
        if signal[r < cut].sum() >= frac * total:
            return r < min(cut + 0.15, 0.95)
    return r < 0.75


def remove_gradient(img, mask):
    """Subtract a plane fitted to tile medians outside `mask`."""
    ny, nx = img.shape
    step = max(48, min(ny, nx) // 40)
    xs, ys, zs = [], [], []
    for y0 in range(0, ny - step, step):
        for x0 in range(0, nx - step, step):
            if mask[y0:y0 + step, x0:x0 + step].any():
                continue
            tile = img[y0:y0 + step, x0:x0 + step]
            zs.append(np.median(tile))
            xs.append(x0 + step / 2.0)
            ys.append(y0 + step / 2.0)
    if len(zs) < 12:
        return img - float(np.median(img)), None
    xs, ys, zs = map(np.asarray, (xs, ys, zs))
    keep = np.ones(len(zs), bool)
    coef = None
    for _ in range(3):
        A = np.column_stack([xs[keep], ys[keep], np.ones(keep.sum())])
        coef, *_ = np.linalg.lstsq(A, zs[keep], rcond=None)
        model = coef[0] * xs + coef[1] * ys + coef[2]
        resid = zs - model
        s = 1.4826 * np.median(np.abs(resid - np.median(resid)))
        if s <= 0:
            break
        keep = np.abs(resid - np.median(resid)) < 2.5 * s
    yy, xx = np.mgrid[0:ny, 0:nx]
    plane = (coef[0] * xx + coef[1] * yy + coef[2]).astype(np.float32)
    return img - plane, coef


def load_masters(session):
    masters = {}
    d = os.path.join(session.root, layout.MASTERS)
    for name in os.listdir(d) if os.path.isdir(d) else []:
        if name.startswith("master-") and name.endswith(".fit"):
            filt = name[len("master-"):-len(".fit")]
            masters[filt] = os.path.join(d, name)
    return masters


def run(session, verbose=True):
    """Produce the final image for whatever palette this session supports."""
    masters = load_masters(session)
    if not masters:
        raise RuntimeError("no masters found - run the register stage first")
    palette = session.palette
    if verbose:
        print(f"      palette {palette} from {sorted(masters)}")

    data, headers = {}, {}
    for filt, path in masters.items():
        with fits.open(path) as hd:
            data[filt] = hd[0].data.astype(np.float32)
            headers[filt] = hd[0].header

    # Background: one mask from the deepest channel, applied to all, so every
    # channel is treated identically and colour is not shifted by the fit.
    deepest = session.filters[0] if session.filters[0] in data \
        else sorted(data)[0]
    mask = target_mask(data[deepest])
    if verbose:
        print(f"      target mask covers {mask.mean():.1%} of the frame")
    for filt in list(data):
        data[filt], coef = remove_gradient(data[filt], mask)

    # Which master feeds which output channel.
    if palette in PALETTE_MAP:
        mapping = PALETTE_MAP[palette]
        channels = [data[mapping[c]] for c in ("Red", "Green", "Blue")]
        lum_src = None
    elif palette in ("LRGB", "RGB", "PARTIAL") and all(
            c in data for c in ("Red", "Green", "Blue")):
        channels = [data["Red"], data["Green"], data["Blue"]]
        lum_src = data.get("Luminance")
    else:
        only = data[deepest]
        stretched, _ = autostretch(only)
        grey = (np.clip(stretched, 0, 1) * 65535).astype(np.uint16)
        return _write(session, np.dstack([grey] * 3), palette, verbose)

    # Colour balance: make the median star colour neutral, which is the cheap
    # substitute for a photometric calibration and is well behaved on a rich
    # field.
    refs = [float(np.percentile(c, 99.5)) for c in channels]
    target = float(np.median(refs))
    channels = [c * (target / r if r > 0 else 1.0) for c, r in
                zip(channels, refs)]

    lum_lin = lum_src if lum_src is not None else np.mean(channels, axis=0)
    lum, params = autostretch(lum_lin)

    # Broadband and narrowband want opposite treatments here.
    #
    # For broadband, all three channels share the luminance's midtone (a
    # "linked" stretch). That preserves the real brightness ratios between
    # them, which is what keeps star colours honest.
    #
    # For narrowband it is wrong. The three lines have wildly different
    # strengths - Ha is typically far brighter than SII and OIII - so a linked
    # stretch renders the whole nebula in whichever colour Ha was mapped to,
    # which in the Hubble palette means overwhelmingly green. Stretching each
    # channel to its OWN sky target instead brings the weak lines up to
    # comparable visual weight, which is what the palette is for. The result is
    # deliberately false colour either way; this makes it a readable false
    # colour.
    unlinked = palette in PALETTE_MAP
    rgb = np.empty(lum.shape + (3,), np.float32)
    for i, ch in enumerate(channels):
        if unlinked:
            rgb[:, :, i], _ = autostretch(ch)
        else:
            sky = float(np.median(ch))
            mad = 1.4826 * float(np.median(np.abs(ch - sky))) or 1.0
            black = sky - 2.8 * mad
            white = float(np.percentile(ch, 99.995))
            rgb[:, :, i] = mtf(
                np.clip((ch - black) / (max(white - black, 1e-6)), 0, 1),
                params["midtone"])
    if unlinked:
        # With each channel independently stretched the composite already
        # carries its own brightness; taking luminance from the linear mean
        # would flatten it back down.
        lum = rgb.mean(axis=2)
    del channels

    # Neutralise the sky: calibrating on stars leaves the background tinted,
    # and without this an empty frame reads brown or blue depending on the moon.
    sky_med = [float(np.median(rgb[:, :, i][~mask])) for i in range(3)]
    neutral_to = float(np.mean(sky_med))
    for i in range(3):
        rgb[:, :, i] = np.clip(rgb[:, :, i] - (sky_med[i] - neutral_to), 0, 1)

    # Denoise colour only; structure comes from the luminance, so this is
    # invisible at normal viewing scale.
    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 * SATURATION, 0, 1)
    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

    # Narrowband palettes are deliberately false colour, so green suppression
    # would fight the palette. It only applies to broadband.
    if palette not in PALETTE_MAP:
        neutralish = 0.5 * (out[:, :, 0] + out[:, :, 2])
        green = out[:, :, 1]
        out[:, :, 1] = np.where(green > neutralish,
                                green * 0.15 + neutralish * 0.85, green)

    return _write(session, (out * 65535 + 0.5).astype(np.uint16), palette,
                  verbose)


def _write(session, rgb16, palette, verbose=True):
    target = session.target.replace(" ", "")
    final = os.path.join(session.root, layout.FINAL)
    os.makedirs(final, exist_ok=True)
    stem = f"{target}-{palette}"
    tifffile.imwrite(os.path.join(final, stem + ".tif"), rgb16,
                     photometric="rgb")
    u8 = (rgb16 / 257).astype(np.uint8)
    Image.fromarray(u8).save(os.path.join(final, stem + ".png"))
    prev = Image.fromarray(u8)
    prev.thumbnail((2400, 2400), Image.LANCZOS)
    prev.save(os.path.join(final, stem + "-preview.jpg"), quality=92)
    if verbose:
        print(f"      -> final/{stem}.png / .tif / -preview.jpg")
    return os.path.join(final, stem + ".png")