"""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.edge_marks import classify_edge_marks
from copienator.paper_ruling import ruling_suppression, _edge_connected_family
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,
                              *, diagnostics: dict | None = None):
    """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 and _edge_connected_family(binary,horizontal,True))
            or (nv and _edge_connected_family(binary,vertical,False))):
        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)
    stages = {} if diagnostics is not None else None
    foreground = _foreground(normalized,dpi,9,paper_pixels, diagnostics=stages)
    clean, holes = _content_mask(foreground,dpi,nh,nv, diagnostics=stages)
    before_refinement = clean.copy() if diagnostics is not None else None
    clean, repeated_holes, refined = _large_blank_ink(deskewed, clean, dpi)
    if diagnostics is not None:
        stages['confirmed_ruling'] = paper_pixels > 0
        stages['refinement_removed'] = (before_refinement > 0) & (clean == 0)
        inverse = cv2.invertAffineTransform(matrix)
        for name, mask in stages.items():
            diagnostics[name] = cv2.warpAffine(mask.astype(np.uint8), inverse,
                (width,height), flags=cv2.INTER_NEAREST, borderValue=0) > 0
    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, *, paper_nearby=None,
               diagnostics=None) -> 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))))
    if paper_nearby is not None:
        w, h = stats[:,2], stats[:,3]
        thin = max(2, round(.45*px))
        aligned = ((np.minimum(w,h) <= thin)
                   & (np.maximum(w,h) >= 3*np.minimum(w,h)))
        overlap = np.bincount(labels[paper_nearby], minlength=n)
        remnants = keep & aligned & (overlap >= .8*stats[:,4])
        keep[remnants] = False
        if diagnostics is not None:
            diagnostics['ruling_residuals'] |= remnants[labels]
    return keep[labels]


