#!/usr/bin/env python3
"""
Detect rectangular form cells in government arrest report images
using morphological operations to isolate horizontal and vertical lines,
then find closed rectangular cells from line intersections.

Iteration 5 (final):
- Two-pass detection: row-band pass (consecutive H lines) + multi-row pass
- Interior-line filter prevents spanning-cell false positives
- Semi-transparent overlay visualization with cell count labels
- Tuned parameters for government arrest report forms
"""

import cv2
import numpy as np
from pathlib import Path

# Paths
BASE_DIR = Path(__file__).parent
DEBUG_DIR = BASE_DIR / "debug"
DEBUG_DIR.mkdir(exist_ok=True)

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

# =============================================================================
# Parameters
# =============================================================================
ADAPTIVE_BLOCK_SIZE = 15
ADAPTIVE_C = 10

H_KERNEL_WIDTH = 40
H_KERNEL_HEIGHT = 1
V_KERNEL_WIDTH = 1
V_KERNEL_HEIGHT = 25

H_DILATE_KERNEL = (5, 1)
V_DILATE_KERNEL = (1, 5)

MIN_H_LINE_LENGTH = 25
MIN_V_LINE_LENGTH = 12
MAX_LINE_THICKNESS = 12

SNAP_TOLERANCE = 8

MIN_CELL_WIDTH = 20
MIN_CELL_HEIGHT = 12
MAX_CELL_WIDTH = 2000
MAX_CELL_HEIGHT = 400   # Allows multi-row cells but not full-page spans


def save_debug(name, step, suffix, image):
    path = DEBUG_DIR / f"{name}_{step}_{suffix}.png"
    cv2.imwrite(str(path), image)
    print(f"  Saved: {path.name}")


def extract_line_segments(binary_image, orientation='horizontal',
                          min_length=30, max_thickness=12):
    contours, _ = cv2.findContours(binary_image, cv2.RETR_EXTERNAL,
                                     cv2.CHAIN_APPROX_SIMPLE)
    segments = []
    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        if orientation == 'horizontal':
            if w >= min_length and h <= max_thickness:
                cy = y + h // 2
                segments.append((x, cy, x + w, cy))
        else:
            if h >= min_length and w <= max_thickness:
                cx = x + w // 2
                segments.append((cx, y, cx, y + h))
    return segments


def snap_coordinates(values, tolerance):
    if not values:
        return {}
    sorted_vals = sorted(set(values))
    clusters = []
    current_cluster = [sorted_vals[0]]
    for v in sorted_vals[1:]:
        if v - current_cluster[-1] <= tolerance:
            current_cluster.append(v)
        else:
            clusters.append(current_cluster)
            current_cluster = [v]
    clusters.append(current_cluster)

    mapping = {}
    for cluster in clusters:
        mean_val = int(round(np.mean(cluster)))
        for v in cluster:
            mapping[v] = mean_val
    return mapping


def merge_ranges(ranges, gap_tolerance=5):
    if not ranges:
        return []
    sorted_ranges = sorted(ranges)
    merged = [list(sorted_ranges[0])]
    for start, end in sorted_ranges[1:]:
        if start <= merged[-1][1] + gap_tolerance:
            merged[-1][1] = max(merged[-1][1], end)
        else:
            merged.append([start, end])
    return [tuple(r) for r in merged]


def range_covered(ranges, start, end, tolerance):
    for r_start, r_end in ranges:
        if r_start <= start + tolerance and r_end >= end - tolerance:
            return True
    return False


def range_touches(ranges, point, tolerance):
    for r_start, r_end in ranges:
        if r_start - tolerance <= point <= r_end + tolerance:
            return True
    return False


def range_overlaps_interior(ranges, start, end, margin=5):
    interior_start = start + margin
    interior_end = end - margin
    if interior_start >= interior_end:
        return False
    for r_start, r_end in ranges:
        overlap_start = max(r_start, interior_start)
        overlap_end = min(r_end, interior_end)
        if overlap_end - overlap_start > margin:
            return True
    return False


