"""Colour-aware suppression of pixels supported by recurring paper geometry.

This only identifies paper ruling. It does not classify pagination or introduce
an independent faint-writing detector.
"""

import cv2
import numpy as np

from copienator.paper_background import _ruling, _runs, _skew


def _maximum(image):
    first, second, third = cv2.split(image)
    return cv2.max(cv2.max(first, second), third)


def _minimum(image):
    first, second, third = cv2.split(image)
    return cv2.min(cv2.min(first, second), third)


def _thick_strokes(binary, dpi):
    """Protect substantial excess thickness without keeping grid junctions."""
    height, width = binary.shape
    widths = []
    for horizontal in (True, False):
        dimension = width if horizontal else height
        length = max(25, round(dimension*.12))
        kernel = np.ones((1,length) if horizontal else (length,1), np.uint8)
        lines = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel)
        counts = np.count_nonzero(lines, axis=1 if horizontal else 0)
        widths.extend(b-a for a,b in _runs(counts > dimension*.3))
    thickness = max(3, int(np.ceil(np.quantile(widths,.95)))+2) if widths else 3
    kernel = np.ones((thickness,thickness), np.uint8)
    cores = cv2.erode(binary, kernel)
    count, labels, stats, _ = cv2.connectedComponentsWithStats(cores, 8)
    keep = np.zeros(count, bool)
    px = dpi/25.4
    x, _, w, h, area = stats[1:].T
    keep[1:] = ((area >= .3*px*px) & (np.maximum(w,h) >= 1.5*px)
                & (np.minimum(w,h) >= 2))
    # A slightly thicker full-height margin rule is still a paper candidate.
    # It will only be removed if its colour also passes the ruling model below.
    side_line = ((h > height*.7) & (w < 2*px)
                 & ((x < width*.3) | (x > width*.7)))
    keep[1:] &= ~side_line
    return cv2.dilate(keep[labels].astype(np.uint8), kernel) > 0


def _edge_connected_family(binary, bands, horizontal):
    """Require several ruling bands reaching a physical edge of the scan.

    Regularly spaced fraction bars inside an answer are not enough evidence.
    """
    dimension = binary.shape[1] if horizontal else binary.shape[0]
    length = max(25, round(dimension*.12))
    lines = cv2.morphologyEx(binary, cv2.MORPH_OPEN,
        np.ones((1,length) if horizontal else (length,1), np.uint8))
    edge = max(3, round(dimension*.08))
    if horizontal:
        touching = np.any(lines[:,:edge] | lines[:,-edge:], axis=1)
        positions = np.any(bands, axis=1)
    else:
        touching = np.any(lines[:edge] | lines[-edge:], axis=0)
        positions = np.any(bands, axis=0)
    return sum(np.any(touching[a:b]) for a,b in _runs(positions)) >= 5


