"""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
from skimage.restoration import denoise_tv_chambolle

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 assemble(session, data, palette, gradient=True, neutralise=True,
             denoise=True, saturation=SATURATION, hdr=False,
             protect_compact=False, verbose=True):
    """Build an RGB image in [0,1] from per-filter linear masters.

    The flags exist so the same code can produce the untouched baseline and the
    fully corrected version, which is what makes the two comparable: the only
    differences between the deliverables are the ones named here.
    """
    deepest = session.filters[0] if session.filters[0] in data         else sorted(data)[0]
    mask = target_mask(data[deepest])

    if gradient:
        for filt in list(data):
            data[filt], _ = remove_gradient(data[filt], mask)

    if palette in PALETTE_MAP:
        mapping = PALETTE_MAP[palette]
        if not all(v in data for v in mapping.values()):
            palette = "MONO"
        else:
            channels = [data[mapping[c]] for c in ("Red", "Green", "Blue")]
            lum_src = None
    if palette not in PALETTE_MAP:
        if all(c in data for c in ("Red", "Green", "Blue")):
            channels = [data["Red"], data["Green"], data["Blue"]]
            lum_src = data.get("Luminance")
        else:
            grey, _ = autostretch(data[deepest])
            return np.dstack([np.clip(grey, 0, 1)] * 3)

    refs = [float(np.percentile(c, 99.5)) for c in channels]
    tgt = float(np.median(refs))
    channels = [c * (tgt / 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)

    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:
        lum = rgb.mean(axis=2)

    if hdr:
        # The white point above is set by field stars, which are far brighter
        # than an extended target, so the target's whole tonal range lands in
        # the top few percent of the curve and reads as a flat blob. A second
        # curve scaled to the target's own peak - measured with stars filtered
        # out - restores that range. Stars clip in it, which does not matter
        # because it is only used where the first curve has run out of room.
        h = max(64, min(lum_lin.shape) // 8)
        ny, nx = lum_lin.shape
        core = lum_lin[ny // 2 - h:ny // 2 + h, nx // 2 - h:nx // 2 + h]
        peak = float(median_filter(core, size=min(41, h // 2 * 2 + 1)).max())
        if peak > params["black"]:
            bright = mtf(np.clip((lum_lin - params["black"]) /
                                 max(peak * 1.15 - params["black"], 1e-6),
                                 0, 1), 0.35)
            w = gaussian_filter(
                np.clip((lum - 0.55) / 0.35, 0, 1).astype(np.float32), 8.0)
            lum = np.clip(lum * (1 - w) + bright * w, 0, 1)
            if verbose:
                print(f"      HDR blend over {float((w > 0.05).mean()):.1%} "
                      f"of the frame (target peak {peak:.0f} ADU)")

    if neutralise:
        sky_med = [float(np.median(rgb[:, :, i][~mask])) for i in range(3)]
        to = float(np.mean(sky_med))
        for i in range(3):
            rgb[:, :, i] = np.clip(rgb[:, :, i] - (sky_med[i] - to), 0, 1)

    if denoise:
        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

    if protect_compact:
        # Smooth the sky, but never a compact source. Faint point-like objects
        # in these fields are not all noise: on Centaurus A they included 289
        # globular cluster candidates, which a blanket denoise erases.
        import finals
        rms = 1.4826 * float(np.median(np.abs(lum - np.median(lum)))) or 1e-3
        smask, nsrc = finals.star_mask(lum.astype(np.float32), rms, max_r=12)
        keep = np.clip(np.clip((lum - 0.10) * 4.0, 0, 1) + smask, 0, 1)
        smooth = denoise_tv_chambolle(lum, weight=0.012)
        lum = np.clip(lum * keep + smooth * (1 - keep), 0, 1)
        if verbose:
            print(f"      sky denoise, {nsrc} compact sources protected")

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

    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 out


def run(session, verbose=True):
    """Produce the four deliverable images for this session."""
    import finals
    return finals.build(session, verbose=verbose)