"""The four deliverable images, for any session.

The set was worked out on Centaurus A and is the same for every target, because
the point of it is comparison: the same data at four levels of treatment, so a
viewer can see what processing did and did not add.

    1-stacked     aligned and averaged, then a standard stretch. Nothing else:
                  no outlier rejection, no gradient removal, no colour work.
                  The honest starting point.
    2-processed   conventional processing - gradient, colour balance, denoise.
    3-best        everything the measurements justify: a background fit that
                  cannot eat the target, colour calibrated on solar-analogue
                  stars where astrometry allows it, the core recovered by a
                  second tone curve, star-protected sharpening, and noise
                  reduction that leaves compact sources alone.
    4-closeup     image 3 cropped to the target's own measured extent.

Plus a side-by-side comparison of the first three.

Image 3's list is not a style. Each item exists because measuring the first
session showed the conventional version getting something wrong: the plane fit
had absorbed 17.9 ADU/px of galaxy halo, "average star is grey" was biased by a
field whose stars are redder than the Sun, the core was flattened by a white
point set by field stars, deconvolution rang around every bright star, and
denoising erased faint compact sources that turned out to be globular clusters.
"""
import os

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 denoise_tv_chambolle, richardson_lucy

import colour
import layout

RL_ITERS = 10
PSF_BOX = 25


def _stem(session, tag):
    return f"{session.target.replace(' ', '')}-{tag}"


def _save(session, rgb01, tag, verbose=True):
    final = os.path.join(session.root, layout.FINAL)
    os.makedirs(final, exist_ok=True)
    stem = _stem(session, tag)
    arr = np.clip(rgb01, 0, 1)
    u8 = (arr * 255 + 0.5).astype(np.uint8)
    Image.fromarray(u8).save(os.path.join(final, stem + ".png"))
    tifffile.imwrite(os.path.join(final, stem + ".tif"),
                     (arr * 65535 + 0.5).astype(np.uint16), photometric="rgb")
    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")


