"""Detect top and bottom content bounds in scanned student work.

Strong coloured strokes are never erased merely because they coincide with
paper ruling. The detector supports coloured handwriting and dark ruled scans;
it remains conservative for pencil-only scans and heavily saturated grids.
"""

from functools import lru_cache

import cv2
import numpy as np

from copienator.paper_background import _skew, _ruling, _foreground, _content_mask


@lru_cache(maxsize=8)
def _paper_kernel(sigma: float) -> np.ndarray:
    """Match OpenCV's uint8 Gaussian coefficients, including error diffusion.

    See getGaussianKernelFixedPoint_ED in OpenCV's smooth.dispatch.cpp.
    The 8-bit coefficients allow exact sums in the faster float filter path.
    """
    size = round(sigma*6+1) | 1
    kernel = cv2.getGaussianKernel(size, sigma).ravel()
    fixed = np.zeros(size, np.float32)
    error = 0.0
    for i in range(size//2):
        value = kernel[i]*256+error
        weight = round(value)
        error = value-weight
        fixed[i] = fixed[-1-i] = weight
    fixed[size//2] = 256-fixed.sum()
    fixed /= 256
    return fixed


def _paper_blur(gray: np.ndarray, sigma: float) -> np.ndarray:
    kernel = _paper_kernel(sigma)
    blurred = cv2.sepFilter2D(gray, cv2.CV_32F, kernel, kernel)
    # GaussianBlur rounds positive half-integers upward, rather than to even.
    np.add(blurred, .5, out=blurred)
    np.floor(blurred, out=blurred)
    return blurred.astype(np.uint8)


def _large_blank_ink(gray: np.ndarray, clean: np.ndarray, dpi: float):
    """Refine confirmed ruling, only accepting substantial extra blank margins.

    Short directional openings tolerate locally bent/broken paper lines. A
    physical component-size threshold rejects their remaining tiny fragments.
    This is deliberately limited to already-confirmed ruled paper.
    """
    px = dpi / 25.4
    height, width = gray.shape
    dark = 255-gray
    length = max(9, round(2.5*px))
    tolerance = max(3, round(.6*px) | 1)
    lines = []
    for horizontal in (True, False):
        broadened = cv2.dilate(dark, np.ones(
            (tolerance, 1) if horizontal else (1, tolerance), np.uint8))
        lines.append(cv2.morphologyEx(broadened, cv2.MORPH_OPEN, np.ones(
            (1, length) if horizontal else (length, 1), np.uint8)))
    residual = cv2.subtract(dark, np.maximum(*lines))
    n, labels, stats, centers = cv2.connectedComponentsWithStats(
        (residual > 35).astype(np.uint8), 8)
    keep = np.zeros(n, bool)
    xx = stats[1:, 0]
    ww, hh, area = stats[1:, 2], stats[1:, 3], stats[1:, 4]
    density = area / (ww*hh)
    shortest, longest = np.minimum(ww, hh), np.maximum(ww, hh)
    compact_ink = ((area >= .8*px*px) & (shortest >= .6*px)
                   & (density > .3) & (longest < 4*shortest))
    # Ruling suppression can fragment faint pencil handwriting into sparse,
    # elongated components. Admit those moderately more readily in the central
    # 80% of the sheet. The outermost 9% deliberately uses a stricter filter:
    # punched holes, torn binding edges, and page numbers usually occur there.
    center_x = xx + ww/2
    central = (center_x > width*.1) & (center_x < width*.9)
    central_ink = (central & (area >= .55*px*px) & (shortest >= .45*px)
                   & (density > .18) & (longest < 6*shortest))
    outer = (center_x < width*.09) | (center_x > width*.91)
    outer_ink = (outer & (area >= 1.2*px*px) & (shortest >= .8*px)
                 & (density >= .4) & (longest < 3*shortest))
    keep[1:] = np.where(outer, outer_ink, compact_ink | central_ink)
    # Use side columns as a prior, then require repeated size and alignment. An
    # isolated note in the same column remains eligible to protect the margin.
    candidates = np.flatnonzero(
        ((xx < 15*px) | (xx+ww > width-15*px))
        & (ww > px) & (ww < 9*px) & (hh > px) & (hh < 12*px)) + 1
    if len(candidates) > 128:
        # Bound the matching cost and retain ambiguous, very noisy margins.
        return clean, 0, False
    holes = set()
    for i in candidates:
        matches = []
        for j in candidates:
            if (abs(centers[i, 0]-centers[j, 0]) < 2*px
                    and .6 < stats[j, 2]/stats[i, 2] < 1.6
                    and .6 < stats[j, 3]/stats[i, 3] < 1.6):
                matches.append(j)
        if len(matches) >= 3 and np.ptp(centers[matches, 1]) > height*.35:
            holes.update(matches)
    keep[list(holes)] = False
    refined = keep[labels]
    # Directional opening also removes long fraction bars. Protect very dark,
    # thick straight strokes independently, even if ruling crosses their ends.
    long_strokes = cv2.morphologyEx((gray < 50).astype(np.uint8), cv2.MORPH_OPEN,
                                   np.ones((1, max(9, round(width*.1))), np.uint8))
    count, lab, st, _ = cv2.connectedComponentsWithStats(long_strokes, 8)
    bars = np.zeros(count, bool)
    bw, bh, ba = st[1:, 2], st[1:, 3], st[1:, 4]
    bars[1:] = ((bw > width*.1) & (bh >= .3*px) & (bh < height*.015)
                & (bw > 8*bh) & (ba/(bw*bh) > .5))
    strong_ruling, strong_lines = _ruling((gray < 50).astype(np.uint8)*255, True)
    if strong_lines:
        # A family of equally dark parallel lines is paper, not fraction bars.
        overlap = np.bincount(lab[strong_ruling > 0], minlength=count)
        bars &= overlap < st[:, 4]*.5
    refined |= bars[lab]
    ys = np.flatnonzero(np.any(refined, axis=1))
    original = np.flatnonzero(np.any(clean, axis=1))
    if not len(ys) or not len(original):
        return clean, 0, False
    # Leave ordinary small crops to the more permissive detector. Keep a
    # recovery neighbourhood around the refined bounds for broken/faint strokes.
    top, bottom = max(0, int(ys.min()-2*px)), min(height, int(ys.max()+1+2*px))
    result = clean.copy()
    changed = False
    if top-original.min() >= 30*px:
        result[:top] = 0
        changed = True
    if original.max()+1-bottom >= 30*px:
        result[bottom:] = 0
        changed = True
    return result, len(holes), changed


def _neutral_paper_foreground(gray: np.ndarray, chroma: np.ndarray, dpi: float):
    """Clean confirmed dark ruling before it can seed whole pages.

    Returns None for ordinary ink components or unconfirmed paper geometry.
    The mask is transformed back to the original displayed pixel coordinates.
    """
    height, width = gray.shape
    # Only neutral darkness is relevant here: long blue equations are not
    # evidence of dark paper. Broken ruling can form several medium-sized
    # components instead of one page-spanning component. The broader threshold
    # also admits faded gray grids; periodic ruling must still be confirmed.
    _, _, stats, _ = cv2.connectedComponentsWithStats(
        ((gray < 160) & (chroma < 30)).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 > .1) >= 3):
        return None, dict(paper_cleanup=False)
    angle = _skew(gray, angle_step=.5)
    matrix = cv2.getRotationMatrix2D((width/2, height/2), angle, 1)
    corners = np.array([[0,0],[width,0],[0,height],[width,height]], dtype=float)
    corners = cv2.transform(corners[None], matrix)[0]
    origin = np.floor(corners.min(axis=0))
    size = np.ceil(corners.max(axis=0)-origin).astype(int)
    matrix[:, 2] -= origin
    deskewed = cv2.warpAffine(gray, matrix, tuple(size), borderValue=255)
    background = _paper_blur(deskewed, dpi/8)
    normalized = cv2.divide(deskewed, np.maximum(background,1), scale=255)
    block = max(15, int(dpi/5) | 1)
    binary = cv2.adaptiveThreshold(normalized,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                   cv2.THRESH_BINARY_INV,block,9)
    horizontal, nh = _ruling(binary,True)
    vertical, nv = _ruling(binary,False)
    if not (nh or nv):
        return None, dict(paper_cleanup=False)
    # The permissive mask retains faint, isolated marks.
    # Bands identify where ruling is expected, but only actual dark pixels
    # inside them may be suppressed. Erasing the complete band loses faint
    # writing alongside a dark grid line.
    paper_pixels = (horizontal|vertical) & ((normalized < 160).astype(np.uint8)*255)
    clean, holes = _content_mask(_foreground(normalized,dpi,9,paper_pixels),dpi,nh,nv)
    clean, repeated_holes, refined = _large_blank_ink(deskewed, clean, dpi)
    restored = cv2.warpAffine(clean, cv2.invertAffineTransform(matrix),
                              (width,height), flags=cv2.INTER_NEAREST, borderValue=0)
    return restored > 0, dict(paper_cleanup=True, angle_deg=round(angle,3),
                             horizontal_lines=nh,vertical_lines=nv,
                             edge_artifacts=holes+repeated_holes,
                             large_blank_refinement=refined)


def _seed_mask(mask: np.ndarray, px: float) -> np.ndarray:
    n, labels, stats, _ = cv2.connectedComponentsWithStats(mask.astype(np.uint8), 8)
    keep = np.zeros(n, bool)
    keep[1:] = ((stats[1:, 4] >= max(4, .12*px*px))
                & (stats[1:, 2] >= max(2, round(.35*px)))
                & (stats[1:, 3] >= max(2, round(.35*px))))
    return keep[labels]


def detect_bounds(rgb: np.ndarray, dpi: float = 200, padding_mm: float = 6,
                  min_crop_mm: float = 5) -> dict:
    """Locate strong ink, recover adjacent faint strokes, and retain padding.

    Neutral punched-hole shadows usually have neither sufficient chroma nor
    sufficient darkness to seed a region. Nothing is discarded merely because
    it is in a side margin. Two seed thresholds expose unstable boundaries.
    """
    px = dpi/25.4
    h, w = rgb.shape[:2]
    red, green, blue = cv2.split(rgb)
    lowest = cv2.min(cv2.min(red, green), blue)
    chroma = cv2.subtract(cv2.max(cv2.max(red, green), blue), lowest)
    gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)
    darkness = 255-lowest
    length = max(15, 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))
    # Strong strokes bypass the line-background estimate entirely: even a long
    # isolated black or coloured fraction bar must survive. Local contrast is
    # used only to recover weaker surrounding strokes.
    neutral_ink = gray < 95
    cleaned, paper_info = _neutral_paper_foreground(gray, chroma, dpi)
    if cleaned is not None:
        neutral_ink = cleaned
    weak = ((chroma > 60) | ((gray < 175) & (residual > 25))).astype(np.uint8)
    radius = max(1, round(2*px))
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2*radius+1, 2*radius+1))
    extents = []
    counts = []
    same_thresholds = not np.any((chroma > 100) & (chroma <= 115) & ~neutral_ink)
    for threshold in (100, 115):
        if threshold == 115 and same_thresholds:
            counts.append(counts[0])
            if extents:
                extents.append(extents[0])
            continue
        seeds = _seed_mask((chroma > threshold) | neutral_ink, px)
        counts.append(int(np.count_nonzero(seeds)))
        if not np.any(seeds):
            continue
        # Limit weak recovery to a physical neighbourhood: faint grid lines
        # connected to a letter cannot grow into a full-page foreground mask.
        nearby = cv2.dilate(seeds.astype(np.uint8), kernel)
        candidate = (weak & nearby) | seeds.astype(np.uint8)
        n, labels = cv2.connectedComponents(candidate, 8)
        seeded_labels = np.zeros(n, bool)
        seeded_labels[np.unique(labels[seeds])] = True
        seeded_labels[0] = False
        ys = np.flatnonzero(np.any(seeded_labels[labels], axis=1))
        extents.append((int(ys.min()), int(ys.max())+1))
    result = dict(top_px=0, bottom_px=h, angle_deg=None,
                  horizontal_lines=0, vertical_lines=0, edge_artifacts=0,
                  detector="ink", seed_pixels=counts, status="review-no-ink-seeds")
    result.update(paper_info)
    if not extents:
        return result
    pad = padding_mm*px
    top = max(0, int(np.floor(min(e[0] for e in extents)-pad)))
    bottom = min(h, int(np.ceil(max(e[1] for e in extents)+pad)))
    uncertain = []
    if len(extents) < 2:
        uncertain.append('seed-threshold')
    else:
        for edge, name in ((0, 'top'), (1, 'bottom')):
            if abs(extents[0][edge]-extents[1][edge]) > max(pad, 3*px):
                uncertain.append(name)
    if top < min_crop_mm*px:
        top = 0
    if h-bottom < min_crop_mm*px:
        bottom = h
    result.update(top_px=top, bottom_px=bottom,
                  status=('review-'+'-'.join(uncertain) if uncertain else
                          'cropped' if top or bottom < h else 'unchanged'))
    return result