#!/usr/bin/env python3
"""Normalize white/black points of all raster images in a PDF.

For every embedded image, a shared linear stretch maps a percentile-defined
dark point to pure black and a bright point to pure white, clipping nothing
in between. A single transform is applied to all color channels so mid-tone
color balance is preserved (hue unchanged).

The PDF structure, text, vector content, and page layout are left untouched:
only the image pixel data is swapped in place.

To keep memory bounded on large scanned books, the PDF is processed in three
stages and never more than one page's worth of decoded pixels is held at a
time:

  1. split:      each page is extracted to its own temporary PDF file
  2. process:    each page file's images are stretched and the page is saved
  3. merge:      all processed page files are joined into the output PDF

Note: an image shared by several pages is duplicated into each page during
splitting, so the output may be slightly larger than the input.

Dependencies: pip install pymupdf numpy
"""

import argparse
import os
import tempfile

import numpy as np
import pymupdf

try:
    import cv2
    HAVE_CV2 = True
except ImportError:
    HAVE_CV2 = False


def stretch(arr, black_pct, white_pct, per_channel, gain_cap):
    """Return (new_array, lo, hi, changed).

    A linear stretch maps the black/white percentiles toward full range, but
    the per-pixel gain is capped at `gain_cap` (if given) and anchored so the
    white point lands at ~245, so hazy scans don't get their noise amplified
    into visible grain / color blotches."""
    h, w, n = arr.shape
    n_color = min(n, 3)
    color = arr[..., :n_color].astype(np.float32)
    alpha = arr[..., n_color:] if n > n_color else None

    if per_channel and n_color == 3:
        lo = np.percentile(color, black_pct, axis=(0, 1))
        hi = np.percentile(color, white_pct, axis=(0, 1))
    elif n_color == 1:
        lo = np.percentile(color, black_pct)
        hi = np.percentile(color, white_pct)
    else:
        luma = 0.299 * color[..., 0] + 0.587 * color[..., 1] + 0.114 * color[..., 2]
        lo = np.percentile(luma, black_pct)
        hi = np.percentile(luma, white_pct)

    lo = float(np.min(lo))
    hi = float(np.max(hi))
    if hi - lo < 1e-3:
        return arr, lo, hi, False

    gain = 255.0 / (hi - lo)
    if gain_cap:
        gain = min(gain, gain_cap)
    offset = 245.0 - gain * hi
    out = np.clip(gain * color + offset, 0.0, 255.0).astype(np.uint8)
    result = np.empty((h, w, n), dtype=np.uint8)
    result[..., :n_color] = out
    if alpha is not None:
        result[..., n_color:] = alpha
    return result, lo, hi, True


def local_contrast(arr, clip):
    """CLAHE on luminance, applied to all channels via a common multiplier so
    color balance is preserved. Boosts local text contrast without amplifying
    global lighting gradients or chroma noise."""
    luma = (0.299 * arr[..., 0] + 0.587 * arr[..., 1] + 0.114 * arr[..., 2]).astype(np.uint8)
    eq = cv2.createCLAHE(clipLimit=clip, tileGridSize=(8, 8)).apply(luma)
    ratio = (eq.astype(np.float32) + 1) / (luma.astype(np.float32) + 1)
    return np.clip(arr.astype(np.float32) * ratio[..., None], 0, 255).astype(np.uint8)


def deartifact(arr, strength):
    """Mildly suppress JPEG2000 ringing/smearing around text edges.

    NL-means denoise preserves edges while flattening the low-amplitude
    ringing that the contrast stretch would otherwise amplify, followed by a
    gentle unsharp that restores edge crispness lost to the source encoding.
    """
    if arr.shape[2] == 3:
        out = cv2.fastNlMeansDenoisingColored(arr, None, strength, strength, 7, 21)
    else:
        out = cv2.fastNlMeansDenoising(arr, None, strength, 7, 21)
    blur = cv2.GaussianBlur(out, (0, 0), 2.0)
    return cv2.addWeighted(out, 1.2, blur, -0.2, 0)


def to_pixmap(arr, gray):
    h, w = arr.shape[:2]
    cs = pymupdf.csGRAY if gray else pymupdf.csRGB
    has_alpha = arr.shape[2] > (1 if gray else 3)
    return pymupdf.Pixmap(cs, w, h, arr.tobytes(), has_alpha)


