#!/usr/bin/env python3
"""
Hybrid form cell detection (v3).

Combines:
- v2-style grid-mask construction + gap/junction repair + contour candidate generation
- v1-style explicit rectangle validation via snapped/merged horizontal+vertical line ranges

Key improvements over v1/v2:
- Junction repair: dilate only at near-intersections (not entire lines)
- Three-pass vertical isolation (long/short/tiny) for short dividers
- Contour candidates avoid O(n^4) exhaustive search
- 4-side validation prevents merged-region false positives
- Auto-scaled kernels based on image width

Based on GPT 5.2 review of v1 and v2.
Debug outputs written to temp/boxes/debug3.
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Sequence, Tuple, Optional

import cv2
import numpy as np


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

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


@dataclass(frozen=True)
class Params:
    # Resize normalization
    target_width: int = 2200

    # Adaptive threshold
    adaptive_block: int = 15
    adaptive_c: int = 10

    # Noise cleanup for line fragments
    min_line_area: int = 120

    # Directional dilation to help H/V meet at intersections
    # (increased from 4 — need enough overlap to create junctions)
    h_dilate_v: int = 8
    v_dilate_h: int = 8

    # Junction repair (patch "almost-touching" intersections)
    junction_dilate: int = 6

    # Final grid sealing
    grid_dilate_size: int = 3
    grid_dilate_iter: int = 1

    # Contour candidate filtering
    min_cell_area: int = 350
    max_cell_area_ratio: float = 0.40
    min_cell_w: int = 12
    min_cell_h: int = 10
    max_aspect: float = 30.0
    min_rectangularity: float = 0.35

    # Border margin for inverse mask
    border_margin: int = 6

    # Validation tolerances (relaxed — v_lines are often fragmented)
    snap_tolerance: int = 12
    edge_tolerance_boost: int = 4

    # Max cell height (prevents full-page spans)
    max_cell_h: int = 500

    # Dedup
    dedup_iou: float = 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) -> tuple:
    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: float, lo: int, hi: int) -> int:
    return int(max(lo, min(hi, round(v))))


def remove_small_components(mask: np.ndarray, min_area: int) -> np.ndarray:
    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 snap_coordinates(values: Sequence[int], tolerance: int) -> Dict[int, int]:
    if not values:
        return {}
    sorted_vals = sorted(set(int(v) for v in values))
    clusters: List[List[int]] = []
    current = [sorted_vals[0]]
    for v in sorted_vals[1:]:
        if v - current[-1] <= tolerance:
            current.append(v)
        else:
            clusters.append(current)
            current = [v]
    clusters.append(current)

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


def merge_ranges(ranges: Sequence[Tuple[int, int]], gap_tolerance: int = 5) -> List[Tuple[int, int]]:
    if not ranges:
        return []
    sr = sorted((min(a, b), max(a, b)) for a, b in ranges)
    merged: List[List[int]] = [list(sr[0])]
    for start, end in sr[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: Sequence[Tuple[int, int]], start: int, end: int, tol: int) -> bool:
    for a, b in ranges:
        if a <= start + tol and b >= end - tol:
            return True
    return False


def range_touches(ranges: Sequence[Tuple[int, int]], point: int, tol: int) -> bool:
    for a, b in ranges:
        if a - tol <= point <= b + tol:
            return True
    return False


def build_line_structures(
    h_mask: np.ndarray,
    v_mask: np.ndarray,
    snap_tol: int,
    min_h_len: int,
    min_v_len: int,
    max_thickness: int = 24,
) -> Tuple[Dict[int, List[Tuple[int, int]]], Dict[int, List[Tuple[int, int]]], List[int], List[int]]:
    """Extract segments from H/V masks, snap, merge into line lookup structures."""

    def extract_segments(binary, orientation):
        contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        segs = []
        for cnt in contours:
            x, y, w, h = cv2.boundingRect(cnt)
            if orientation == "h":
                if w >= min_h_len and h <= max_thickness:
                    cy = y + h // 2
                    segs.append((x, cy, x + w, cy))
            else:
                if h >= min_v_len and w <= max_thickness:
                    cx = x + w // 2
                    segs.append((cx, y, cx, y + h))
        return segs

    h_segments = extract_segments(h_mask, "h")
    v_segments = extract_segments(v_mask, "v")

    all_ys = ([s[1] for s in h_segments] +
              [s[1] for s in v_segments] + [s[3] for s in v_segments])
    all_xs = ([s[0] for s in v_segments] +
              [s[0] for s in h_segments] + [s[2] for s in h_segments])

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

    h_lines: Dict[int, List[Tuple[int, int]]] = {}
    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)
        h_lines.setdefault(sy, []).append((sx1, sx2))

    v_lines: Dict[int, List[Tuple[int, int]]] = {}
    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)
        v_lines.setdefault(sx, []).append((sy1, sy2))

    for y in list(h_lines.keys()):
        h_lines[y] = merge_ranges(h_lines[y], gap_tolerance=snap_tol)
    for x in list(v_lines.keys()):
        v_lines[x] = merge_ranges(v_lines[x], gap_tolerance=snap_tol)

    return h_lines, v_lines, sorted(h_lines.keys()), sorted(v_lines.keys())


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 = max(0, ax2 - ax1) * max(0, ay2 - ay1)
    area_b = max(0, bx2 - bx1) * max(0, 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 contour_candidates(inverse, img_area, p):
    contours, _ = cv2.findContours(inverse, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)
    out = []
    max_area = img_area * p.max_cell_area_ratio
    for cnt in contours:
        area = cv2.contourArea(cnt)
        if area < p.min_cell_area or area > max_area:
            continue
        x, y, w, h = cv2.boundingRect(cnt)
        if w < p.min_cell_w or h < p.min_cell_h:
            continue
        aspect = max(w, h) / max(min(w, h), 1)
        if aspect > p.max_aspect:
            continue
        rectangularity = float(area) / float(max(w * h, 1))
        if rectangularity < p.min_rectangularity:
            continue
        out.append((x, y, x + w, y + h))
    return dedup_boxes(out, p.dedup_iou)


def validate_candidate(box, h_lines, v_lines, unique_ys, unique_xs,
                        edge_tol, snap_tol, max_cell_h):
    """Snap candidate to nearest grid lines and check all 4 sides exist."""
    x1, y1, x2, y2 = box
    if x2 <= x1 or y2 <= y1 or (y2 - y1) > max_cell_h:
        return None

    def nearest(vals, v):
        if not vals:
            return None
        idx = int(np.searchsorted(vals, v))
        candidates = []
        if 0 <= idx < len(vals):
            candidates.append(vals[idx])
        if 0 <= idx - 1 < len(vals):
            candidates.append(vals[idx - 1])
        if not candidates:
            return None
        best = min(candidates, key=lambda t: abs(t - v))
        return best if abs(best - v) <= snap_tol else None

    sx1 = nearest(unique_xs, x1)
    sx2 = nearest(unique_xs, x2)
    sy1 = nearest(unique_ys, y1)
    sy2 = nearest(unique_ys, y2)
    if None in (sx1, sx2, sy1, sy2) or sx2 <= sx1 or sy2 <= sy1:
        return None

    # 4-side validation: H sides must be covered, V sides just need presence
    # (V lines are often fragmented, so relax to "any segment exists near this x")
    if not range_covered(h_lines.get(sy1, []), sx1, sx2, edge_tol):
        return None
    if not range_covered(h_lines.get(sy2, []), sx1, sx2, edge_tol):
        return None

    # For verticals: check nearby x values too (within snap_tol)
    def has_vertical_near(x, y_top, y_bot, tol):
        for vx in range(x - tol, x + tol + 1):
            ranges = v_lines.get(vx, [])
            if ranges and (range_touches(ranges, y_top, edge_tol) or
                           range_touches(ranges, y_bot, edge_tol)):
                return True
        return False

    if not has_vertical_near(sx1, sy1, sy2, snap_tol // 2):
        return None
    if not has_vertical_near(sx2, sy1, sy2, snap_tol // 2):
        return None

    return (sx1, sy1, sx2, sy2)


def process_image(name, image_path, p):
    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 for stable parameters
    img_norm, scale = resize_to_target(img, p.target_width)
    h, w = img_norm.shape[:2]
    img_area = h * w
    print(f"  Original: {img.shape[1]}x{img.shape[0]} -> Normalized: {w}x{h} (scale={scale:.3f})")

    # Auto-scale kernels
    k = w / 2200.0
    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(15 * k, 10, 30)
    v_tiny = clamp_int(10 * k, 7, 16)
    h_close_w = clamp_int(15 * k, 10, 30)
    v_close_h = clamp_int(10 * k, 6, 18)
    print(f"  Kernels: H={h_long}/{h_short}, V={v_long}/{v_short}/{v_tiny}")

    # Step 1: Grayscale + binary
    gray = cv2.cvtColor(img_norm, cv2.COLOR_BGR2GRAY)
    binary = cv2.adaptiveThreshold(
        gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY_INV, p.adaptive_block, p.adaptive_c
    )
    save_debug(name, "01", "binary", binary)

    # Step 2: Horizontal lines (two-pass)
    hk_long = cv2.getStructuringElement(cv2.MORPH_RECT, (h_long, 1))
    hk_short = cv2.getStructuringElement(cv2.MORPH_RECT, (h_short, 1))
    h_lines_img = 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_img = cv2.morphologyEx(h_lines_img, cv2.MORPH_CLOSE, h_close)
    h_lines_img = remove_small_components(h_lines_img, p.min_line_area)
    save_debug(name, "02", "h_lines", h_lines_img)

    # Step 3: Vertical lines (THREE-pass: long/short/tiny)
    vk_long = cv2.getStructuringElement(cv2.MORPH_RECT, (1, v_long))
    vk_short = cv2.getStructuringElement(cv2.MORPH_RECT, (1, v_short))
    vk_tiny = cv2.getStructuringElement(cv2.MORPH_RECT, (1, v_tiny))
    v_lines_img = cv2.bitwise_or(
        cv2.bitwise_or(
            cv2.morphologyEx(binary, cv2.MORPH_OPEN, vk_long),
            cv2.morphologyEx(binary, cv2.MORPH_OPEN, vk_short),
        ),
        cv2.morphologyEx(binary, cv2.MORPH_OPEN, vk_tiny),
    )
    v_close = cv2.getStructuringElement(cv2.MORPH_RECT, (1, v_close_h))
    v_lines_img = cv2.morphologyEx(v_lines_img, cv2.MORPH_CLOSE, v_close)
    v_lines_img = remove_small_components(v_lines_img, p.min_line_area)
    save_debug(name, "03", "v_lines", v_lines_img)

    # Step 4: Directional dilation (make H thicker vertically, V thicker horizontally)
    h_dk = cv2.getStructuringElement(cv2.MORPH_RECT, (1, p.h_dilate_v))
    v_dk = cv2.getStructuringElement(cv2.MORPH_RECT, (p.v_dilate_h, 1))
    h_fat = cv2.dilate(h_lines_img, h_dk)
    v_fat = cv2.dilate(v_lines_img, v_dk)

    # Step 5: Junction repair — dilate only where H and V nearly intersect
    junc = cv2.bitwise_and(h_fat, v_fat)
    if p.junction_dilate > 0:
        jk = cv2.getStructuringElement(cv2.MORPH_RECT, (p.junction_dilate, p.junction_dilate))
        junc = cv2.dilate(junc, jk)
    save_debug(name, "04", "junctions", junc)

    grid = cv2.bitwise_or(cv2.bitwise_or(h_fat, v_fat), junc)

    # Step 6: Final uniform dilation
    if p.grid_dilate_size > 0 and p.grid_dilate_iter > 0:
        dk = cv2.getStructuringElement(cv2.MORPH_RECT, (p.grid_dilate_size, p.grid_dilate_size))
        grid = cv2.dilate(grid, dk, iterations=p.grid_dilate_iter)
    save_debug(name, "05", "grid", grid)

    # Step 7: Invert to get cell regions
    inverse = cv2.bitwise_not(grid)
    m = p.border_margin
    inverse[:m, :] = 0
    inverse[-m:, :] = 0
    inverse[:, :m] = 0
    inverse[:, -m:] = 0
    save_debug(name, "06", "inverse", inverse)

    # Build line structures for validation (from clean H/V masks, not fat grid)
    min_h_len = clamp_int(25 * k, 18, 80)
    min_v_len = clamp_int(12 * k, 8, 40)
    h_struct, v_struct, unique_ys, unique_xs = build_line_structures(
        h_lines_img, v_lines_img,
        snap_tol=p.snap_tolerance, min_h_len=min_h_len, min_v_len=min_v_len,
    )
    print(f"  Grid: {len(unique_ys)} H-lines, {len(unique_xs)} V-lines")

    # Generate candidates from contours
    candidates = contour_candidates(inverse, img_area, p)
    print(f"  Contour candidates: {len(candidates)}")

    # Validate each candidate (snap + 4-side check)
    edge_tol = p.snap_tolerance + p.edge_tolerance_boost
    max_cell_h = clamp_int(p.max_cell_h * k, 200, 1200)
    valid = []
    for b in candidates:
        vb = validate_candidate(b, h_struct, v_struct, unique_ys, unique_xs,
                                edge_tol, p.snap_tolerance, max_cell_h)
        if vb is not None:
            valid.append(vb)

    valid = dedup_boxes(valid, p.dedup_iou)
    valid.sort(key=lambda c: (c[1], c[0]))
    print(f"  Validated cells: {len(valid)} (rejected {len(candidates) - len(valid)})")

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

    # Draw result
    cell_img = draw_cells(img, valid)
    save_debug(name, "07", "cells", cell_img)

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

    return valid


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


if __name__ == "__main__":
    p = Params()
    print("Hybrid Form Cell Detection (v3) - Junction Repair + Contour Candidates + 4-Side Validation")
    print(f"Debug dir: {DEBUG_DIR}\n")
    results = {}
    for name, path in IMAGES:
        cells = process_image(name, path, p)
        results[name] = cells
    print(f"\n{'='*60}")
    print("SUMMARY")
    for name, cells in results.items():
        print(f"  {name}: {len(cells)} cells")