def deduplicate_cells(cells, tolerance=5):
    if not cells:
        return []
    cells = sorted(cells, key=lambda c: (c[2]-c[0]) * (c[3]-c[1]))
    kept = []
    for cell in cells:
        is_dup = False
        for existing in kept:
            if (abs(cell[0] - existing[0]) <= tolerance and
                abs(cell[1] - existing[1]) <= tolerance and
                abs(cell[2] - existing[2]) <= tolerance and
                abs(cell[3] - existing[3]) <= tolerance):
                is_dup = True
                break
        if not is_dup:
            kept.append(cell)
    return kept


def build_line_structures(h_segments, v_segments, snap_tol):
    """Snap coordinates, build merged line lookup structures."""
    h_ys = [s[1] for s in h_segments]
    v_xs = [s[0] for s in v_segments]
    h_x_starts = [s[0] for s in h_segments]
    h_x_ends = [s[2] for s in h_segments]
    v_y_starts = [s[1] for s in v_segments]
    v_y_ends = [s[3] for s in v_segments]

    all_ys = h_ys + v_y_starts + v_y_ends
    all_xs = v_xs + h_x_starts + h_x_ends

    y_snap = snap_coordinates(all_ys, snap_tol)
    x_snap = snap_coordinates(all_xs, snap_tol)

    snapped_h = []
    for x1, y, x2, _ in h_segments:
        sy = y_snap.get(y, y)
        sx1 = x_snap.get(x1, x1)
        sx2 = x_snap.get(x2, x2)
        snapped_h.append((sx1, sy, sx2, sy))

    snapped_v = []
    for x, y1, _, y2 in v_segments:
        sx = x_snap.get(x, x)
        sy1 = y_snap.get(y1, y1)
        sy2 = y_snap.get(y2, y2)
        snapped_v.append((sx, sy1, sx, sy2))

    h_lines = {}
    for x1, y, x2, _ in snapped_h:
        h_lines.setdefault(y, []).append((min(x1, x2), max(x1, x2)))
    v_lines = {}
    for x, y1, _, y2 in snapped_v:
        v_lines.setdefault(x, []).append((min(y1, y2), max(y1, y2)))

    for y in h_lines:
        h_lines[y] = merge_ranges(h_lines[y])
    for x in v_lines:
        v_lines[x] = merge_ranges(v_lines[x])

    unique_ys = sorted(h_lines.keys())
    unique_xs = sorted(v_lines.keys())

    return h_lines, v_lines, unique_ys, unique_xs


def find_cells_from_segments(h_segments, v_segments, snap_tol,
                              min_w, min_h, max_w, max_h, img_width, img_height):
    h_lines, v_lines, unique_ys, unique_xs = build_line_structures(
        h_segments, v_segments, snap_tol
    )

    print(f"  Grid: {len(unique_ys)} horizontal lines, {len(unique_xs)} vertical lines")
    edge_tol = snap_tol + 2
    interior_margin = 5
    all_cells = []

    # =========================================================================
    # Pass 1: Row-bands (consecutive horizontal line pairs)
    # =========================================================================
    for i in range(len(unique_ys) - 1):
        y_top = unique_ys[i]
        y_bot = unique_ys[i + 1]
        band_h = y_bot - y_top
        if band_h < min_h or band_h > max_h:
            continue

        # Find vertical lines spanning this band
        band_xs = []
        for x in unique_xs:
            v_ranges = v_lines.get(x, [])
            if (range_touches(v_ranges, y_top, edge_tol) and
                range_touches(v_ranges, y_bot, edge_tol)):
                band_xs.append(x)

        # Consecutive vertical pairs form cells
        for j in range(len(band_xs) - 1):
            x_left = band_xs[j]
            x_right = band_xs[j + 1]
            cell_w = x_right - x_left
            if cell_w < min_w or cell_w > max_w:
                continue
            if not range_covered(h_lines.get(y_top, []), x_left, x_right, edge_tol):
                continue
            if not range_covered(h_lines.get(y_bot, []), x_left, x_right, edge_tol):
                continue
            all_cells.append((x_left, y_top, x_right, y_bot))

    print(f"  Pass 1 (row-bands): {len(all_cells)} cells")

    # =========================================================================
    # Pass 2: Multi-row cells (span 2+ consecutive H lines)
    # Only kept if no interior H or V lines cross through them.
    # =========================================================================
    multi_row_cells = []
    for i, y_top in enumerate(unique_ys):
        for j in range(i + 2, len(unique_ys)):
            y_bot = unique_ys[j]
            cell_h = y_bot - y_top
            if cell_h < min_h or cell_h > max_h:
                continue

            for k, x_left in enumerate(unique_xs):
                for l in range(k + 1, len(unique_xs)):
                    x_right = unique_xs[l]
                    cell_w = x_right - x_left
                    if cell_w < min_w or cell_w > max_w:
                        continue

                    # All 4 sides must exist
                    if not range_covered(h_lines.get(y_top, []), x_left, x_right, edge_tol):
                        continue
                    if not range_covered(h_lines.get(y_bot, []), x_left, x_right, edge_tol):
                        continue
                    v_left = v_lines.get(x_left, [])
                    if not (range_touches(v_left, y_top, edge_tol) and
                            range_touches(v_left, y_bot, edge_tol)):
                        continue
                    v_right = v_lines.get(x_right, [])
                    if not (range_touches(v_right, y_top, edge_tol) and
                            range_touches(v_right, y_bot, edge_tol)):
                        continue

                    # Interior check: reject if any H line spans the cell width
                    has_interior = False
                    for y in unique_ys:
                        if y <= y_top + interior_margin or y >= y_bot - interior_margin:
                            continue
                        if range_covered(h_lines.get(y, []), x_left, x_right, edge_tol):
                            has_interior = True
                            break

                    if not has_interior:
                        # Check interior vertical lines
                        for x in unique_xs:
                            if x <= x_left + interior_margin or x >= x_right - interior_margin:
                                continue
                            v_ranges = v_lines.get(x, [])
                            if range_overlaps_interior(v_ranges, y_top, y_bot, margin=interior_margin):
                                has_interior = True
                                break

                    if not has_interior:
                        multi_row_cells.append((x_left, y_top, x_right, y_bot))

    print(f"  Pass 2 (multi-row): {len(multi_row_cells)} cells")
    all_cells.extend(multi_row_cells)

    # Deduplicate
    all_cells = list(set(all_cells))
    all_cells = deduplicate_cells(all_cells, tolerance=snap_tol)

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

    return all_cells