def process_page(doc, page_index, black_pct, white_pct, per_channel, jpg_quality,
                 deartifact_strength, gain_cap, clahe_clip, remove_watermark, tmpdir):
    """Stretch every image on one page of a single-page doc.

    Returns (n_changed, n_skipped)."""
    page = doc[page_index]
    n_changed = n_skipped = 0
    seen = set()
    for img in page.get_images(full=True):
        xref = img[0]
        if xref in seen:
            continue
        seen.add(xref)
        if img[1] != 0:
            if remove_watermark:
                page.delete_image(xref)
                print(f"    xref {xref}: soft-masked image (watermark); removed", flush=True)
            else:
                print(f"    xref {xref}: soft-masked image (watermark); left untouched", flush=True)
                n_skipped += 1
            continue
        try:
            pix = pymupdf.Pixmap(doc, xref)
            if pix.colorspace is None:
                continue
            if pix.colorspace == pymupdf.csCMYK:
                pix = pymupdf.Pixmap(pymupdf.csRGB, pix)
            gray = pix.colorspace == pymupdf.csGRAY
            arr = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, pix.n).copy()
            new, lo, hi, changed = stretch(arr, black_pct, white_pct, per_channel, gain_cap)
            if clahe_clip and HAVE_CV2:
                new = local_contrast(new, clahe_clip)
            if deartifact_strength and HAVE_CV2:
                new = deartifact(new, deartifact_strength)
        except Exception as e:
            print(f"    xref {xref}: {e}; skipped", flush=True)
            continue
        if not changed:
            n_skipped += 1
            print(f"    xref {xref}: {pix.width}x{pix.height} already full-range ({lo:.0f}..{hi:.0f}); skipped", flush=True)
            continue
        new_pix = to_pixmap(new, gray)
        jpg = os.path.join(tmpdir, f"img_{xref}.jpg")
        new_pix.save(jpg, jpg_quality=jpg_quality)
        page.replace_image(xref, filename=jpg)
        os.remove(jpg)
        n_changed += 1
        print(f"    xref {xref}: {pix.width}x{pix.height} stretched {lo:.0f}..{hi:.0f} -> 0..255 (JPEG q{jpg_quality})", flush=True)
    return n_changed, n_skipped


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("pdf", help="input PDF")
    ap.add_argument("-o", "--output", help="output PDF (default: <input>-normalized.pdf)")
    ap.add_argument("--black-pct", type=float, default=5.0,
                    help="lower brightness percentile used as black point (default 5.0)")
    ap.add_argument("--white-pct", type=float, default=95.0,
                    help="upper brightness percentile used as white point (default 95.0)")
    ap.add_argument("--per-channel", action="store_true",
                    help="stretch each RGB channel independently (default: shared map preserves color balance)")
    ap.add_argument("--gain-cap", type=float, default=1.8,
                    help="cap the contrast gain (default 1.8; higher = more aggressive, more noise/color blotches)")
    ap.add_argument("--clahe", type=float, default=2.0, metavar="CLIP",
                    help="CLAHE local-contrast clip limit on luminance (default 2.0; 0 = off)")
    ap.add_argument("--jpg-quality", type=int, default=85,
                    help="JPEG quality for re-encoded images (default 85; lower = smaller file)")
    ap.add_argument("--deartifact", type=int, default=0, metavar="H",
                    help="reduce JPEG2000 ringing/smearing: NL-means denoise with strength H "
                         "(try 3-5; adds ~1.5 s per image, requires opencv-python)")
    ap.add_argument("--remove-watermark", action="store_true",
                    help="delete soft-masked background images instead of leaving them untouched")
    ap.add_argument("--keep-temp", metavar="DIR",
                    help="keep the per-page temporary files in DIR instead of deleting them")
    args = ap.parse_args()

    if args.deartifact and not HAVE_CV2:
        print("--deartifact requires opencv-python (pip install opencv-python-headless)")
        return

    src = pymupdf.open(args.pdf)
    n_pages = src.page_count
    print(f"PDF: {n_pages} pages", flush=True)
    if n_pages == 0:
        print("Empty PDF; nothing to do.")
        return

    out = args.output or args.pdf.rsplit(".", 1)[0] + "-normalized.pdf"
    tmp = tempfile.TemporaryDirectory(prefix="normalize-",
                                      dir=args.keep_temp or None)

    n_changed = n_skipped = 0
    with tmp:
        # Stage 1+2: extract each page, process it, save it to disk.
        for i in range(n_pages):
            page_doc = pymupdf.open()
            try:
                page_doc.insert_pdf(src, from_page=i, to_page=i)
                if i % 25 == 0:
                    print(f"  page {i+1}/{n_pages} ...", flush=True)
                c, s = process_page(page_doc, 0, args.black_pct, args.white_pct, args.per_channel,
                                    args.jpg_quality, args.deartifact, args.gain_cap,
                                    args.clahe, args.remove_watermark, tmp.name)
                n_changed += c
                n_skipped += s
                page_doc.save(os.path.join(tmp.name, f"page_{i:05d}.pdf"), garbage=1, deflate=True)
            finally:
                page_doc.close()

        # Stage 3: join all processed pages into the final PDF.
        print("Merging pages ...", flush=True)
        merged = pymupdf.open()
        try:
            for i in range(n_pages):
                page_doc = pymupdf.open(os.path.join(tmp.name, f"page_{i:05d}.pdf"))
                try:
                    merged.insert_pdf(page_doc)
                finally:
                    page_doc.close()
            merged.save(out, garbage=4, deflate=True)
        finally:
            merged.close()

    src.close()
    if n_changed == 0:
        print(f"No images needed adjustment ({n_skipped} already full-range).")
    print(f"Saved: {out} ({n_changed} images replaced, {n_skipped} skipped)", flush=True)


if __name__ == "__main__":
    main()