#!/usr/bin/env python3
"""compare_8_vs_12.py - 8-bit against 12-bit on a dark, high-contrast scene.

    ~/.venvs/fgc-plots/bin/python compare_8_vs_12.py <BayerRG12p.tif> <RGB8.png> <out_dir> \
        [<october lossless c12 .jxl> <label> ...]

Inputs: the camera's BayerRG12p frame converted to colour by the viewer and saved as 16-bit TIFF
(12-bit values in the top bits), and the camera's own RGB8 frame of the same view taken seconds
later (another capture: trees move, so no pixel-by-pixel comparison with it). The 8-bit version of
the same frame is the 12-bit value >> 4 (what BayerRG8 would deliver). Optional extra lossless
12-bit colour files from earlier campaigns get the same noise measurement, for comparison.

Writes to <out_dir>:
  overview.jpg          the frame upright, display gamma, crop boxes
  noise_vs_signal.png   noise against brightness, every image measured the same way
  dark_crop.png         dark crop, camera RGB8 / 8-bit / 12-bit, +2 and +4 stops
  dark_compressed.png   dark crop after JPEG XL, +4 stops: today's camera RGB8 file (d0.8 e4),
                        and at d1.0 the 8-bit file, the 8-bit values in a 12-bit file, 12-bit
  mid_crop.png          mid-tone crop, 8-bit / 12-bit, contrast x4 around its mean (control)
  same_setting.png      dark crop, camera RGB8 and 12-bit, uncompressed and both at d0.8 e4,
                        shown x1, +2 and +4 stops
  same_setting_zoom.png the middle 200 x 200 pixels of the two d0.8 e4 crops at +4 stops, 2x
  results.json          noise, levels, sizes and errors per brightness region
"""
import json
import os
import sys

import cv2
import imagecodecs
import matplotlib
import matplotlib.ticker
import numpy as np

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402

tif, png, out = sys.argv[1:4]
extra = list(zip(sys.argv[4::2], sys.argv[5::2]))
os.makedirs(out, exist_ok=True)

c12 = cv2.imread(tif, cv2.IMREAD_UNCHANGED)[:, :, ::-1] >> 4         # 0-4095, uint16
cam8 = cv2.imread(png, cv2.IMREAD_UNCHANGED)[:, :, ::-1].copy()       # 0-255, uint8
c8 = (c12 >> 4).astype(np.uint8)                                      # same frame, top 8 bits
c8in12 = c8.astype(np.uint16) * 16 + 8   # the same 8-bit values, declared as a 12-bit file
# (integers only: + 8 rather than + 7.5; the 0.5 is taken off again after decoding)
H, W = c12.shape[:2]


def to12(a):
    """Any decoded image on the 12-bit scale; 8-bit v stands for the middle of 16v ... 16v + 15."""
    return a.astype(np.float32) * 16 + 7.5 if a.dtype == np.uint8 else a.astype(np.float32)


ref = c12.astype(np.float32)
lum_s = cv2.GaussianBlur(ref[..., 1].astype(np.float32), (0, 0), 8)
REGIONS = {"dark (< 300 DN)": lum_s < 300,
           "mid (300-1000 DN)": (lum_s >= 300) & (lum_s < 1000),
           "bright (>= 1000 DN)": lum_s >= 1000}

# --- noise against brightness -------------------------------------------------------------
P = 24
EDGES = np.array([100, 130, 160, 200, 250, 300, 400, 500, 700, 1000, 1400, 2000, 2400, 2800, 3300])