def draw_cells_overlay(img, cells, name):
    """Draw cells with semi-transparent fill and numbered labels."""
    cell_img = img.copy()
    overlay = img.copy()
    colors = [
        (0, 200, 0),    # Green
        (200, 0, 0),    # Blue
        (0, 140, 255),  # Orange
        (200, 0, 200),  # Magenta
        (0, 220, 220),  # Yellow
        (220, 220, 0),  # Cyan
        (128, 0, 200),  # Purple
        (0, 128, 128),  # Teal
    ]

    for i, (x1, y1, x2, y2) in enumerate(cells):
        color = colors[i % len(colors)]
        cv2.rectangle(overlay, (x1, y1), (x2, y2), color, -1)

    # Blend
    alpha = 0.20
    cell_img = cv2.addWeighted(overlay, alpha, cell_img, 1 - alpha, 0)

    # Draw outlines and labels
    for i, (x1, y1, x2, y2) in enumerate(cells):
        color = colors[i % len(colors)]
        cv2.rectangle(cell_img, (x1, y1), (x2, y2), color, 2)

        # Small label with cell number
        label = str(i)
        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = 0.35
        thickness = 1
        (tw, th), _ = cv2.getTextSize(label, font, font_scale, thickness)
        # Position label in top-left corner of cell
        lx = x1 + 2
        ly = y1 + th + 2
        # Only draw label if cell is big enough
        if (x2 - x1) > tw + 4 and (y2 - y1) > th + 4:
            cv2.rectangle(cell_img, (lx - 1, ly - th - 1), (lx + tw + 1, ly + 1),
                         (255, 255, 255), -1)
            cv2.putText(cell_img, label, (lx, ly), font, font_scale, (0, 0, 0), thickness)

    return cell_img


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

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

    h, w = img.shape[:2]
    print(f"  Image size: {w} x {h}")

    # Step 1: Grayscale
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    save_debug(name, "01", "gray", gray)

    # Step 2: Adaptive binary threshold
    binary = cv2.adaptiveThreshold(
        gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C,
        cv2.THRESH_BINARY_INV, ADAPTIVE_BLOCK_SIZE, ADAPTIVE_C
    )
    save_debug(name, "02", "binary", binary)

    # Step 3: Horizontal lines
    h_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (H_KERNEL_WIDTH, H_KERNEL_HEIGHT))
    h_lines_img = cv2.morphologyEx(binary, cv2.MORPH_OPEN, h_kernel)
    h_dilate = cv2.getStructuringElement(cv2.MORPH_RECT, H_DILATE_KERNEL)
    h_lines_img = cv2.dilate(h_lines_img, h_dilate, iterations=1)
    save_debug(name, "03", "h_lines", h_lines_img)

    # Step 4: Vertical lines
    v_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (V_KERNEL_WIDTH, V_KERNEL_HEIGHT))
    v_lines_img = cv2.morphologyEx(binary, cv2.MORPH_OPEN, v_kernel)
    v_dilate = cv2.getStructuringElement(cv2.MORPH_RECT, V_DILATE_KERNEL)
    v_lines_img = cv2.dilate(v_lines_img, v_dilate, iterations=1)
    save_debug(name, "04", "v_lines", v_lines_img)

    # Step 5: Combined
    combined = cv2.add(h_lines_img, v_lines_img)
    save_debug(name, "05", "combined_lines", combined)

    # Step 6: Extract line segments
    h_segments = extract_line_segments(h_lines_img, 'horizontal',
                                        MIN_H_LINE_LENGTH, MAX_LINE_THICKNESS)
    v_segments = extract_line_segments(v_lines_img, 'vertical',
                                        MIN_V_LINE_LENGTH, MAX_LINE_THICKNESS)

    print(f"  Found {len(h_segments)} horizontal segments, {len(v_segments)} vertical segments")

    # Draw segments
    seg_img = np.ones_like(img) * 255
    for x1, y1, x2, y2 in h_segments:
        cv2.line(seg_img, (x1, y1), (x2, y2), (255, 0, 0), 1)
    for x1, y1, x2, y2 in v_segments:
        cv2.line(seg_img, (x1, y1), (x2, y2), (0, 0, 255), 1)
    save_debug(name, "06", "segments", seg_img)

    # Step 7: Find cells
    cells = find_cells_from_segments(
        h_segments, v_segments, SNAP_TOLERANCE,
        MIN_CELL_WIDTH, MIN_CELL_HEIGHT, MAX_CELL_WIDTH, MAX_CELL_HEIGHT,
        w, h
    )

    print(f"  Detected {len(cells)} cells")

    # Draw cells overlay
    cell_img = draw_cells_overlay(img, cells, name)
    save_debug(name, "07", "cells", cell_img)

    # Stats
    if cells:
        widths = [c[2] - c[0] for c in cells]
        heights = [c[3] - c[1] for c in cells]
        print(f"  Cell width range: {min(widths)} - {max(widths)} px")
        print(f"  Cell height range: {min(heights)} - {max(heights)} px")
        small = sum(1 for w in widths if w < 100)
        medium = sum(1 for w in widths if 100 <= w <= 500)
        large = sum(1 for w in widths if w > 500)
        print(f"  Size distribution: {small} small (<100px), {medium} medium, {large} large (>500px)")

    return cells


