#!/usr/bin/env python3
"""
Form cell detection via line extension + contour approach (v8).

Based on v7 with OCR-based cell merging: if an OCR text element spans
2+ horizontally adjacent cells, merge those cells. This handles form
fields where a single text entry crosses an internal divider line.

Requires searchable PDFs (generated by convert_to_pdf.py).

Debug outputs written to temp/boxes/debug8.
"""

from __future__ import annotations

import sys
from pathlib import Path
from typing import List, Tuple

import cv2
import numpy as np

# Add project root so we can import natural_pdf
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))


BASE_DIR = Path(__file__).parent
DEBUG_DIR = BASE_DIR / "debug8"
DEBUG_DIR.mkdir(exist_ok=True)

IMAGES = [
    ("boxes-1", BASE_DIR / "boxes-1.png", BASE_DIR / "boxes-1.pdf"),
    ("boxes-2", BASE_DIR / "boxes-2.png", BASE_DIR / "boxes-2.pdf"),
]

# Parameters
TARGET_WIDTH = 2200

ADAPTIVE_BLOCK = 15
ADAPTIVE_C = 10

MIN_LINE_AREA_H = 100
MIN_LINE_AREA_V = 100  # Moderate — contextual filter handles the rest

# Extension amount: how far to extend each line beyond its detected bounds
LINE_EXTEND_PX = 15

# Contextual filter: V-line endpoints must be within this many pixels of an H-line
V_CONTEXT_TOLERANCE = 12

# Hybrid context: V-lines taller than this only need ONE endpoint near H-line;
# shorter ones need BOTH (short = more likely text fragment)
V_HEIGHT_BOTH_THRESHOLD = 40  # normalized px — below this, require both endpoints

# Contour filtering
MIN_CELL_AREA = 350
MAX_CELL_AREA_RATIO = 0.40
MIN_CELL_W = 45
MIN_CELL_H = 18
MAX_ASPECT = 50.0
MIN_RECTANGULARITY = 0.35

BORDER_MARGIN = 6
DEDUP_IOU = 0.85


def save_debug(name: str, step: str, suffix: str, img: np.ndarray) -> None:
    path = DEBUG_DIR / f"{name}_{step}_{suffix}.png"
    cv2.imwrite(str(path), img)
    print(f"  Saved: {path.name}")


def resize_to_target(img: np.ndarray, target_w: int):
    h, w = img.shape[:2]
    if target_w <= 0 or abs(w - target_w) < 10:
        return img, 1.0
    scale = target_w / float(w)
    new_h = int(round(h * scale))
    interp = cv2.INTER_AREA if scale < 1 else cv2.INTER_CUBIC
    return cv2.resize(img, (target_w, new_h), interpolation=interp), scale


def clamp_int(v, lo, hi):
    return int(max(lo, min(hi, round(v))))


def remove_small_components(mask, min_area):
    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask, connectivity=8)
    out = np.zeros_like(mask)
    for i in range(1, num_labels):
        if stats[i, cv2.CC_STAT_AREA] >= min_area:
            out[labels == i] = 255
    return out