def ruling_suppression(rgb: np.ndarray, dpi: float):
    """Return a mask in original coordinates and evidence for paper cleanup.

    Require a periodic horizontal or vertical family before modelling any line.
    Compare actual RGB pixels with locally persistent line colour, retain darker
    or different-colour strokes, and guard unusually thick strokes. A separate
    full-height margin rule is eligible only alongside that confirmed family.
    """
    height, width = rgb.shape[:2]
    lowest = _minimum(rgb)
    chroma = cv2.subtract(_maximum(rgb), lowest)
    _, _, stats, _ = cv2.connectedComponentsWithStats(
        ((lowest < 215) & ((chroma > 20) | (lowest < 160))).astype(np.uint8), 8)
    spans = np.maximum(stats[1:,2]/width, stats[1:,3]/height)
    if not (np.any(spans > .35) or np.count_nonzero(spans > .05) >= 3):
        return np.zeros((height,width), bool), {}

    angle = _skew(lowest, angle_step=.5)
    matrix = cv2.getRotationMatrix2D((width/2,height/2), angle, 1)
    corners = cv2.transform(np.array([[[0,0],[width,0],[0,height],
                                       [width,height]]], dtype=float), matrix)[0]
    origin = np.floor(corners.min(axis=0))
    size = np.ceil(corners.max(axis=0)-origin).astype(int)
    matrix[:,2] -= origin
    image = cv2.warpAffine(rgb, matrix, tuple(size), borderValue=(255,255,255))
    dark = 255-image
    dark16 = dark.astype(np.int16)
    darkness = _maximum(dark)
    binary = cv2.adaptiveThreshold(_minimum(image), 255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV,
        max(15,int(dpi/5) | 1), 5)
    horizontal, nh = _ruling(binary, True)
    vertical, nv = _ruling(binary, False)
    if not ((nh and _edge_connected_family(binary,horizontal,True))
            or (nv and _edge_connected_family(binary,vertical,False))):
        return np.zeros((height,width), bool), {}

    thick = _thick_strokes(binary, dpi)
    removed = np.zeros(binary.shape, bool)
    protected = np.zeros(binary.shape, bool)
    for is_horizontal, bands, count in ((True,horizontal,nh),(False,vertical,nv)):
        length = max(25, round(5.2*dpi/25.4))
        tolerance = max(3, round(.8*dpi/25.4) | 1)
        transverse = np.ones((tolerance,1) if is_horizontal else (1,tolerance), np.uint8)
        along = np.ones((1,length) if is_horizontal else (length,1), np.uint8)
        # Small transverse tolerance follows local scan warp. The RGB model is
        # still tested against actual pixels; it does not erase entire bands.
        broadened = cv2.dilate(dark, transverse)
        model = cv2.morphologyEx(broadened, cv2.MORPH_OPEN, along)
        support = _maximum(model) > 9
        if count:
            support &= bands > 0
        close = _maximum(np.abs(dark16-model.astype(np.int16))) < 40
        sample = support & (darkness > 9) & close & ~thick
        # Neutral, red, green and blue line families are modelled separately so
        # a red margin rule cannot be mistaken for unusually dark blue writing.
        first, second, third = cv2.split(model)
        classes = np.full(binary.shape, 2, np.uint8)
        classes[second <= third] = 1
        classes[(first <= second) & (first <= third)] = 0
        classes[cv2.subtract(_maximum(model),_minimum(model)) < 25] = 3
        credible = np.zeros(binary.shape, bool)
        for colour in range(4):
            selected = sample & (classes == colour)
            if selected.sum() < 100:
                continue
            yy, xx = np.nonzero(selected)
            spread = np.ptp(yy if is_horizontal else xx)
            axis_size = image.shape[0] if is_horizontal else image.shape[1]
            margin_rule = (not is_horizontal and np.ptp(yy) > image.shape[0]*.7
                and (np.median(xx) < image.shape[1]*.3 or np.median(xx) > image.shape[1]*.7))
            if spread < axis_size*.35 and not margin_rule:
                continue
            # An isolated long fraction bar must not define its own paper colour.
            # The upper quantile admits darker major rules within an established
            # family while rejecting excess ink on most of its pixels.
            ceiling = np.floor(np.quantile(dark[selected],.99,axis=0)+25).astype(np.int16)
            credible |= ((classes == colour)
                         & (_maximum(dark16-ceiling) <= 0))
        match = close & credible & (darkness > 9) & ~thick
        removed |= support & match
        protected |= support & ~match & (binary > 0)

    inverse = cv2.invertAffineTransform(matrix)
    result = cv2.warpAffine(removed.astype(np.uint8), inverse,
                           (width,height), flags=cv2.INTER_LINEAR) > 0
    guard = cv2.warpAffine(protected.astype(np.uint8), inverse,
                          (width,height), flags=cv2.INTER_LINEAR) > 0
    # Cover one-pixel deskew/interpolation remnants, except next to writing that
    # failed the paper test. Confirmed pixels of the other orientation remain
    # removable at line crossings.
    kernel = np.ones((3,3), np.uint8)
    result |= ((cv2.dilate(result.astype(np.uint8),kernel) > 0)
               & ~(cv2.dilate(guard.astype(np.uint8),kernel) > 0))
    return result, dict(ruling_cleanup=True, ruling_angle_deg=round(angle,3),
                       ruling_horizontal_lines=nh, ruling_vertical_lines=nv)