def noise_floor(g):
    """Per-pixel noise of a colour image's green channel, on the 12-bit scale.

    Standard deviation of the residual after a Gaussian (sigma 2 px) local mean, per 24 x 24 patch;
    patches grouped by brightness and the quietest 10% of each group averaged (texture only adds).
    Plain standard deviation, not a median-based one, so it also works on 8-bit steps. Neighbouring
    pixels of a de-Bayered image are correlated, which this measure tolerates better than
    neighbour differences.
    """
    g = g.astype(np.float32)
    r = g - cv2.GaussianBlur(g, (0, 0), 2.0)
    h, w = g.shape[0] // P * P, g.shape[1] // P * P
    sd = r[:h, :w].reshape(h // P, P, w // P, P).std(axis=(1, 3)).ravel()
    mn = g[:h, :w].reshape(h // P, P, w // P, P).mean(axis=(1, 3)).ravel()
    rows = []
    for a, b in zip(EDGES[:-1], EDGES[1:]):
        m = (mn >= a) & (mn < b)
        if m.sum() < 50:
            continue
        rows.append((float(np.median(mn[m])), float(np.sort(sd[m])[:max(1, m.sum() // 10)].mean())))
    return rows


# The residual of white noise under this high-pass is a little below its true sigma: calibrate.
wn = np.random.default_rng(0).normal(0, 10, (1200, 1200)) + 1000
CAL = 10 / noise_floor(wn)[0][1]
curves = {"12-bit (this frame)": [(s, n * CAL) for s, n in noise_floor(ref[..., 1])],
          "8-bit, same frame (12-bit >> 4)": [(s, n * CAL) for s, n in noise_floor(to12(c8)[..., 1])],
          "camera RGB8 (other capture)": [(s, n * CAL) for s, n in noise_floor(to12(cam8)[..., 1])]}
for path, label in extra:
    a = imagecodecs.jpegxl_decode(open(path, "rb").read()).astype(np.float32)
    a = a / 16 if a.max() > 4095.5 else a
    curves[label] = [(s, n * CAL) for s, n in noise_floor(a[..., 1])]

fig, ax = plt.subplots(figsize=(6.6, 3.6))
style = {"12-bit (this frame)": ("tab:blue", "-", "o"),
         "8-bit, same frame (12-bit >> 4)": ("tab:orange", "-", "s"),
         "camera RGB8 (other capture)": ("tab:purple", "-", "^")}
other = iter([("tab:green", "--", "."), ("tab:red", "--", ".")])
for label, rows in curves.items():
    c, ls, mk = style.get(label) or next(other)
    ax.plot([r[0] for r in rows], [r[1] for r in rows], color=c, ls=ls, marker=mk, ms=4, label=label)
ax.axhline(16, color="k", ls="--", lw=1, label="one 8-bit step (16)")
ax.axhline(16 / np.sqrt(12), color="gray", ls=":", lw=1, label="8-bit rounding error (4.6)")
ax.set_xscale("log"); ax.set_yscale("log")
ticks = matplotlib.ticker
ax.xaxis.set_major_locator(ticks.FixedLocator([150, 200, 300, 500, 700, 1000, 1500, 2000, 3000]))
ax.yaxis.set_major_locator(ticks.FixedLocator([5, 7, 10, 15, 20, 30, 50, 70]))
for axis in (ax.xaxis, ax.yaxis):
    axis.set_major_formatter(ticks.ScalarFormatter()); axis.set_minor_formatter(ticks.NullFormatter())
ax.set_xlabel("signal, 12-bit units (green)"); ax.set_ylabel("noise σ, 12-bit units")
ax.set_title("Noise vs brightness in the colour image (quietest 10% of 24×24 patches)", fontsize=10)
ax.grid(True, which="both", alpha=0.3)
ax.legend(fontsize=7, loc="upper left")
fig.tight_layout(); fig.savefig(os.path.join(out, "noise_vs_signal.png"), dpi=100); plt.close(fig)

# --- JPEG XL at the tower's settings (effort 3), 12-bit and 8-bit of the same frame ----------
DIST = [0, 0.3, 0.5, 1.0]
CS = 400
dark_xy = (40, 360)       # shaded forest, about 200 DN
mid_xy = (3800, 1300)     # distant valley in haze, about 1400 DN


def crop(a, xy):
    x, y = xy
    return a[y:y + CS, x:x + CS]


# Errors against the 12-bit frame, per region. Per pixel: everything, including discarded noise.
# Averaged: after a Gaussian average (sigma 4 px, about 200 pixels) of both images - the error in
# local brightness that remains once noise is averaged away, i.e. lost scene information such as a
# faint, even brightening.
ref_avg = cv2.GaussianBlur(ref, (0, 0), 4)
size, dark_dec, errors = {}, {}, []
SOURCES = (("12-bit", c12, 12), ("8-bit", c8, 8), ("8-bit values, 12-bit file", c8in12, 12),
           ("camera RGB8", cam8, 8))
for name, img, bps in SOURCES:
    for d in DIST:
        kw = dict(lossless=True) if d == 0 else dict(distance=d)
        data = imagecodecs.jpegxl_encode(np.ascontiguousarray(img), effort=3, bitspersample=bps, **kw)
        size[(name, d)] = len(data) / 1e6
        print(f"{name} d{d}: {len(data) / 1e6:.2f} MB", flush=True)
        if name == "camera RGB8":  # another capture: sizes only
            continue
        a = to12(imagecodecs.jpegxl_decode(data)).astype(np.float32)
        if img is c8in12:
            a -= 0.5
        if name == "12-bit" and d == 0:
            assert np.array_equal(a, ref), "12-bit lossless must be exact"
        dark_dec[(name, d)] = crop(a, dark_xy).copy()
        e_px = a - ref
        e_av = cv2.GaussianBlur(a, (0, 0), 4) - ref_avg
        row = dict(source=name, distance=d, size_mb=round(size[(name, d)], 2))
        for rn, m in REGIONS.items():
            row[rn] = dict(rms_px=round(float(np.sqrt((e_px[m] ** 2).mean())), 2),
                           rms_avg=round(float(np.sqrt((e_av[m] ** 2).mean())), 2))
        errors.append(row)
        del a, e_px, e_av
# Both formats with today's setting (distance 0.8, effort 4). The two frames are separate captures,
# so each is scored against its own uncompressed original.
same_setting, same_crops = [], {}
for name, img, bps in (("camera RGB8", cam8, 8), ("BayerRG12p", c12, 12)):
    data = imagecodecs.jpegxl_encode(np.ascontiguousarray(img), effort=4, distance=0.8, bitspersample=bps)
    own = to12(img)
    a = to12(imagecodecs.jpegxl_decode(data))
    same_crops[name] = (crop(own, dark_xy).copy(), crop(a, dark_xy).copy())
    e_px = a - own
    e_av = cv2.GaussianBlur(a, (0, 0), 4) - cv2.GaussianBlur(own, (0, 0), 4)
    row = dict(source=name, distance=0.8, effort=4, size_mb=round(len(data) / 1e6, 2))
    for rn, m in REGIONS.items():
        row[rn] = dict(rms_px=round(float(np.sqrt((e_px[m] ** 2).mean())), 2),
                       rms_avg=round(float(np.sqrt((e_av[m] ** 2).mean())), 2))
    same_setting.append(row)
    print(json.dumps(row), flush=True)
    del own, a, e_px, e_av
size[("camera RGB8 today e4", 0.8)] = same_setting[0]["size_mb"]
dark_dec[("camera RGB8 today", 0.8)] = same_crops["camera RGB8"][1]

# --- crops ---------------------------------------------------------------------------------
def upright(a):
    return np.ascontiguousarray(np.rot90(a, 1))  # camera on its side: sky was on the right


def view(a12, stops):
    return np.clip(a12 * 2.0 ** stops / 16, 0, 255).astype(np.uint8)


def grid(tiles, col_labels, row_labels, path, titles=None):
    """Tiles in rows; labels per column and row, or a title per tile."""
    nr, nc = len(tiles), len(tiles[0])
    fig, axs = plt.subplots(nr, nc, figsize=(nc * 3.0, nr * 3.25), squeeze=False)
    for i in range(nr):
        for j in range(nc):
            axs[i][j].imshow(tiles[i][j], interpolation="nearest")
            axs[i][j].set_xticks([]); axs[i][j].set_yticks([])
            if titles:
                axs[i][j].set_title(titles[i][j], fontsize=8)
            elif i == 0:
                axs[i][j].set_title(col_labels[j], fontsize=8)
            if j == 0 and row_labels:
                axs[i][j].set_ylabel(row_labels[i], fontsize=8)
    fig.tight_layout(pad=0.4); fig.savefig(path, dpi=150); plt.close(fig)


d12, d8, dc = crop(ref, dark_xy), crop(to12(c8), dark_xy), crop(to12(cam8), dark_xy)
grid([[upright(view(a, s)) for a in (dc, d8, d12)] for s in (2, 4)],
     ["camera RGB8 (today, other capture)", "8-bit (same frame, 12-bit >> 4)", "12-bit"],
     ["+2 stops (×4)", "+4 stops (×16)"], os.path.join(out, "dark_crop.png"))
keys = [[("camera RGB8 today", 0.8), ("8-bit", 1.0)], [("8-bit values, 12-bit file", 1.0), ("12-bit", 1.0)]]
grid([[upright(view(dark_dec[k], 4)) for k in row] for row in keys], None, None,
     os.path.join(out, "dark_compressed.png"),
     titles=[["(a) camera RGB8 file, d0.8 e4 (today)", "(b) 8-bit file, d1.0"],
             ["(c) 8-bit values in a 12-bit file, d1.0", "(d) 12-bit file, d1.0"]])
tiles = [same_crops["camera RGB8"][0], same_crops["camera RGB8"][1],
         same_crops["BayerRG12p"][0], same_crops["BayerRG12p"][1]]
grid([[upright(view(a, s)) for a in tiles] for s in (0, 2, 4)],
     ["camera RGB8, uncompressed", "camera RGB8, d0.8 e4", "BayerRG12p, uncompressed", "BayerRG12p, d0.8 e4"],
     ["as stored (×1)", "+2 stops (×4)", "+4 stops (×16)"], os.path.join(out, "same_setting.png"))
z = slice(100, 300)
grid([[np.kron(upright(view(same_crops[n][1], 4))[z, z], np.ones((2, 2, 1), np.uint8))
       for n in ("camera RGB8", "BayerRG12p")]],
     ["camera RGB8, d0.8 e4", "BayerRG12p, d0.8 e4"], ["+4 stops (×16), 2× zoom"],
     os.path.join(out, "same_setting_zoom.png"))
mid_mean = crop(ref, mid_xy).mean(axis=(0, 1))  # per channel
grid([[upright(np.clip((crop(a, mid_xy) - mid_mean) * 4 / 16 + 128, 0, 255).astype(np.uint8))
       for a in (to12(c8), ref)]],
     ["8-bit (same frame)", "12-bit"], ["contrast ×4 around the mean"], os.path.join(out, "mid_crop.png"))

# Overview: upright, quarter size, display gamma (sRGB) so the forest is visible.
small = cv2.resize(ref, (W // 4, H // 4), interpolation=cv2.INTER_AREA) / 4095
small = np.where(small <= 0.0031308, 12.92 * small, 1.055 * small ** (1 / 2.4) - 0.055)
ov = cv2.cvtColor((np.clip(small, 0, 1) * 255).astype(np.uint8), cv2.COLOR_RGB2BGR)
for (x, y), c in ((dark_xy, (255, 160, 0)), (mid_xy, (0, 0, 255))):
    cv2.rectangle(ov, (x // 4, y // 4), ((x + CS) // 4, (y + CS) // 4), c, 2)
cv2.imwrite(os.path.join(out, "overview.jpg"), upright(ov), [cv2.IMWRITE_JPEG_QUALITY, 88])


def levels(a):
    g = np.round(a[..., 1]).astype(int)
    return int(len(np.unique(g)))


res = dict(
    frame=dict(width=W, height=H, green_range_12bit=[int(c12[..., 1].min()), int(c12[..., 1].max())],
               share_of_frame={k: round(float(m.mean()) * 100, 1) for k, m in REGIONS.items()}),
    noise_calibration=round(CAL, 3),
    noise={k: [[round(s), round(n, 1)] for s, n in v] for k, v in curves.items()},
    dark_crop=dict(x=dark_xy[0], y=dark_xy[1], size=CS, mean_green_12bit=round(float(d12[..., 1].mean()), 1),
                   green_range_12bit=[int(d12[..., 1].min()), int(d12[..., 1].max())],
                   distinct_green_levels={"12-bit": levels(d12), "8-bit same frame": levels(d8),
                                          "camera RGB8": levels(dc)}),
    mid_crop=dict(x=mid_xy[0], y=mid_xy[1], size=CS,
                  mean_green_12bit=round(float(crop(ref, mid_xy)[..., 1].mean()), 1)),
    sizes_mb={f"{n} d{d}": round(v, 2) for (n, d), v in size.items()},
    errors=errors,
    same_setting=same_setting,
)
json.dump(res, open(os.path.join(out, "results.json"), "w"), indent=1)
print(json.dumps(res, indent=1))