#!/usr/bin/env python3
"""make_figures.py - image figures for docs/raw12-pipeline-measurements.md.

    ~/.venvs/fgc-plots/bin/python make_figures.py <jxl_dir> <out_dir> [detail_x,detail_y shadow_x,shadow_y]

<jxl_dir> holds the example .jxl files `raw12_bench examples` wrote (its jxl/
folder, copied off the tower). Decodes them with imagecodecs and OpenCV, so it
runs on the laptop venv. The reference is c12_ea_lossless.jxl (the 12-bit
frame, exact); its >>4 is the 8-bit view of the same moment. Writes:
  thumbs/<name>.jpg      downscaled full frame (1/4)
  detail/<name>.png      480x320 100% crop of a detailed region (8-bit view)
  diff/<name>.png        |decoded - reference| in that crop, x8 (12-bit examples only)
  shadow/<name>.png      480x320 100% crop of a dark region, pushed +2 stops
  shadow_8_vs_12.png     that shadow, 8-bit vs 12-bit of one frame, +2 stops
  crops.jpg              where the crops are
  levels.json            crop positions, distinct levels in the shadow
Crops are chosen automatically unless given: detail = most structure after
blurring away pixel noise, without clipped sky; shadow = darkest textured region.
"""
import glob
import json
import os
import sys

import cv2
import imagecodecs
import numpy as np

src, out = sys.argv[1], sys.argv[2]
for d in ("thumbs", "detail", "diff", "shadow"):
    os.makedirs(os.path.join(out, d), exist_ok=True)
CW, CH = 480, 320


def load(path):
    """Decoded example as HxWx3 RGB on the 12-bit scale (float)."""
    a = imagecodecs.jpegxl_decode(open(path, "rb").read()).astype(np.float64)
    name = os.path.basename(path)
    eight = name.startswith(("rgb8", "c8"))
    if a.max() > 4095.5:
        a = a / 16.0
    if eight:
        a = a * 4095.0 / 255.0
    if a.ndim == 2:
        m = np.round(a).astype(np.uint16)
        if "planes" in name:
            h, w = m.shape
            hh, hw = h // 2, w // 2
            mo = np.empty_like(m)
            mo[0::2, 0::2], mo[0::2, 1::2] = m[:hh, :hw], m[:hh, hw:]
            mo[1::2, 0::2], mo[1::2, 1::2] = m[hh:, :hw], m[hh:, hw:]
            m = mo
        a = cv2.cvtColor(m, cv2.COLOR_BayerBG2RGB_EA).astype(np.float64)
    return a, eight


def view8(rgb12, gain=1.0):
    return cv2.cvtColor(np.clip(rgb12 * gain / 4095.0 * 255.0, 0, 255).astype(np.uint8),
                        cv2.COLOR_RGB2BGR)


ref, _ = load(os.path.join(src, "c12_ea_lossless.jxl"))
h, w = ref.shape[:2]
lum = 0.2126 * ref[..., 0] + 0.7152 * ref[..., 1] + 0.0722 * ref[..., 2]

if len(sys.argv) > 4:
    dx, dy = map(int, sys.argv[3].split(","))
    sx, sy = map(int, sys.argv[4].split(","))
else:
    smooth = cv2.GaussianBlur(lum.astype(np.float32), (0, 0), 2.0)
    structure = np.abs(cv2.Laplacian(smooth, cv2.CV_32F))
    best_d, best_s = (-1, None), (1e9, None)
    for y in range(CH, h - 2 * CH, CH // 2):
        for x in range(CW, w - 2 * CW, CW // 2):
            win = lum[y:y + CH, x:x + CW]
            st = float(structure[y:y + CH, x:x + CW].mean())
            if (win >= 4000).mean() < 0.01 and st > best_d[0]:
                best_d = (st, (x, y))
            if st > 0.5 * float(structure.mean()) and win.mean() < best_s[0]:
                best_s = (float(win.mean()), (x, y))
    (dx, dy), (sx, sy) = best_d[1], best_s[1]

ref_detail = ref[dy:dy + CH, dx:dx + CW]
for path in sorted(glob.glob(os.path.join(src, "*.jxl"))):
    name = os.path.splitext(os.path.basename(path))[0]
    img, eight = load(path)
    small = cv2.resize(img, (w // 4, h // 4), interpolation=cv2.INTER_AREA)
    cv2.imwrite(os.path.join(out, "thumbs", name + ".jpg"), view8(small),
                [cv2.IMWRITE_JPEG_QUALITY, 88])
    cv2.imwrite(os.path.join(out, "detail", name + ".png"), view8(img[dy:dy + CH, dx:dx + CW]))
    cv2.imwrite(os.path.join(out, "shadow", name + ".png"), view8(img[sy:sy + CH, sx:sx + CW], 4.0))
    if not eight:  # 8-bit examples are other captures: no pixel-by-pixel difference
        diff = np.abs(img[dy:dy + CH, dx:dx + CW] - ref_detail).mean(axis=2)
        cv2.imwrite(os.path.join(out, "diff", name + ".png"),
                    np.clip(diff * 8 / 16, 0, 255).astype(np.uint8))

# 8-bit vs 12-bit of the same frame: the 8-bit view is the 12-bit >> 4.
s12 = ref[sy:sy + CH, sx:sx + CW]
s8 = np.floor(s12 / 16.0) * 16.0
gap = np.full((CH, 8, 3), 255, np.uint8)
cv2.imwrite(os.path.join(out, "shadow_8_vs_12.png"),
            np.hstack([view8(s8, 4.0), gap, view8(s12, 4.0)]))
# Where the crops are.
th = view8(cv2.resize(ref, (w // 4, h // 4), interpolation=cv2.INTER_AREA))
for (x, y), c in (((dx, dy), (0, 0, 255)), ((sx, sy), (255, 160, 0))):
    cv2.rectangle(th, (x // 4, y // 4), ((x + CW) // 4, (y + CH) // 4), c, 2)
cv2.imwrite(os.path.join(out, "crops.jpg"), th, [cv2.IMWRITE_JPEG_QUALITY, 88])

g12 = np.round(s12[..., 1]).astype(int)
levels = dict(
    detail_crop=dict(x=dx, y=dy, w=CW, h=CH), shadow_crop=dict(x=sx, y=sy, w=CW, h=CH),
    shadow_mean_12bit=round(float(s12.mean()), 1),
    shadow_distinct_levels_8bit=int(len(np.unique(g12 >> 4))),
    shadow_distinct_levels_12bit=int(len(np.unique(g12))),
    shadow_green_range_12bit=[int(g12.min()), int(g12.max())],
)
json.dump(levels, open(os.path.join(out, "levels.json"), "w"), indent=1)
print(json.dumps(levels))