#!/usr/bin/env python3
"""Find lines of ink on a scanned page that the text layer has no line for.

Engine-independent: it never runs OCR. It binarises the page image, finds
horizontal bands of ink (text lines) by projection, and asks, for each band,
whether any line box in the Internet Archive's own hOCR for that page covers it.
A band with ink across a text-like width and no covering hOCR line is a
CANDIDATE. Candidates are then looked at by eye; the detector decides nothing.

usage: detect.py page.jpg page.hocr [--crops DIR] [--json]
"""
import json, re, sys, os
import numpy as np
from PIL import Image

def hocr_lines(path):
    s = open(path, encoding="utf-8", errors="replace").read()
    page = re.search(r'class="ocr_page"[^>]*title="[^"]*bbox (\d+) (\d+) (\d+) (\d+)', s)
    W, H = int(page.group(3)), int(page.group(4))
    lines = []
    for m in re.finditer(r'<span class="ocr_(?:line|caption|header|textfloat)"[^>]*title="[^"]*?bbox (\d+) (\d+) (\d+) (\d+)', s):
        lines.append(tuple(int(v) for v in m.groups()))
    words = [tuple(int(v) for v in m.groups()) for m in
             re.finditer(r'class="ocrx_word"[^>]*title="[^"]*?bbox (\d+) (\d+) (\d+) (\d+)', s)]
    return W, H, lines, words

def ink_mask(img):
    g = np.asarray(img.convert("L"), dtype=np.float32)
    # Otsu threshold
    hist, _ = np.histogram(g, bins=256, range=(0, 256))
    total = g.size; sum_all = np.dot(np.arange(256), hist)
    wb = 0; sb = 0; best = 0; t = 128
    for i in range(256):
        wb += hist[i]
        if wb == 0: continue
        wf = total - wb
        if wf == 0: break
        sb += i * hist[i]
        mb = sb / wb; mf = (sum_all - sb) / wf
        v = wb * wf * (mb - mf) ** 2
        if v > best: best, t = v, i
    return g < t

def bands(mask, x0, x1, min_h, min_frac):
    """Horizontal ink bands inside columns x0..x1."""
    sub = mask[:, x0:x1]
    prof = sub.sum(axis=1) / max(1, (x1 - x0))
    on = prof > 0.01
    out = []; y = 0; H = len(on)
    while y < H:
        if on[y]:
            y0 = y
            while y < H and on[y]: y += 1
            if y - y0 >= min_h:
                cols = sub[y0:y].any(axis=0)
                xs = np.nonzero(cols)[0]
                span = (xs[-1] - xs[0]) / (x1 - x0) if len(xs) else 0
                # fraction of columns that carry ink: text is broken into glyphs
                fill = cols.mean()
                out.append(dict(y0=int(y0), y1=int(y), xs0=int(x0 + xs[0]), xs1=int(x0 + xs[-1]),
                                span=float(span), fill=float(fill)))
        else:
            y += 1
    return out

