#!/usr/bin/env python3
"""Extract a plate image from the original scan, dropping the printed page
number (top) and the typeset caption (bottom), which are re-set in LaTeX.

usage: crop_plate.py PDF_PAGE [--top FRAC --bottom FRAC] [--keep-caption]
Writes plates/pNNNN.png (1-bit PNG when the source is bitonal).
Heuristic defaults: ignore the outer 4% (scanner edge), the top 7.5%
(page number) and the bottom 9% (caption). Override per page via manifest.
"""
import sys, argparse, subprocess, io, math
from PIL import Image, ImageOps
import numpy as np

import os
SRC = os.environ.get("ROSS_SRC", "source/10731406.pdf")
ap = argparse.ArgumentParser()
ap.add_argument("page", type=int)
ap.add_argument("--top", type=float, default=0.075)
ap.add_argument("--bottom", type=float, default=0.09)
ap.add_argument("--edge", type=float, default=0.04)
ap.add_argument("--dpi", type=int, default=300)
ap.add_argument("--box", default=None,
                help="explicit crop as LEFT,TOP,RIGHT,BOTTOM fractions of the raw page; skips the scanner-bar and page-number heuristics. Use for dark plates (a photographed manuscript), where \"dark = scanner bar\" does not hold")
ap.add_argument("--trim", default=None,
                help="post-rotation trim as TOP,RIGHT,BOTTOM,LEFT fractions, e.g. 0,0,0.06,0")
ap.add_argument("--rotate", type=int, default=0, choices=[0,90,180,270],
                help="rotate the crop counter-clockwise (lossless); plates the original prints sideways")
ap.add_argument("--raw", action="store_true",
                help="with --box: take the box exactly — no shadow sweep, deskew, bar sweep or autocrop. Use for a framed plate, whose own border rules the edge heuristics mistake for scanner bars")
ap.add_argument("--rule", type=float, default=0.012,
                help="dark runs thinner than this fraction are plate rules, not scanner bars")
a = ap.parse_args()

png = subprocess.run(["pdftoppm", "-png", "-r", str(a.dpi), "-f", str(a.page),
                      "-l", str(a.page), SRC], capture_output=True).stdout
im = Image.open(io.BytesIO(png)).convert("L")
W, H = im.size
# strip scanner black bars: any column/row that is >40% dark is an edge;
# keep only the widest run of "paper" columns/rows.
dark = np.array(im) < 100
def paper_run(profile, rule=None):
    ok = profile < 0.40
    # A thin dark run is a ruled line belonging to the plate (engraved column
    # rules, table borders), not a scanner bar: close it so it cannot split the
    # paper run. Scanner bars are an order of magnitude thicker.
    rule = int((rule if rule is not None else a.rule) * len(ok))
    i = 0
    while i < len(ok):
        if not ok[i]:
            j = i
            while j < len(ok) and not ok[j]: j += 1
            interior = i > 0.05*len(ok) and j < 0.95*len(ok)
            if j - i <= rule and interior: ok[i:j] = True
            i = j
        else: i += 1
    best=(0,0); i=0
    while i < len(ok):
        if ok[i]:
            j=i
            while j < len(ok) and ok[j]: j+=1
            if j-i > best[1]-best[0]: best=(i,j)
            i=j
        else: i+=1
    return best
if a.box:
    l, t, r, b = (float(x) for x in a.box.split(","))
    im = im.crop((int(W*l), int(H*t), int(W*r), int(H*b)))
else:
    c0,c1 = paper_run(dark.mean(0)); r0,r1 = paper_run(dark.mean(1))
    im = im.crop((c0,r0,c1,r1)); W,H = im.size
    box = (int(W*a.edge), int(H*a.top), int(W*(1-a.edge)), int(H*(1-a.bottom)))
    im = im.crop(box)
# A thin, very dark run just inside an edge is a scanner bar the paper-run
# heuristic kept (it is thinner than --rule, so it was read as a plate rule).
# Sweep it off before autocropping.
def strip_edge_bars(im, frac=0.06, maxthick=30):
    for _ in range(2):
        a = np.array(im) < 100
        W, H = im.size
        l, r, t, b = 0, W, 0, H
        cols = a.mean(0); rows = a.mean(1)
        for i in range(int(W*frac)):
            if cols[i] > 0.5: l = i + 1
        for i in range(W - 1, W - int(W*frac) - 1, -1):
            if cols[i] > 0.5: r = i
        for i in range(int(H*frac)):
            if rows[i] > 0.5: t = i + 1
        for i in range(H - 1, H - int(H*frac) - 1, -1):
            if rows[i] > 0.5: b = i
        if (l, r, t, b) == (0, W, 0, H): break
        if l > maxthick + int(W*frac) or W - r > maxthick + int(W*frac): break
        im = im.crop((l, t, r, b))
    return im