if __name__ == "__main__":
    print("Form Cell Detection via Morphological Line Isolation (v5)")
    print(f"Parameters:")
    print(f"  Adaptive threshold: block={ADAPTIVE_BLOCK_SIZE}, C={ADAPTIVE_C}")
    print(f"  H kernel: {H_KERNEL_WIDTH}x{H_KERNEL_HEIGHT}, dilate {H_DILATE_KERNEL}")
    print(f"  V kernel: {V_KERNEL_WIDTH}x{V_KERNEL_HEIGHT}, dilate {V_DILATE_KERNEL}")
    print(f"  Min line lengths: H={MIN_H_LINE_LENGTH}, V={MIN_V_LINE_LENGTH}")
    print(f"  Snap tolerance: {SNAP_TOLERANCE}")
    print(f"  Cell size range: {MIN_CELL_WIDTH}-{MAX_CELL_WIDTH} x {MIN_CELL_HEIGHT}-{MAX_CELL_HEIGHT}")

    all_cells = {}
    for name, path in IMAGES:
        cells = process_image(name, path)
        if cells:
            all_cells[name] = cells

    print(f"\n{'='*60}")
    print("SUMMARY")
    print(f"{'='*60}")
    for name, cells in all_cells.items():
        print(f"  {name}: {len(cells)} cells detected")