def extract_and_extend_lines(line_mask, orientation, extend_px, min_length,
                              max_thickness=20, h_line_mask=None, context_tol=8,
                              both_threshold=40):
    """
    Extract line contours, get bounding boxes, EXTEND along primary axis,
    return list of (x, y, w, h) extended bounding boxes.

    For vertical lines: hybrid contextual filter using the H-line pixel mask.
    - Short V-lines (h < both_threshold) must touch H-line pixels at BOTH
      endpoints — these are likely text fragments.
    - Taller V-lines only need ONE endpoint — these are likely real dividers
      that may not perfectly reach both H-lines.
    """
    contours, _ = cv2.findContours(line_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    extended = []
    img_h = line_mask.shape[0]
    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        if orientation == "h":
            if w < min_length or h > max_thickness:
                continue
            x_new = max(0, x - extend_px)
            w_new = w + 2 * extend_px
            extended.append((x_new, y, w_new, h))
        else:
            if h < min_length or w > max_thickness:
                continue
            # Hybrid contextual filter using H-line pixel mask
            if h_line_mask is not None:
                cx = x + w // 2
                top_y = y
                bot_y = min(y + h, img_h - 1)
                x_lo = max(0, cx - context_tol)
                x_hi = min(h_line_mask.shape[1], cx + context_tol + 1)
                top_window = h_line_mask[max(0, top_y - context_tol):top_y + context_tol + 1, x_lo:x_hi]
                bot_window = h_line_mask[max(0, bot_y - context_tol):min(img_h, bot_y + context_tol + 1), x_lo:x_hi]
                near_top = np.any(top_window > 0) if top_window.size > 0 else False
                near_bot = np.any(bot_window > 0) if bot_window.size > 0 else False
                if h < both_threshold:
                    # Short lines: require BOTH endpoints (likely text if not)
                    if not (near_top and near_bot):
                        continue
                else:
                    # Tall lines: require at least ONE endpoint
                    if not (near_top or near_bot):
                        continue
            y_new = max(0, y - extend_px)
            h_new = h + 2 * extend_px
            extended.append((x, y_new, w, h_new))
    return extended


def redraw_lines(shape, h_boxes, v_boxes):
    """Redraw extended line bounding boxes onto a clean mask."""
    mask = np.zeros(shape[:2], dtype=np.uint8)
    for x, y, w, h in h_boxes:
        cv2.rectangle(mask, (x, y), (x + w, y + h), 255, -1)
    for x, y, w, h in v_boxes:
        cv2.rectangle(mask, (x, y), (x + w, y + h), 255, -1)
    return mask


def iou(a, b):
    ax1, ay1, ax2, ay2 = a
    bx1, by1, bx2, by2 = b
    ox = max(0, min(ax2, bx2) - max(ax1, bx1))
    oy = max(0, min(ay2, by2) - max(ay1, by1))
    inter = ox * oy
    if inter <= 0:
        return 0.0
    area_a = (ax2 - ax1) * (ay2 - ay1)
    area_b = (bx2 - bx1) * (by2 - by1)
    return float(inter) / float(max(area_a + area_b - inter, 1))


def dedup_boxes(boxes, thresh):
    if not boxes:
        return []
    boxes = sorted(boxes, key=lambda b: (b[2] - b[0]) * (b[3] - b[1]))
    keep = []
    for b in boxes:
        if all(iou(b, k) <= thresh for k in keep):
            keep.append(b)
    return keep


def find_cells_from_contours(inverse, img_area):
    contours, _ = cv2.findContours(inverse, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)
    out = []
    max_area = img_area * MAX_CELL_AREA_RATIO
    for cnt in contours:
        area = cv2.contourArea(cnt)
        if area < MIN_CELL_AREA or area > max_area:
            continue
        x, y, w, h = cv2.boundingRect(cnt)
        if w < MIN_CELL_W or h < MIN_CELL_H:
            continue
        aspect = max(w, h) / max(min(w, h), 1)
        if aspect > MAX_ASPECT:
            continue
        rectangularity = float(area) / float(max(w * h, 1))
        if rectangularity < MIN_RECTANGULARITY:
            continue
        out.append((x, y, x + w, y + h))
    return dedup_boxes(out, DEDUP_IOU)


def draw_cells(img, cells):
    out = img.copy()
    overlay = img.copy()
    colors = [
        (0, 200, 0), (200, 0, 0), (0, 140, 255), (200, 0, 200),
        (0, 220, 220), (220, 220, 0), (128, 0, 200), (0, 128, 128),
    ]
    for i, (x1, y1, x2, y2) in enumerate(cells):
        cv2.rectangle(overlay, (x1, y1), (x2, y2), colors[i % len(colors)], -1)
    out = cv2.addWeighted(overlay, 0.18, out, 0.82, 0)
    for i, (x1, y1, x2, y2) in enumerate(cells):
        color = colors[i % len(colors)]
        cv2.rectangle(out, (x1, y1), (x2, y2), color, 2)
        label = str(i)
        font = cv2.FONT_HERSHEY_SIMPLEX
        (tw, th), _ = cv2.getTextSize(label, font, 0.35, 1)
        if (x2 - x1) > tw + 6 and (y2 - y1) > th + 6:
            lx, ly = x1 + 2, y1 + th + 2
            cv2.rectangle(out, (lx - 1, ly - th - 1), (lx + tw + 1, ly + 1), (255, 255, 255), -1)
            cv2.putText(out, label, (lx, ly), font, 0.35, (0, 0, 0), 1)
    return out


def get_ocr_boxes(pdf_path):
    """Load OCR text bounding boxes from a searchable PDF via natural-pdf.
    Returns list of (x1, y1, x2, y2) in pixel coordinates matching the
    original image dimensions."""
    from natural_pdf import PDF

    pdf = PDF(str(pdf_path))
    page = pdf.pages[0]
    pw, ph = page.width, page.height

    # Get the original image dimensions to compute scale
    # PDF coordinates are in points; we need pixel coordinates
    elements = page.find_all("text")
    boxes = []
    for el in elements:
        text = el.extract_text().strip()
        if not text:
            continue
        boxes.append((el.x0, el.top, el.x1, el.bottom, text))
    pdf.close()
    return boxes, pw, ph


def boxes_overlap_h(a, b, y_tolerance=0.5):
    """Check if two cells are on the same row (vertically overlapping)
    and horizontally adjacent."""
    ax1, ay1, ax2, ay2 = a
    bx1, by1, bx2, by2 = b
    # Check vertical overlap: they share the same row
    ah = ay2 - ay1
    bh = by2 - by1
    min_h = min(ah, bh)
    overlap_y = min(ay2, by2) - max(ay1, by1)
    if overlap_y < min_h * y_tolerance:
        return False
    return True


def cells_same_row(a, b, tolerance=5):
    """Check if two cells are on the same row: same y0 AND same y1.
    Cells formed by the same pair of horizontal lines share exact top/bottom.
    Parent cells spanning multiple rows will NOT match."""
    return abs(a[1] - b[1]) <= tolerance and abs(a[3] - b[3]) <= tolerance


def text_overlaps_cell(text_box, cell, min_px=2):
    """Check if a text bounding box overlaps a cell."""
    tx1, ty1, tx2, ty2 = text_box
    cx1, cy1, cx2, cy2 = cell
    ox = min(tx2, cx2) - max(tx1, cx1)
    oy = min(ty2, cy2) - max(ty1, cy1)
    return ox >= min_px and oy >= min_px


def merge_cells_by_ocr(cells, ocr_boxes, img_w, img_h, pdf_w, pdf_h):
    """Merge cells that share the same row AND overlap the same OCR text element.

    Simple logic: if cells share the same y and the same text element
    crosses both of them, they should be one cell.

    ocr_boxes: list of (x1, y1, x2, y2, text) in PDF coordinates
    cells: list of (x1, y1, x2, y2) in pixel coordinates

    Returns merged cell list.
    """
    sx = img_w / pdf_w
    sy = img_h / pdf_h

    scaled_ocr = []
    for ox1, oy1, ox2, oy2, text in ocr_boxes:
        scaled_ocr.append((ox1 * sx, oy1 * sy, ox2 * sx, oy2 * sy, text))

    n = len(cells)
    parent = list(range(n))

    def find(i):
        while parent[i] != i:
            parent[i] = parent[parent[i]]
            i = parent[i]
        return i

    def union(i, j):
        ri, rj = find(i), find(j)
        if ri != rj:
            parent[ri] = rj

    # Debug: dump cells and OCR for a specific row
    debug_y_min, debug_y_max = img_h * 0.62, img_h * 0.72  # IBR/UCR area
    debug_cells = [(i, cells[i]) for i in range(n) if cells[i][1] >= debug_y_min and cells[i][1] <= debug_y_max]
    if debug_cells:
        print(f"  DEBUG IBR/UCR row (y={debug_y_min:.0f}..{debug_y_max:.0f}):")
        for i, (cx1, cy1, cx2, cy2) in debug_cells:
            print(f"    Cell {i}: x={cx1}..{cx2} y={cy1}..{cy2}")
        for tx1, ty1, tx2, ty2, text in scaled_ocr:
            if ty1 >= debug_y_min - 20 and ty1 <= debug_y_max + 20:
                overlaps = [i for i, c in debug_cells if text_overlaps_cell((tx1, ty1, tx2, ty2), c)]
                if overlaps:
                    same_row_groups = []
                    for a in overlaps:
                        for b in overlaps:
                            if a < b and cells_same_row(cells[a], cells[b]):
                                same_row_groups.append((a, b))
                    print(f"    OCR '{text[:50]}' x={tx1:.0f}..{tx2:.0f} y={ty1:.0f}..{ty2:.0f} -> cells {overlaps}, same_row_pairs={same_row_groups}")

    merge_count = 0
    for tx1, ty1, tx2, ty2, text in scaled_ocr:
        # Find all cells this text overlaps
        overlapping = []
        for i in range(n):
            if text_overlaps_cell((tx1, ty1, tx2, ty2), cells[i]):
                overlapping.append(i)

        if len(overlapping) < 2:
            continue

        # Merge pairs with same y0/y1 (same row = same horizontal lines)
        for a_pos in range(len(overlapping)):
            for b_pos in range(a_pos + 1, len(overlapping)):
                ai = overlapping[a_pos]
                bi = overlapping[b_pos]
                if cells_same_row(cells[ai], cells[bi]):
                    if find(ai) != find(bi):
                        union(ai, bi)
                        merge_count += 1

    # Build merged cells from union-find groups
    from collections import defaultdict
    groups = defaultdict(list)
    for i in range(n):
        groups[find(i)].append(i)

    merged = []
    for group_indices in groups.values():
        xs1 = [cells[i][0] for i in group_indices]
        ys1 = [cells[i][1] for i in group_indices]
        xs2 = [cells[i][2] for i in group_indices]
        ys2 = [cells[i][3] for i in group_indices]
        merged.append((min(xs1), min(ys1), max(xs2), max(ys2)))

    merged.sort(key=lambda c: (c[1], c[0]))
    print(f"  OCR merge: {merge_count} pair merges, {len(cells)} -> {len(merged)} cells")
    return merged


def remove_nested_cells(cells, containment_thresh=0.85):
    """Remove cells that are mostly contained within other cells.
    Keep only the innermost (leaf) cells — if cell A is inside cell B,
    drop B (the parent) since its children cover its area."""
    n = len(cells)
    if n == 0:
        return cells

    # For each cell, check if it contains other cells
    is_parent = [False] * n
    for i in range(n):
        ax1, ay1, ax2, ay2 = cells[i]
        a_area = (ax2 - ax1) * (ay2 - ay1)
        if a_area <= 0:
            continue
        for j in range(n):
            if i == j:
                continue
            bx1, by1, bx2, by2 = cells[j]
            b_area = (bx2 - bx1) * (by2 - by1)
            if b_area <= 0 or b_area >= a_area:
                continue
            # How much of cell j is inside cell i?
            ox = max(0, min(ax2, bx2) - max(ax1, bx1))
            oy = max(0, min(ay2, by2) - max(ay1, by1))
            overlap = ox * oy
            if overlap / b_area >= containment_thresh:
                # Cell j is mostly inside cell i → cell i is a parent
                is_parent[i] = True
                break

    kept = [c for c, p in zip(cells, is_parent) if not p]
    removed = n - len(kept)
    if removed > 0:
        print(f"  Removed {removed} parent/nested cells, {len(kept)} leaf cells remain")
    return kept


def process_image(name, image_path, pdf_path=None):
    print(f"\n{'='*60}")
    print(f"Processing: {name}")
    print(f"{'='*60}")

    img = cv2.imread(str(image_path))
    if img is None:
        print(f"  ERROR: Could not load {image_path}")
        return []

    # Resize
    img_norm, scale = resize_to_target(img, TARGET_WIDTH)
    h, w = img_norm.shape[:2]
    img_area = h * w
    k = w / 2200.0
    print(f"  Original: {img.shape[1]}x{img.shape[0]} -> Normalized: {w}x{h}")

    # Auto-scale kernels
    h_long = clamp_int(40 * k, 35, 80)
    h_short = clamp_int(25 * k, 18, 45)
    v_long = clamp_int(30 * k, 18, 55)
    v_short = clamp_int(18 * k, 12, 35)
    h_close_w = clamp_int(15 * k, 10, 30)
    v_close_h = clamp_int(10 * k, 6, 18)
    extend = clamp_int(LINE_EXTEND_PX * k, 6, 20)
    print(f"  Kernels: H={h_long}/{h_short}, V={v_long}/{v_short}, extend={extend}px")

    # Step 1: Binary
    gray = cv2.cvtColor(img_norm, cv2.COLOR_BGR2GRAY)
    binary = cv2.adaptiveThreshold(
        gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY_INV, ADAPTIVE_BLOCK, ADAPTIVE_C
    )
    save_debug(name, "01", "binary", binary)

    # Step 2: Horizontal lines (two-pass + gap close)
    hk_long = cv2.getStructuringElement(cv2.MORPH_RECT, (h_long, 1))
    hk_short = cv2.getStructuringElement(cv2.MORPH_RECT, (h_short, 1))
    h_lines = cv2.bitwise_or(
        cv2.morphologyEx(binary, cv2.MORPH_OPEN, hk_long),
        cv2.morphologyEx(binary, cv2.MORPH_OPEN, hk_short),
    )
    h_close = cv2.getStructuringElement(cv2.MORPH_RECT, (h_close_w, 1))
    h_lines = cv2.morphologyEx(h_lines, cv2.MORPH_CLOSE, h_close)
    h_lines = remove_small_components(h_lines, MIN_LINE_AREA_H)
    save_debug(name, "02", "h_lines", h_lines)

    # Step 3: Vertical lines (two-pass only — tiny kernel picks up text)
    vk_long = cv2.getStructuringElement(cv2.MORPH_RECT, (1, v_long))
    vk_short = cv2.getStructuringElement(cv2.MORPH_RECT, (1, v_short))
    v_lines = cv2.bitwise_or(
        cv2.morphologyEx(binary, cv2.MORPH_OPEN, vk_long),
        cv2.morphologyEx(binary, cv2.MORPH_OPEN, vk_short),
    )
    v_close = cv2.getStructuringElement(cv2.MORPH_RECT, (1, v_close_h))
    v_lines = cv2.morphologyEx(v_lines, cv2.MORPH_CLOSE, v_close)
    v_lines = remove_small_components(v_lines, MIN_LINE_AREA_V)
    save_debug(name, "03", "v_lines", v_lines)

    # Step 4: Extract line bounding boxes and EXTEND along primary axis
    min_h_len = clamp_int(20 * k, 15, 60)
    min_v_len = clamp_int(20 * k, 12, 40)

    # First extract H-lines (no context filter needed — H isolation is clean)
    h_boxes = extract_and_extend_lines(h_lines, "h", extend, min_h_len)

    # Extract V-lines with hybrid contextual filter
    context_tol = clamp_int(V_CONTEXT_TOLERANCE * k, 6, 16)
    both_thresh = clamp_int(V_HEIGHT_BOTH_THRESHOLD * k, 25, 60)
    v_boxes = extract_and_extend_lines(v_lines, "v", extend, min_v_len,
                                        h_line_mask=h_lines, context_tol=context_tol,
                                        both_threshold=both_thresh)
    print(f"  Extended lines: {len(h_boxes)} H, {len(v_boxes)} V (context_tol={context_tol})")

    # Step 5: Redraw extended lines -> watertight grid
    watertight = redraw_lines(img_norm.shape, h_boxes, v_boxes)
    save_debug(name, "04", "watertight_grid", watertight)

    # Step 6: Invert -> cell regions
    inverse = cv2.bitwise_not(watertight)
    m = BORDER_MARGIN
    inverse[:m, :] = 0
    inverse[-m:, :] = 0
    inverse[:, :m] = 0
    inverse[:, -m:] = 0
    save_debug(name, "05", "inverse", inverse)

    # Step 7: Find cells from contours
    cells = find_cells_from_contours(inverse, img_area)
    print(f"  Detected cells: {len(cells)}")

    # Scale back to original coordinates
    if scale != 1.0:
        inv_s = 1.0 / scale
        cells = [(int(round(x1*inv_s)), int(round(y1*inv_s)),
                  int(round(x2*inv_s)), int(round(y2*inv_s)))
                 for x1, y1, x2, y2 in cells]

    # Sort top-to-bottom, left-to-right
    cells.sort(key=lambda c: (c[1], c[0]))

    # Draw pre-merge result
    cell_img = draw_cells(img, cells)
    save_debug(name, "06", "cells_pre_merge", cell_img)

    # Step 8: OCR-based cell merging
    if pdf_path and pdf_path.exists():
        ocr_boxes, pdf_w, pdf_h = get_ocr_boxes(pdf_path)
        orig_h, orig_w = img.shape[:2]
        print(f"  OCR boxes: {len(ocr_boxes)} text elements")

        # Debug: draw OCR boxes on image to visualize alignment
        ocr_debug = img.copy()
        sx_dbg = orig_w / pdf_w
        sy_dbg = orig_h / pdf_h
        for ox1, oy1, ox2, oy2, text in ocr_boxes:
            px1 = int(ox1 * sx_dbg)
            py1 = int(oy1 * sy_dbg)
            px2 = int(ox2 * sx_dbg)
            py2 = int(oy2 * sy_dbg)
            cv2.rectangle(ocr_debug, (px1, py1), (px2, py2), (0, 0, 255), 1)
        # Also draw cell borders
        for cx1, cy1, cx2, cy2 in cells:
            cv2.rectangle(ocr_debug, (cx1, cy1), (cx2, cy2), (0, 255, 0), 2)
        save_debug(name, "06b", "ocr_vs_cells", ocr_debug)

        cells = merge_cells_by_ocr(cells, ocr_boxes, orig_w, orig_h, pdf_w, pdf_h)

        # Draw post-merge result
        cell_img = draw_cells(img, cells)
        save_debug(name, "07", "cells_merged", cell_img)
    else:
        print(f"  No PDF found for OCR merge — skipping")

    if cells:
        ws = [c[2] - c[0] for c in cells]
        hs = [c[3] - c[1] for c in cells]
        print(f"  Width range: {min(ws)}..{max(ws)} px")
        print(f"  Height range: {min(hs)}..{max(hs)} px")

    return cells


if __name__ == "__main__":
    print("Form Cell Detection via Line Extension + Contours (v8) + OCR Merge")
    print(f"Debug dir: {DEBUG_DIR}\n")
    results = {}
    for name, img_path, pdf_path in IMAGES:
        if img_path.exists():
            cells = process_image(name, img_path, pdf_path)
            results[name] = cells
        else:
            print(f"  Missing: {img_path}")
    print(f"\n{'='*60}")
    print("SUMMARY")
    for name, cells in results.items():
        print(f"  {name}: {len(cells)} cells")
