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

Based on Gemini 3.1 Pro's key insight: instead of dilating lines (which
thickens them), EXTEND their bounding boxes along the primary axis.
This creates a watertight grid without distorting cell boundaries.

Pipeline:
1. Adaptive threshold -> binary
2. Multi-pass morphological opening -> isolate H and V lines
3. Extract line contours and get their bounding boxes
4. Extend each line's bbox along its primary axis by N pixels
5. Redraw extended lines onto a clean mask -> watertight grid
6. Invert -> contours -> cells

Key advantage over v2: lines are lengthened (not thickened), so they
reliably meet at intersections without inflating cell boundaries.

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

from __future__ import annotations

from pathlib import Path
from typing import List, Tuple

import cv2
import numpy as np


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

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

# 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

# Contour filtering
MIN_CELL_AREA = 350
MAX_CELL_AREA_RATIO = 0.40
MIN_CELL_W = 12
MIN_CELL_H = 10
MAX_ASPECT = 30.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):
    """
    Extract line contours, get bounding boxes, EXTEND along primary axis,
    return list of (x, y, w, h) extended bounding boxes.

    For vertical lines: optionally filter by context using the H-line pixel mask.
    A V-line's endpoints must have actual H-line pixels nearby, not just be near
    any y-coordinate. This removes text fragments while keeping real dividers.
    """
    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
            # 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)
                # Check if H-line pixels exist near endpoints in a small window
                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 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 process_image(name, image_path):
    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(8 * k, 5, 30)

    # 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 contextual filter: endpoints must touch actual H-line pixels
    context_tol = clamp_int(V_CONTEXT_TOLERANCE * k, 6, 16)
    v_boxes = extract_and_extend_lines(v_lines, "v", extend, min_v_len,
                                        h_line_mask=h_lines, context_tol=context_tol)
    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 result
    cell_img = draw_cells(img, cells)
    save_debug(name, "06", "cells", cell_img)

    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 (v4)")
    print(f"Debug dir: {DEBUG_DIR}\n")
    results = {}
    for name, path in IMAGES:
        if path.exists():
            cells = process_image(name, path)
            results[name] = cells
        else:
            print(f"  Missing: {path}")
    print(f"\n{'='*60}")
    print("SUMMARY")
    for name, cells in results.items():
        print(f"  {name}: {len(cells)} cells")