def measure_psf(img):
    """Median-stack isolated unsaturated 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
    xs, ys = o["x"], o["y"]
    cuts = []
    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:      # isolated only
            continue
        cut = sub[int(ys[i]) - h:int(ys[i]) + h + 1,
                  int(xs[i]) - h:int(xs[i]) + h + 1].astype(np.float64)
        if cut.shape != (PSF_BOX, PSF_BOX) or cut.max() <= 0:
            continue
        cuts.append(cut / cut.sum())
        if len(cuts) >= 120:
            break
    if len(cuts) < 8:
        return None
    psf = np.median(np.stack(cuts), axis=0)
    psf[psf < 0] = 0
    return (psf / psf.sum()).astype(np.float32)


def star_mask(img, rms, max_r=18):
    """Feathered mask over stars, to keep deconvolution and denoise off them."""
    bkg = sep.Background(img, bw=64, bh=64, fw=3, fh=3)
    o = sep.extract(img - bkg.back(), 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)
    R = 20
    yy, xx = np.mgrid[-R:R + 1, -R:R + 1]
    rr = np.hypot(xx, yy)
    for x, y, npix in zip(o["x"], o["y"], o["npix"]):
        r = float(np.clip(2.5 * np.sqrt(npix / np.pi) + 3.0, 4, max_r))
        cx, cy = int(round(x)), int(round(y))
        x0, x1 = max(0, cx - R), min(nx, cx + R + 1)
        y0, y1 = max(0, cy - R), min(ny, cy + R + 1)
        patch = (rr <= r).astype(np.float32)[(y0 - cy + R):(y1 - cy + R),
                                             (x0 - cx + R):(x1 - cx + R)]
        np.maximum(m[y0:y1, x0:x1], patch, out=m[y0:y1, x0:x1])
    return np.clip(gaussian_filter(m, 2.5), 0, 1), len(o)


def deconvolve(img, psf):
    """Richardson-Lucy on the signal, never on a star.

    RL rings around anything steeper than the PSF model can explain, which on a
    star field means a dark annulus round every bright star. Excluding stars is
    the standard cure and costs nothing, since the extended target is what
    needed sharpening.
    """
    if psf is None:
        return img
    rms = 1.4826 * float(np.median(np.abs(img - np.median(img)))) or 1.0
    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
    signal = np.clip((img - 3.0 * rms) / (20.0 * rms), 0.0, 1.0)
    smask, _ = star_mask(img, rms)
    w = (signal * (1.0 - smask)).astype(np.float32)
    return (out * w + img * (1.0 - w)).astype(np.float32)


def object_extent(img, mask):
    """Bounding box of the target, from where the signal actually is."""
    ys, xs = np.where(mask)
    if len(xs) == 0:
        ny, nx = img.shape
        return 0, nx, 0, ny
    return int(xs.min()), int(xs.max()), int(ys.min()), int(ys.max())


def build(session, verbose=True):
    """Produce all four images plus the comparison. Returns their paths."""
    masters = colour.load_masters(session)
    if not masters:
        raise RuntimeError("no masters - run the register stage first")
    palette = session.palette
    out = {}

    # ---------------------------------------------------------------- image 1
    # The baseline: masters as stacked, one standard stretch, nothing else.
    raw = {f: fits.getdata(p).astype(np.float32) for f, p in masters.items()}
    img1 = colour.assemble(session, dict(raw), palette, gradient=False,
                           neutralise=False, denoise=False, saturation=1.0,
                           verbose=False)
    out["1-stacked"] = _save(session, img1, "1-stacked", verbose)
    del img1

    # ---------------------------------------------------------------- image 2
    img2 = colour.assemble(session, {f: v.copy() for f, v in raw.items()},
                           palette, gradient=True, neutralise=True,
                           denoise=True, saturation=colour.SATURATION,
                           verbose=verbose)
    out["2-processed"] = _save(session, img2, "2-processed", verbose)
    del img2

    # ---------------------------------------------------------------- image 3
    data = {f: v.copy() for f, v in raw.items()}
    del raw
    deepest = session.filters[0] if session.filters[0] in data \
        else sorted(data)[0]
    mask = colour.target_mask(data[deepest])

    # Sharpen the deep channel only, with stars protected.
    psf = measure_psf(data[deepest])
    if psf is not None:
        if verbose:
            print(f"      deconvolving {deepest} (star-protected)")
        data[deepest] = deconvolve(data[deepest], psf)

    img3 = colour.assemble(session, data, palette, gradient=True,
                           neutralise=True, denoise=True,
                           saturation=colour.SATURATION + 0.05,
                           hdr=True, protect_compact=True, verbose=verbose)
    out["3-best"] = _save(session, img3, "3-best", verbose)

    # ---------------------------------------------------------------- image 4
    x0, x1, y0, y1 = object_extent(data[deepest], mask)
    ny, nx = img3.shape[:2]
    # Margin around the target, then clamped to the frame. Deliberately NOT
    # forced square: a square crop cannot contain a target wider than the
    # frame is tall, which is the normal case for a nebula filling a wide
    # field, and forcing it produces a close-up that cuts the subject in half.
    mx = int(0.08 * (x1 - x0))
    my = int(0.08 * (y1 - y0))
    left, right = max(0, x0 - mx), min(nx, x1 + mx)
    top, bottom = max(0, y0 - my), min(ny, y1 + my)

    covers = (right - left) / nx > 0.95 and (bottom - top) / ny > 0.95
    if covers:
        # The target fills the field. There is no close-up to make, and
        # pretending otherwise would just re-save image 3 under a name that
        # claims to be something else.
        if verbose:
            print("      target fills the frame - no close-up is possible")
    elif right - left > 64 and bottom - top > 64:
        crop = img3[top:bottom, left:right]
        assert left <= x0 and right >= x1 and top <= y0 and bottom >= y1,             "close-up would clip the target"
        out["4-closeup"] = _save(session, crop, "4-closeup", verbose)
        if verbose:
            print(f"      close-up {right-left}x{bottom-top} px, "
                  f"clearance {min(x0-left, right-x1, y0-top, bottom-y1)} px")

    del img3, data

    _comparison(session, out, verbose)
    return out


def _comparison(session, paths, verbose=True):
    """Side-by-side of the first three, so the progression is checkable."""
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    panels = [("1-stacked", "1. Stacked + stretch"),
              ("2-processed", "2. Conventional"),
              ("3-best", "3. Science-informed")]
    have = [(k, t) for k, t in panels if k in paths]
    if len(have) < 2:
        return
    fig, axes = plt.subplots(1, len(have), figsize=(7 * len(have), 7))
    fig.patch.set_facecolor("#111111")
    for ax, (key, title) in zip(np.atleast_1d(axes), have):
        im = Image.open(paths[key])
        im.thumbnail((1400, 1400), Image.LANCZOS)
        ax.imshow(np.asarray(im))
        ax.set_title(title, color="white", fontsize=14)
        ax.set_xticks([]); ax.set_yticks([])
    fig.suptitle(f"{session.target} - {session.telescope} - "
                 f"{session.palette} - same data, three treatments",
                 color="white", fontsize=16)
    fig.tight_layout(rect=[0, 0, 1, 0.94])
    path = os.path.join(session.root, layout.FINAL,
                        _stem(session, "comparison") + ".jpg")
    fig.savefig(path, dpi=90, facecolor=fig.get_facecolor(),
                pil_kwargs={"quality": 90})
    plt.close(fig)
    if verbose:
        print(f"      -> final/{os.path.basename(path)}")