#!/usr/bin/env python3
"""Stack Christopher Go's raw Jupiter capture, and measure what stacking did.

Input: 2022-07-22-2041_2-CG-L-18.avi -- 3,215 frames, 720x620, 8-bit RGGB Bayer,
107 fps, Celestron C14 (356 mm), 22 July 2022 20:41 UT. Used with permission.

Pass 1 scores every frame for sharpness on one Bayer plane (so the colour filter
array cannot leak into the metric). Pass 2 demosaics, registers each frame to a
fixed reference by phase correlation, and averages -- the whole run, and the
sharpest 5% separately. No deconvolution is applied anywhere.

Two things this script gets right that the obvious version gets wrong:
  * ONE reference frame throughout. Switching reference mid-pass silently
    produces two mutually-misaligned groups and a blurrier stack than a single
    frame.
  * The SIGN of the phase-correlation shift is verified against the data rather
    than assumed -- applied the wrong way it pushes frames apart. The check is
    printed; residuals were 29.5 wrong-way, 15.3 unregistered, 4.3 right-way.
"""
import numpy as np, subprocess, os, cv2, sys

W, H = 720, 620
N = W * H
CODE = cv2.COLOR_BayerBG2BGR          # verified: leaves the sampled RGGB sites untouched
SRC = sys.argv[1] if len(sys.argv) > 1 else "2022-07-22-2041_2-CG-L-18.avi"
OUT = os.path.dirname(os.path.abspath(SRC)) or "."


def frames():
    p = subprocess.Popen(["ffmpeg", "-v", "error", "-i", SRC,
                          "-pix_fmt", "gray", "-f", "rawvideo", "-"],
                         stdout=subprocess.PIPE, bufsize=N * 8)
    while True:
        b = p.stdout.read(N)
        if len(b) < N:
            break
        yield np.frombuffer(b, dtype=np.uint8).reshape(H, W)
    p.stdout.close(); p.wait()


def check_demosaic(raw):
    """A correct demosaic returns the sampled value untouched at its own site."""
    d = cv2.cvtColor(raw, CODE)
    assert np.array_equal(d[0::2, 0::2, 2], raw[0::2, 0::2]), "R sites not preserved"
    assert np.array_equal(d[1::2, 1::2, 0], raw[1::2, 1::2]), "B sites not preserved"


def main():
    sharp = []
    for raw in frames():                                   # pass 1
        g = raw[::2, ::2].astype(np.float32)
        sharp.append(np.abs(np.diff(g, axis=1)).mean() + np.abs(np.diff(g, axis=0)).mean())
    sharp = np.array(sharp)
    order = np.argsort(sharp)
    i_best, i_med = int(order[-1]), int(order[len(order) // 2])
    best = set(order[-int(0.05 * sharp.size):].tolist())
    print("frames %d | sharpness best/worst %.3f" % (sharp.size, sharp.max() / sharp.min()))

    ref = None
    for i, raw in enumerate(frames()):
        if i == i_best:
            check_demosaic(raw)
            ref = cv2.cvtColor(raw, CODE).astype(np.float32)
            break
    ref_g = np.ascontiguousarray(ref[..., 1])
    hann = cv2.createHanningWindow((W, H), cv2.CV_32F)

    def shift(img, sx, sy):
        return cv2.warpAffine(img, np.float32([[1, 0, sx], [0, 1, sy]]), (W, H),
                              flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)

    probe = None
    for i, raw in enumerate(frames()):
        if i == 3000:
            probe = np.ascontiguousarray(cv2.cvtColor(raw, CODE).astype(np.float32)[..., 1])
            break
    (dx, dy), _ = cv2.phaseCorrelate(ref_g, probe, hann)
    r_plus = np.abs(shift(probe, dx, dy) - ref_g).mean()
    r_minus = np.abs(shift(probe, -dx, -dy) - ref_g).mean()
    sgn = -1.0 if r_minus < r_plus else 1.0
    print("sign check: unregistered %.2f | +shift %.2f | -shift %.2f -> %+.0f"
          % (np.abs(probe - ref_g).mean(), r_plus, r_minus, sgn))

    sum_all = np.zeros((H, W, 3), np.float64); n_all = 0
    sum_best = np.zeros((H, W, 3), np.float64); n_best = 0
    single = {}
    for i, raw in enumerate(frames()):                     # pass 2
        rgb = cv2.cvtColor(raw, CODE).astype(np.float32)
        (dx, dy), _ = cv2.phaseCorrelate(ref_g, np.ascontiguousarray(rgb[..., 1]), hann)
        reg = shift(rgb, sgn * dx, sgn * dy)
        sum_all += reg; n_all += 1
        if i in best:
            sum_best += reg; n_best += 1
        if i in (i_med, i_best):
            single[i] = rgb.copy()

    def save(name, a):
        cv2.imwrite(os.path.join(OUT, name),
                    (np.clip(a, 0, 255) / 255.0 * 65535).astype(np.uint16))
    save("go_single.png", single[i_med])
    save("go_stack_best.png", sum_best / n_best)
    save("go_stack_all.png", sum_all / n_all)
    print("stacked: best %d, all %d" % (n_best, n_all))


if __name__ == "__main__":
    main()


# ---------------------------------------------------------------------------
# Aperture ladder: take the same capture and impose on every frame the
# diffraction limit and reduced light grasp of a smaller mirror, then stack.
# Extra blur is the quadrature difference of the two diffraction limits;
# noise is raised so per-frame SNR falls with collecting area. The native
# per-frame sigma is measured from a difference of two registered frames
# rather than assumed.  Run with --ladder.
# ---------------------------------------------------------------------------