"""
Cell detection via line mask + contour approach.
Version 4 (final) - refined parameters after 3 iterations.

Approach:
1. Adaptive threshold to binary
2. Two-pass morphological opening to isolate H and V lines
3. Close small gaps within line segments
4. Remove noise (tiny connected components)
5. Directional dilation to connect lines at intersections
6. Additional uniform dilation to seal remaining gaps
7. Invert the grid mask -> white cells on black grid
8. Find contours -> each contour = a cell candidate
9. Filter by area, dimensions, aspect ratio, rectangularity
10. IoU-based dedup to remove near-identical bounding boxes
"""

import cv2
import numpy as np
from pathlib import Path

INPUT_DIR = Path(__file__).parent
DEBUG_DIR = INPUT_DIR / "debug2"
DEBUG_DIR.mkdir(exist_ok=True)

# Parameters (tuned over 3 iterations)
PARAMS = {
    # Two-pass morphological kernel sizes for line isolation
    "h_kernel_long": 40,     # first pass: strong horizontal lines
    "h_kernel_short": 25,    # second pass: shorter horizontal segments
    "v_kernel_long": 30,     # first pass: strong vertical lines
    "v_kernel_short": 15,    # second pass: shorter vertical dividers
    # Close small gaps within individual line segments
    "h_close_width": 15,     # horizontal gap closing
    "v_close_height": 10,    # vertical gap closing
    # Noise removal: min connected component area for line fragments
    "min_line_area": 100,
    # Directional dilation to help lines meet at intersections
    "h_dilate_v": 4,         # dilate H lines in vertical direction
    "v_dilate_h": 4,         # dilate V lines in horizontal direction
    # Additional uniform dilation to close remaining gaps
    "grid_dilate_size": 2,
    "grid_dilate_iter": 1,
    # Contour filtering
    "min_cell_area": 400,
    "max_cell_area_ratio": 0.35,  # fraction of total image area
    "min_cell_width": 10,
    "min_cell_height": 10,
    "max_aspect_ratio": 25,
    "min_rectangularity": 0.4,
    # Adaptive threshold
    "adaptive_block": 15,
    "adaptive_c": 10,
    # Border margin (zero out edges to avoid outer contour)
    "border_margin": 5,
    # Dedup IoU threshold
    "dedup_iou": 0.85,
}


def remove_small_components(mask, min_area):
    """Remove connected components smaller than min_area from a binary mask."""
    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask, connectivity=8)
    cleaned = np.zeros_like(mask)
    for i in range(1, num_labels):
        if stats[i, cv2.CC_STAT_AREA] >= min_area:
            cleaned[labels == i] = 255
    return cleaned


