Copies/copienator/left_margin.py

"""Conservative body-text margin suggestions for the cutleft reviewer."""
from __future__ import annotations

import cv2
import numpy as np
from PIL import Image

from copienator.ink_detection import detect_bounds

ANALYSIS_DPI = 100
WIDTH_TOLERANCE = .08
# Copie01's name/class block reaches about 27% of the full page height.
# Start below it with clearance, including on later pages with repeated names.
BODY_TOP = .30
BODY_BOTTOM = .82
LABEL_BOTTOM = .92
LABEL_LEFT_FRACTION = .05
COHERENT_MARGIN_TOLERANCE_MM = 10
LABEL_MARGIN_GAP_MM = 10
DISCARD_LEFT_FRACTION = .10
STANDARD_MARGIN_MM = 40
MARGIN_CLEARANCE_MM = 20
MIN_MARGIN_MM = MARGIN_CLEARANCE_MM


def trimmed_left_edge(starts: list[float]) -> float | None:
    """Discard the leftmost 10% of occupied bands; retain the next minimum."""
    if not starts:
        return None
    ordered = sorted(starts)
    return round(ordered[int(len(ordered) * DISCARD_LEFT_FRACTION)], 1)


def handwriting_mask(rgb: np.ndarray) -> np.ndarray:
    """Intersect rognage's foreground with a directional ruling filter."""
    px = ANALYSIS_DPI / 25.4
    gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)
    darkness = 255 - rgb.min(axis=2)
    length = round(5.2 * px) | 1
    horizontal = cv2.morphologyEx(darkness, cv2.MORPH_OPEN, np.ones((1, length), np.uint8))
    vertical = cv2.morphologyEx(darkness, cv2.MORPH_OPEN, np.ones((length, 1), np.uint8))
    residual = cv2.subtract(darkness, np.maximum(horizontal, vertical))
    binary = ((residual > 32) & (gray < 200)).astype(np.uint8)
    count, labels, stats, _ = cv2.connectedComponentsWithStats(binary, 8)
    keep = np.zeros(count, bool)
    for index, (_x, _y, width, height, area) in enumerate(stats[1:], 1):
        keep[index] = (area >= max(5, .35 * px * px) and height >= .5 * px
                       and width >= .25 * px and height <= 6 * width and width <= 15 * height)
    diagnostics = {}
    detect_bounds(rgb, ANALYSIS_DPI, diagnostics=diagnostics)
    return keep[labels] & diagnostics['masks']['foreground']


def occupied_band_starts(mask: np.ndarray, bottom_fraction: float) -> list[float]:
    """Ignore empty bands and writing starting beyond the page midpoint."""
    height, width = mask.shape
    band_height = max(1, round(4 * ANALYSIS_DPI / 25.4))
    bottom = round(bottom_fraction * height)
    starts = []
    for y in range(round(BODY_TOP * height), bottom, band_height):
        band = mask[y:min(y + band_height, bottom)]
        xs = np.flatnonzero(band.sum(axis=0) >= 2)
        if band.sum() >= 40 and len(xs) >= 12 and xs[0] <= width / 2:
            starts.append(float(xs[0]) * 25.4 / ANALYSIS_DPI)
    return starts


def body_edge(mask: np.ndarray) -> dict:
    """Keep the body decile; collect a wider 30–92% label-search window."""
    starts = occupied_band_starts(mask, BODY_BOTTOM)
    return {
        'edge_mm': trimmed_left_edge(starts) if len(starts) >= 6 else None,
        'bands': len(starts),
        'band_starts_mm': starts,
        'label_band_starts_mm': occupied_band_starts(mask, LABEL_BOTTOM),
    }


def estimate_page(image: Image.Image, render_dpi: int) -> dict:
    scale = ANALYSIS_DPI / render_dpi
    small = image.convert('RGB').resize(
        (max(1, round(image.width * scale)), max(1, round(image.height * scale))),
        Image.Resampling.LANCZOS,
    )
    return body_edge(handwriting_mask(np.asarray(small)))


def format_groups(images: list[Image.Image], reference_width: float) -> list[list[int]]:
    """Group similar widths; keep secondary narrow subpages in separate groups."""
    groups: list[list[int]] = []
    for index, image in enumerate(images):
        narrow = image.width < .8 * reference_width
        for group in groups:
            width = float(np.median([images[i].width for i in group]))
            if (images[group[0]].width < .8 * reference_width) == narrow and abs(image.width - width) <= width * WIDTH_TOLERANCE:
                group.append(index)
                break
        else:
            groups.append([index])
    return groups