# ---- deskew -----------------------------------------------------------
# Scans are a degree or two out of square: the page edge and any ruled lines
# lean. Estimate the angle by rotating a downsampled copy through a small range
# and taking the angle whose row/column ink profiles are sharpest (a straight
# page concentrates ink into rows and columns, maximising their variance).
# NB rotation resamples — the only step in this tool that does. It is applied
# only when the page is measurably out of square.
def skew_angle(im, limit=2.0, step=0.05):
    # measure on the interior: edge bands and the page frame would otherwise
    # dominate the profile variance and pin the estimate at zero
    W, H = im.size
    g = im.crop((int(W*0.08), int(H*0.06), int(W*0.92), int(H*0.94)))
    g = g.resize((max(g.width//4, 1), max(g.height//4, 1)), Image.BILINEAR)
    best, best_score = 0.0, -1.0
    n = int(limit/step)
    for i in range(-n, n+1):
        deg = round(i*step, 2)
        a = np.array(g.rotate(deg, resample=Image.BILINEAR, fillcolor=255)) < 128
        score = float(a.mean(1).var() + a.mean(0).var())
        if score > best_score: best, best_score = deg, score
    return best

def rule_angle(im):
    """Angle from long near-vertical or near-horizontal rules, when the plate
    has any: compare where a rule sits near one end against the other. More
    sensitive than the profile method for engravings, which have few text
    lines but strong ruled columns."""
    a = np.array(im) < 128
    H, W = a.shape
    def drift(arr, long_axis):
        n = arr.shape[long_axis]
        lo = arr.take(range(int(n*0.10), int(n*0.35)), axis=long_axis)
        hi = arr.take(range(int(n*0.65), int(n*0.90)), axis=long_axis)
        pl, ph = lo.mean(long_axis), hi.mean(long_axis)
        # A rule is several pixels wide, so group contiguous strong lines and
        # compare their CENTRES — matching column to column would pair edge with
        # edge and under-measure the drift.
        def centres(prof):
            idx = [i for i in range(len(prof)) if prof[i] > 0.55]
            if not idx: return []
            out, cur = [], [idx[0]]
            for i in idx[1:]:
                if i - cur[-1] <= 3: cur.append(i)
                else: out.append(sum(cur)/len(cur)); cur = [i]
            out.append(sum(cur)/len(cur))
            return out
        cl, ch = centres(pl), centres(ph)
        if not cl or not ch: return None
        ds = []
        for i in cl:
            j = min(ch, key=lambda k: abs(k-i))
            if abs(j-i) <= 20: ds.append(j-i)
        if not ds: return None
        span = n*0.55
        return math.degrees(math.atan(float(np.median(ds))/span))
    v = drift(a, 0)   # vertical rules: drift measured down the page
    h = drift(a.T, 0) # horizontal rules
    cands = [x for x in (v, h) if x is not None]
    if not cands: return None
    return -cands[0] if abs(cands[0]) >= 0.05 else 0.0

def deskew(im):
    ang = rule_angle(im)
    if ang is None: ang = skew_angle(im)
    if abs(ang) < 0.05: return im, 0.0
    return im.rotate(ang, resample=Image.BICUBIC, expand=True, fillcolor=255), ang

# ---- scanner shadow ---------------------------------------------------------
# A soft grey band down an edge (the gutter shadow) is not dark enough for the
# bar sweep but still prints. Trim edge rows/columns that are mostly dark.
def strip_edge_shadow(im, frac=0.05, dark=0.30):
    # A gutter shadow is often a wedge — dark over part of the edge only — so a
    # whole-column mean misses it. Score each edge column (row) by the darkest
    # tenth of its length instead.
    a = np.array(im) < 150
    W, H = im.size
    l, r, t, b = 0, W, 0, H
    def worst(v, n):           # darkest window of length n along v
        if len(v) < n: return v.mean() if len(v) else 0.0
        c = np.cumsum(np.insert(v, 0, 0.0))
        return float(((c[n:] - c[:-n])/n).max())
    hw, vw = max(H//10, 1), max(W//10, 1)
    # stop at the first light line: only a band touching the edge is shadow,
    # anything past it is the plate's own content
    for i in range(int(W*frac)):
        if worst(a[:, i], hw) > dark: l = i + 1
        else: break
    for i in range(W-1, W-int(W*frac)-1, -1):
        if worst(a[:, i], hw) > dark: r = i
        else: break
    for i in range(int(H*frac)):
        if worst(a[i, :], vw) > dark: t = i + 1
        else: break
    for i in range(H-1, H-int(H*frac)-1, -1):
        if worst(a[i, :], vw) > dark: b = i
        else: break
    return im.crop((l, t, r, b)) if (l, r, t, b) != (0, W, 0, H) else im

if a.raw:
    _skew = 0.0
else:
    im = strip_edge_shadow(im)  # drop the gutter shadow first…
    im, _skew = deskew(im)      # …so it cannot bias the angle estimate
    im = strip_edge_shadow(im)  # …then clear what the rotation brought in

# autocrop to dark content with a small margin
def autocrop(im):
    arr = np.array(im) < 128
    rows = np.where(arr.mean(1) > 0.002)[0]; cols = np.where(arr.mean(0) > 0.002)[0]
    if not (len(rows) and len(cols)): return im
    m = int(0.01*im.width)
    return im.crop((max(cols[0]-m,0), max(rows[0]-m,0),
                    min(cols[-1]+m, im.width), min(rows[-1]+m, im.height)))
if not a.raw:
    im = autocrop(im)          # bring the edges to the ink…
    im = autocrop(strip_edge_bars(im))  # …then sweep off any scanner bar now at an edge
    im = autocrop(strip_edge_shadow(im))  # …and the gutter shadow the autocrop just exposed
if a.rotate:
    im = im.transpose({90: Image.ROTATE_90, 180: Image.ROTATE_180,
                       270: Image.ROTATE_270}[a.rotate])
# --trim T,R,B,L (fractions, after rotation): for skewed scanner edges that are
# too thin for the bar heuristic and too dark for the autocrop to ignore.
if a.trim:
    t, r, b, l = (float(x) for x in a.trim.split(","))
    W, H = im.size
    im = autocrop(im.crop((int(W*l), int(H*t), int(W*(1-r)), int(H*(1-b)))))
uniq = np.unique(np.array(im))
out = f"plates/p{a.page:04d}.png"
if len(uniq) <= 2:           # bitonal source: keep it crisp and small
    im = im.point(lambda v: 255 if v > 128 else 0).convert("1")
im.save(out, optimize=True)
print(out, im.size, im.mode, f"deskew {_skew:+.1f}deg" if _skew else "")