def process_image(img_path: Path, name: str, params: dict):
    print(f"\n{'='*60}")
    print(f"Processing: {name} ({img_path})")
    print(f"{'='*60}")

    img = cv2.imread(str(img_path))
    if img is None:
        print(f"ERROR: Could not read {img_path}")
        return []
    h, w = img.shape[:2]
    total_area = h * w
    print(f"Image size: {w}x{h}")

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # ---- Step 1: Binary threshold ----
    binary = cv2.adaptiveThreshold(
        gray, 255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY_INV,
        params["adaptive_block"],
        params["adaptive_c"],
    )
    cv2.imwrite(str(DEBUG_DIR / f"{name}_01_binary.png"), binary)
    print(f"  [1] Binary threshold")

    # ---- Step 2: Isolate horizontal lines (two-pass) ----
    h_k_long = cv2.getStructuringElement(cv2.MORPH_RECT, (params["h_kernel_long"], 1))
    h_k_short = cv2.getStructuringElement(cv2.MORPH_RECT, (params["h_kernel_short"], 1))
    h_lines = cv2.bitwise_or(
        cv2.morphologyEx(binary, cv2.MORPH_OPEN, h_k_long),
        cv2.morphologyEx(binary, cv2.MORPH_OPEN, h_k_short),
    )
    # Close gaps within H lines
    h_close = cv2.getStructuringElement(cv2.MORPH_RECT, (params["h_close_width"], 1))
    h_lines = cv2.morphologyEx(h_lines, cv2.MORPH_CLOSE, h_close)
    h_lines = remove_small_components(h_lines, params["min_line_area"])
    cv2.imwrite(str(DEBUG_DIR / f"{name}_02_h_lines.png"), h_lines)
    print(f"  [2] H lines (kernels: {params['h_kernel_long']}/{params['h_kernel_short']})")

    # ---- Step 3: Isolate vertical lines (two-pass) ----
    v_k_long = cv2.getStructuringElement(cv2.MORPH_RECT, (1, params["v_kernel_long"]))
    v_k_short = cv2.getStructuringElement(cv2.MORPH_RECT, (1, params["v_kernel_short"]))
    v_lines = cv2.bitwise_or(
        cv2.morphologyEx(binary, cv2.MORPH_OPEN, v_k_long),
        cv2.morphologyEx(binary, cv2.MORPH_OPEN, v_k_short),
    )
    # Close gaps within V lines
    v_close = cv2.getStructuringElement(cv2.MORPH_RECT, (1, params["v_close_height"]))
    v_lines = cv2.morphologyEx(v_lines, cv2.MORPH_CLOSE, v_close)
    v_lines = remove_small_components(v_lines, params["min_line_area"])
    cv2.imwrite(str(DEBUG_DIR / f"{name}_03_v_lines.png"), v_lines)
    print(f"  [3] V lines (kernels: {params['v_kernel_long']}/{params['v_kernel_short']})")

    # ---- Step 4: Combine with directional dilation ----
    h_dilate_k = cv2.getStructuringElement(cv2.MORPH_RECT, (1, params["h_dilate_v"]))
    v_dilate_k = cv2.getStructuringElement(cv2.MORPH_RECT, (params["v_dilate_h"], 1))
    grid_mask = cv2.bitwise_or(
        cv2.dilate(h_lines, h_dilate_k),
        cv2.dilate(v_lines, v_dilate_k),
    )
    cv2.imwrite(str(DEBUG_DIR / f"{name}_04_grid_mask.png"), grid_mask)
    print(f"  [4] Grid mask (directional dilation)")

    # ---- Step 5: Uniform dilation to close remaining gaps ----
    if params["grid_dilate_size"] > 0 and params["grid_dilate_iter"] > 0:
        dk = cv2.getStructuringElement(
            cv2.MORPH_RECT,
            (params["grid_dilate_size"], params["grid_dilate_size"]),
        )
        grid_thick = cv2.dilate(grid_mask, dk, iterations=params["grid_dilate_iter"])
    else:
        grid_thick = grid_mask.copy()
    cv2.imwrite(str(DEBUG_DIR / f"{name}_05_grid_thick.png"), grid_thick)
    print(f"  [5] Grid thick (uniform dilation {params['grid_dilate_size']}x{params['grid_dilate_iter']})")

    # ---- Step 6: Invert (cells = white) ----
    inverse = cv2.bitwise_not(grid_thick)
    margin = params["border_margin"]
    inverse[:margin, :] = 0
    inverse[-margin:, :] = 0
    inverse[:, :margin] = 0
    inverse[:, -margin:] = 0
    cv2.imwrite(str(DEBUG_DIR / f"{name}_06_inverse.png"), inverse)
    print(f"  [6] Inverse (border zeroed, margin={margin})")

    # ---- Step 7: Find contours ----
    contours, _ = cv2.findContours(inverse, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
    print(f"  [7] Found {len(contours)} raw contours")

    contour_img = img.copy()
    cv2.drawContours(contour_img, contours, -1, (0, 255, 0), 1)
    cv2.imwrite(str(DEBUG_DIR / f"{name}_07_contours.png"), contour_img)

    # ---- Step 8: Filter contours ----
    min_area = params["min_cell_area"]
    max_area = total_area * params["max_cell_area_ratio"]

    candidates = []
    for cnt in contours:
        area = cv2.contourArea(cnt)
        if area < min_area or area > max_area:
            continue

        x, y, bw, bh = cv2.boundingRect(cnt)

        if bw < params["min_cell_width"] or bh < params["min_cell_height"]:
            continue

        aspect = max(bw, bh) / max(min(bw, bh), 1)
        if aspect > params["max_aspect_ratio"]:
            continue

        if bw > w * 0.85 and bh > h * 0.85:
            continue

        rect_area = bw * bh
        rectangularity = area / max(rect_area, 1)
        if rectangularity < params["min_rectangularity"]:
            continue

        candidates.append((x, y, bw, bh, area, rectangularity))

    print(f"  [8a] {len(candidates)} candidates after shape filtering")

    # ---- IoU-based dedup ----
    iou_thresh = params["dedup_iou"]
    keep = [True] * len(candidates)
    for i in range(len(candidates)):
        if not keep[i]:
            continue
        x1, y1, w1, h1, a1, _ = candidates[i]
        for j in range(i + 1, len(candidates)):
            if not keep[j]:
                continue
            x2, y2, w2, h2, a2, _ = candidates[j]
            ox = max(0, min(x1 + w1, x2 + w2) - max(x1, x2))
            oy = max(0, min(y1 + h1, y2 + h2) - max(y1, y2))
            inter = ox * oy
            union = w1 * h1 + w2 * h2 - inter
            iou = inter / max(union, 1)
            if iou > iou_thresh:
                # Drop the smaller one
                if a1 < a2:
                    keep[i] = False
                    break
                else:
                    keep[j] = False

    cells = [c for c, k in zip(candidates, keep) if k]
    print(f"  [8b] {len(cells)} cells after IoU dedup (thresh={iou_thresh})")

    # ---- Draw results ----
    colors = [
        (0, 200, 0), (200, 0, 0), (0, 0, 200), (200, 200, 0),
        (0, 200, 200), (200, 0, 200), (100, 200, 0), (0, 100, 200),
    ]
    cell_img = img.copy()
    for i, (x, y, bw, bh, area, rect) in enumerate(cells):
        color = colors[i % len(colors)]
        cv2.rectangle(cell_img, (x, y), (x + bw, y + bh), color, 2)
    cv2.imwrite(str(DEBUG_DIR / f"{name}_08_cells.png"), cell_img)

    # Grid overlay on original
    grid_overlay = img.copy()
    grid_color = np.zeros_like(img)
    grid_color[:, :, 2] = grid_thick
    grid_overlay = cv2.addWeighted(grid_overlay, 0.7, grid_color, 0.5, 0)
    cv2.imwrite(str(DEBUG_DIR / f"{name}_09_grid_overlay.png"), grid_overlay)

    # ---- Stats ----
    if cells:
        areas = [c[4] for c in cells]
        widths = [c[2] for c in cells]
        heights = [c[3] for c in cells]
        print(f"\n  Results:")
        print(f"    Total cells: {len(cells)}")
        print(f"    Area range: {min(areas):.0f} - {max(areas):.0f}")
        print(f"    Width range: {min(widths)} - {max(widths)}")
        print(f"    Height range: {min(heights)} - {max(heights)}")
        small = sum(1 for a in areas if a < 2000)
        medium = sum(1 for a in areas if 2000 <= a < 20000)
        large = sum(1 for a in areas if a >= 20000)
        print(f"    Size breakdown: small(<2k)={small}  medium(2k-20k)={medium}  large(>20k)={large}")

        # Coverage: what fraction of image area is covered by cell bounding boxes
        cell_mask = np.zeros((h, w), dtype=np.uint8)
        for x, y, bw, bh, _, _ in cells:
            cell_mask[y:y+bh, x:x+bw] = 255
        coverage = np.sum(cell_mask > 0) / total_area
        print(f"    Page coverage: {coverage:.1%}")
    else:
        print("  No cells detected!")

    return cells


if __name__ == "__main__":
    images = [
        (INPUT_DIR / "boxes-1.png", "boxes1"),
        (INPUT_DIR / "boxes-2.png", "boxes2"),
    ]

    all_results = {}
    for img_path, name in images:
        if img_path.exists():
            all_results[name] = process_image(img_path, name, PARAMS)
        else:
            print(f"WARNING: {img_path} not found")

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