def analyse_margins(images: list[Image.Image], render_dpi: int, reference_width: float) -> dict:
    """Estimate body margins, then retain left labels supported across coherent pages."""
    starts = [0] * len(images)
    groups = []
    page_estimates = {}

    def page_estimate(index):
        if index not in page_estimates:
            page_estimates[index] = estimate_page(images[index], render_dpi)
        return page_estimates[index]

    for indices in format_groups(images, reference_width):
        # Spread samples through the copy rather than relying on the header page.
        sample = sorted({indices[0], indices[min(1, len(indices) - 1)],
                         indices[len(indices) // 2], indices[-1]})
        estimates = {index: page_estimate(index) for index in sample}
        values = [item['edge_mm'] for item in estimates.values() if item['edge_mm'] is not None]
        # Trim across the copy's format group, not separately on each page:
        # one marginal comment must not veto the other pages' shared margin.
        band_starts = [start for item in estimates.values() for start in item['band_starts_mm']]
        status = 'insufficient-ink'
        edge = trimmed_left_edge(band_starts)
        left_mm = 0.0
        if len(values) >= min(2, len(indices)):
            status = 'near-edge' if edge < MIN_MARGIN_MM else 'margin'
            left_mm = (0.0 if edge < MIN_MARGIN_MM else
                       max(edge / 2, edge - MARGIN_CLEARANCE_MM))
        # Use ALL coherent pages, not just the four body-margin samples.
        # A sparse/low/outlying page cannot supply evidence for this format's
        # established body margin. Its normal fallback behaviour remains below.
        coherent = [index for index in indices
                    if edge is not None and page_estimate(index)['edge_mm'] is not None
                    and page_estimate(index)['edge_mm'] >= MIN_MARGIN_MM
                    and abs(page_estimate(index)['edge_mm'] - edge) <= COHERENT_MARGIN_TOLERANCE_MM]
        label_starts = [start for index in coherent for start in
                        page_estimate(index).get('label_band_starts_mm',
                                                 page_estimate(index)['band_starts_mm'])]
        label_edge = (round(sorted(label_starts)[int(len(label_starts) * LABEL_LEFT_FRACTION)], 1)
                      if label_starts else None)
        labels_in_margin = (status == 'margin' and len(coherent) >= min(2, len(indices))
                            and label_edge is not None and edge - label_edge >= LABEL_MARGIN_GAP_MM)
        body_left_mm = left_mm
        if labels_in_margin:
            status = 'labels-in-margin'
            left_mm = 0.0
        # Secondary subpages can start in the middle of an answer. Preserve
        # their left edge even when the only visible writing is indented.
        if images[indices[0]].width < .8 * reference_width:
            status = 'narrow-page'
            left_mm = 0.0
        start = round(left_mm * render_dpi / 25.4)
        for index in indices:
            starts[index] = start
        groups.append({
            'pages': [index + 1 for index in indices],
            'status': status,
            'edge_mm': edge,
            'left_mm': round(left_mm, 1),
            'body_left_mm': round(body_left_mm, 1),
            'labels_in_margin': labels_in_margin,
            'coherent_pages': [index + 1 for index in coherent],
            'label_edge_mm': label_edge,
            'label_bands': len(label_starts),
            'bands': len(band_starts),
            'discarded_bands': int(len(band_starts) * DISCARD_LEFT_FRACTION),
            'samples': [{'page': index + 1, **item} for index, item in estimates.items()],
        })
    # Inspect actual neighbours, including pages outside the format samples.
    # Blank/uncertain pages do not establish or extend a no-margin run.
    no_margin_run = 0
    no_margin_from_page = None
    for index in range(len(images)):
        edge = page_estimate(index)['edge_mm']
        no_margin_run = no_margin_run + 1 if edge is not None and edge < MIN_MARGIN_MM else 0
        if no_margin_run >= 2:
            run_start = index - 1
            no_margin_from_page = run_start + 1
            starts[run_start:] = [0] * (len(images) - run_start)
            break

    pages = []
    for index in range(len(images)):
        estimate = page_estimate(index)
        edge = estimate['edge_mm']
        pages.append({'page': index + 1, **estimate,
                      'status': 'unknown' if edge is None else
                      ('low' if edge < MIN_MARGIN_MM else 'normal')})

    return {'version': 9, 'pages': pages, 'starts_px': starts, 'groups': groups,
            'no_margin_from_page': no_margin_from_page,
            'discard_left_fraction': DISCARD_LEFT_FRACTION, 'margin_clearance_mm': MARGIN_CLEARANCE_MM,
            'cut_fraction_of_margin': .5, 'cut_rule': 'maximum',
            'standard_margin_mm': STANDARD_MARGIN_MM, 'minimum_margin_mm': MIN_MARGIN_MM,
            'body_height_fraction': [BODY_TOP, BODY_BOTTOM],
            'label_height_fraction': [BODY_TOP, LABEL_BOTTOM],
            'label_left_fraction': LABEL_LEFT_FRACTION,
            'coherent_margin_tolerance_mm': COHERENT_MARGIN_TOLERANCE_MM,
            'label_margin_gap_mm': LABEL_MARGIN_GAP_MM,
            'maximum_band_start_width_fraction': .5}