def run(img, lines, words, crops=None, base=None, mask=None, W=None, H=None, threshold=0.50):
    """The rule. lines/words are hOCR bboxes in hOCR pixel space (W x H); the
    image may be any scale of the same page."""
    if W is None: W, H = img.width, img.height
    sx, sy = img.width / W, img.height / H
    if mask is None: mask = ink_mask(img)
    # text block: the extent of the hOCR words. Uncovered ink is sought only
    # inside the block the OCR itself found text in, so page edges, gutters and
    # thumb shadows outside it do not count.
    if words:
        bx0 = min(w[0] for w in words) * sx; bx1 = max(w[2] for w in words) * sx
        by0 = min(w[1] for w in words) * sy; by1 = max(w[3] for w in words) * sy
    else:
        bx0, bx1, by0, by1 = 0.1 * img.width, 0.9 * img.width, 0.05 * img.height, 0.95 * img.height
    x0, x1 = int(bx0), int(bx1)
    # the median hOCR line height sets the scale
    hs = sorted((l[3] - l[1]) * sy for l in lines if l[3] > l[1]) or [0.012 * img.height]
    lh = hs[len(hs) // 2]
    # coverage: every hOCR line and word box, padded by a fifth of a line
    pad = int(0.2 * lh)
    cov = np.zeros(mask.shape, bool)
    for (lx0, ly0, lx1, ly1) in list(lines) + list(words):
        cov[max(0, int(ly0*sy) - pad):int(ly1*sy) + pad, max(0, int(lx0*sx) - pad):int(lx1*sx) + pad] = True
    y_lo, y_hi = max(0, int(by0 - lh)), min(img.height, int(by1 + lh))
    unc = (mask & ~cov)[:, x0:x1]
    # rows where uncovered ink is spread across the text block, smoothed over a
    # third of a line so the gaps between glyph rows do not split a line
    k = max(1, int(lh / 3))
    cs = np.vstack([np.zeros((1, unc.shape[1]), np.int32), np.cumsum(unc, axis=0, dtype=np.int32)])
    Hh = unc.shape[0]
    ys = np.arange(Hh)
    lo = np.clip(ys - k, 0, Hh); hi = np.clip(ys + k + 1, 0, Hh)
    win = (cs[hi] - cs[lo]) > 0
    # the uncovered ink must be dense along a run of the row: within some
    # window of `wlen` columns (six line heights, or 15% of the block if that
    # is wider), at least `threshold` of the columns carry uncovered ink. A
    # line dropped inside one newspaper column qualifies as well as a line
    # dropped across a whole book page.
    wlen = int(min(unc.shape[1], max(6 * lh, 0.15 * unc.shape[1]))) if unc.shape[1] else 1
    if unc.shape[1]:
        C = np.hstack([np.zeros((Hh, 1), np.int32), np.cumsum(win, axis=1, dtype=np.int32)])
        Wsum = C[:, wlen:] - C[:, :-wlen]
        spread = Wsum.max(axis=1) / wlen
        argx = Wsum.argmax(axis=1)
    else:
        spread = np.zeros(Hh)
    # raw (unsmoothed) uncovered-ink rows, to tell a rule from a line of text
    raw = unc.mean(axis=1) if unc.shape[1] else np.zeros(Hh)
    hot = spread >= threshold
    cands = []; y = y_lo
    while y < y_hi:
        if hot[y]:
            y0 = y
            while y < y_hi and hot[y]: y += 1
            h = y - y0
            rawh = int((raw[y0:y] * unc.shape[1] >= 0.5 * wlen * 0.05).sum()) if unc.shape[1] else 0
            if h >= 0.5 * lh and rawh >= 0.35 * lh:
                xs = np.nonzero(unc[y0:y].any(axis=0))[0]
                peak = y0 + int(np.argmax(spread[y0:y]))
                wx0 = x0 + int(argx[peak])
                cands.append(dict(y0=int(y0), y1=int(y), xs0=int(x0 + xs[0]), xs1=int(x0 + xs[-1]),
                                  max_spread=round(float(spread[y0:y].max()), 3),
                                  height_over_lh=round(float(h / lh), 2), est_lines=max(1, round(h / lh)), raw_rows=rawh, win=[wx0, wx0 + wlen]))
        else:
            y += 1
    res = dict(line_height=round(float(lh), 1), block=[x0, int(by0), x1, int(by1)], candidates=cands,
               scale=[sx, sy])
    if crops and cands:
        os.makedirs(crops, exist_ok=True)
        rgb = np.asarray(img.convert("RGB")).copy()
        # tint: ink the text layer covers goes grey-blue, ink it does not goes red
        full_unc = np.zeros(mask.shape, bool); full_unc[:, x0:x1] = unc
        rgb[mask & cov] = (rgb[mask & cov] * 0.5 + np.array([40, 60, 140]) * 0.5).astype(np.uint8)
        rgb[full_unc] = (220, 20, 20)
        for i, c in enumerate(cands):
            pad2 = int(2 * lh)
            cx0 = max(0, c["win"][0] - 2 * wlen // 3); cx1 = min(img.width, c["win"][1] + 2 * wlen // 3)
            if cx1 - cx0 < 12 * lh:
                mid = (cx0 + cx1) // 2; cx0 = max(0, int(mid - 6 * lh)); cx1 = min(img.width, int(mid + 6 * lh))
            box = (cx0, max(0, c["y0"] - pad2), cx1, min(img.height, c["y1"] + pad2))
            im = Image.fromarray(rgb[box[1]:box[3], box[0]:box[2]])
            plain = img.crop(box).convert("RGB")
            both = Image.new("RGB", (im.width, im.height * 2 + 6), (255, 255, 255))
            both.paste(plain, (0, 0)); both.paste(im, (0, im.height + 6))
            if both.width > 1200: both = both.resize((1200, max(1, int(both.height * 1200 / both.width))))
            elif both.width < 700: both = both.resize((700, max(1, int(both.height * 700 / both.width))))
            p = os.path.join(crops, f"{base}_c{i}.jpg"); both.save(p, quality=85); c["crop"] = os.path.basename(p)
            c["crop_box"] = list(box)
    return res

def main():
    a = sys.argv[1:]
    img_p, hocr_p = a[0], a[1]
    crops = a[a.index("--crops") + 1] if "--crops" in a else None
    W, H, lines, words = hocr_lines(hocr_p)
    img = Image.open(img_p)
    base = os.path.splitext(os.path.basename(img_p))[0]
    res = run(img, lines, words, crops=crops, base=base, W=W, H=H)
    res.update(image=img_p, hocr=hocr_p, n_hocr_lines=len(lines))
    print(json.dumps(res, indent=1))

if __name__ == "__main__":
    main()