def detect_bounds(rgb: np.ndarray, dpi: float = 200, padding_mm: float = 6,
                  min_crop_mm: float = 5, *, diagnostics: dict | None = None,
                  colour_cleanup: bool = False, edge_mark_cleanup: bool = True) -> 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.

    Small isolated side-margin marks near the physical top/bottom edge are
    excluded using same-page position and separation. Marks near text remain.
    edge_mark_cleanup=False is available for evaluation comparisons.

    Experimental colour cleanup is opt-in because its full-resolution model
    is too expensive for the default PDF processing path.

    Optional diagnostics are written to a separate caller-owned dictionary.
    Its masks use the original image coordinates and never influence decisions
    or enter the JSON-serializable return value.
    """
    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))
    # The directional estimate alone cannot veto a strong fraction bar. Only
    # independently confirmed paper geometry may suppress a strong candidate.
    neutral_ink = gray < 95
    if diagnostics is not None:
        diagnostics.clear()
        diagnostics['masks'] = {}
        diagnostics['thresholds'] = []
    masks = diagnostics['masks'] if diagnostics is not None else None
    cleaned, paper_info = _neutral_paper_foreground(gray, chroma, dpi, diagnostics=masks)
    if cleaned is not None:
        # Adaptive neutral cleanup includes faint paper fragments. Such pixels
        # need local contrast or substantial darkness before they seed a crop.
        neutral_ink = cleaned & ((residual > 25) | (gray < 95))
    if colour_cleanup:
        suppressed, ruling_info = ruling_suppression(rgb, dpi)
    else:
        suppressed, ruling_info = np.zeros((h,w), bool), {}
    paper_info.update(ruling_info)
    neutral_ink &= ~suppressed
    paper_nearby = (cv2.dilate(suppressed.astype(np.uint8),
                    np.ones((3,3),np.uint8)) > 0) if np.any(suppressed) else None
    if masks is not None:
        masks['paper_suppressed'] = masks.get('paper_suppressed', np.zeros((h,w), bool)) | suppressed
        masks['confirmed_ruling'] = masks.get('confirmed_ruling', np.zeros((h,w), bool)) | suppressed
        masks['ruling_residuals'] = np.zeros((h,w), bool)
        masks['isolated_edge_marks'] = np.zeros((h,w), bool)
    weak = ((chroma > 60) | ((gray < 175) & (residual > 25))).astype(np.uint8)
    weak[suppressed] = 0
    if masks is not None:
        masks['strong_seeds'] = np.zeros((h,w), bool)
        masks['foreground'] = np.zeros((h,w), bool)
        masks['weak_candidates'] = weak > 0
    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])
            if diagnostics is not None:
                diagnostics['thresholds'].append(dict(
                    diagnostics['thresholds'][0], threshold=115, reused=True))
            continue
        seeds = _seed_mask(((chroma > threshold) | neutral_ink) & ~suppressed, px,
                           paper_nearby=paper_nearby, diagnostics=masks)
        counts.append(int(np.count_nonzero(seeds)))
        if not np.any(seeds):
            if diagnostics is not None:
                diagnostics['thresholds'].append(dict(threshold=threshold,
                    seed_pixels=0, extent=None, limiting_components={'top': [], 'bottom': []}))
            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)
        if edge_mark_cleanup:
            n, labels, stats, _ = cv2.connectedComponentsWithStats(candidate, 8)
        else:
            n, labels = cv2.connectedComponents(candidate, 8)
        seeded_labels = np.zeros(n, bool)
        seeded_labels[np.unique(labels[seeds])] = True
        seeded_labels[0] = False
        removed, groups = (classify_edge_marks(labels, stats, seeded_labels, rgb, dpi)
                           if edge_mark_cleanup else ([], []))
        if removed:
            remaining = seeded_labels.copy()
            remaining[removed] = False
            # Keep the original ink bounds if no foreground would remain.
            if np.any(remaining):
                if masks is not None:
                    discarded = np.zeros(n,bool)
                    discarded[removed] = True
                    masks['isolated_edge_marks'] |= discarded[labels]
                seeded_labels = remaining
        retained = seeded_labels[labels]
        if masks is not None:
            masks['strong_seeds'] |= seeds & retained
        ys = np.flatnonzero(np.any(retained, axis=1))
        extents.append((int(ys.min()), int(ys.max())+1))
        if diagnostics is not None:
            masks['foreground'] |= retained
            diagnostics['thresholds'].append(_explain_extent(retained, seeds,
                chroma, neutral_ink, threshold, extents[-1], px))
            diagnostics['thresholds'][-1]['edge_mark_groups'] = groups
    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 diagnostics is not None:
        masks['recovered_weak'] = masks['foreground'] & ~masks['strong_seeds']
        masks['unretained_weak'] = masks['weak_candidates'] & ~masks['foreground']
        for name in ('confirmed_ruling', 'paper_suppressed', 'rejected_artifacts',
                     'repeated_holes', 'refinement_removed'):
            masks.setdefault(name, np.zeros((h,w), bool))
        diagnostics['shape'] = [h,w]
        diagnostics['mask_pixels'] = {name: int(mask.sum()) for name, mask in masks.items()}
    if not extents:
        if diagnostics is not None:
            diagnostics['edges'] = {edge: {'reason': 'no-reliable-ink-seeds',
                'boundary_px': 0 if edge == 'top' else h} for edge in ('top', 'bottom')}
        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)
    proposed_top, proposed_bottom = top, bottom
    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'))
    if diagnostics is not None:
        diagnostics['edges'] = {}
        for edge, boundary, proposed, extent in (
                ('top', top, proposed_top, min(e[0] for e in extents)),
                ('bottom', bottom, proposed_bottom, max(e[1] for e in extents))):
            removal = boundary if edge == 'top' else h-boundary
            reason = 'crop-to-retained-ink-with-padding' if removal else (
                'below-minimum-crop' if boundary != proposed else 'ink-or-padding-reaches-edge')
            diagnostics['edges'][edge] = dict(boundary_px=boundary,
                proposed_boundary_px=proposed, ink_extent_px=extent,
                reason=reason, threshold_disagreement=edge in uncertain,
                limiting_thresholds=[i['threshold'] for i in diagnostics['thresholds']
                    if i['extent'] is not None and i['extent'][0 if edge == 'top' else 1] == extent])
    return result


def _explain_extent(retained, seeds, chroma, neutral, threshold, extent, px):
    """Describe actual boundary-setting regions; infer no handwriting semantics."""
    _, labels, stats, _ = cv2.connectedComponentsWithStats(retained.astype(np.uint8), 8)
    limiting = {}
    for edge, row in (('top', extent[0]), ('bottom', extent[1]-1)):
        ids = [int(i) for i in np.unique(labels[row]) if i]
        components = []
        for identifier in ids:
            x,y,w,h,area = map(int, stats[identifier])
            region = labels[y:y+h,x:x+w] == identifier
            strong = seeds[y:y+h,x:x+w] & region
            components.append(dict(component_id=identifier, bbox_px=[x,y,x+w,y+h],
                bbox_mm=[round(v/px,3) for v in (x,y,x+w,y+h)], area_px=area,
                seed_pixels=int(strong.sum()),
                coloured_seed_pixels=int((strong & (chroma[y:y+h,x:x+w] > threshold)).sum()),
                neutral_seed_pixels=int((strong & neutral[y:y+h,x:x+w]).sum()),
                reason='retained-component-touches-' + edge + '-ink-extent'))
        limiting[edge] = components
    return dict(threshold=threshold, seed_pixels=int(seeds.sum()),
                extent=list(extent), limiting_components=limiting)