import cv2
import numpy as np
import json
import sys

DPI = 150
PT_PER_PX = 72.0 / DPI
PAGE_H_PT = 842.0  # A4 height in points

def detect(page_png, page_num):
    img = cv2.imread(page_png)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # binarize (boxes are thin black lines on white)
    _, th = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY_INV)

    # isolate horizontal and vertical lines
    horiz = th.copy()
    vert = th.copy()
    h_size = max(1, img.shape[1] // 60)
    v_size = max(1, img.shape[0] // 60)
    horiz_struct = cv2.getStructuringElement(cv2.MORPH_RECT, (h_size, 1))
    vert_struct = cv2.getStructuringElement(cv2.MORPH_RECT, (1, v_size))
    horiz = cv2.erode(horiz, horiz_struct); horiz = cv2.dilate(horiz, horiz_struct)
    vert = cv2.erode(vert, vert_struct); vert = cv2.dilate(vert, vert_struct)
    grid = cv2.bitwise_and(horiz, vert) | cv2.bitwise_or(horiz, vert)
    grid_lines = cv2.add(horiz, vert)

    contours, _ = cv2.findContours(grid_lines, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
    boxes = []
    for c in contours:
        x, y, w, h = cv2.boundingRect(c)
        area = w * h
        # small square-ish cells: comb boxes are roughly 20-40px wide/tall at 150dpi
        if 15 <= w <= 60 and 15 <= h <= 60 and 0.5 <= w / h <= 2.0:
            boxes.append((x, y, w, h))

    # de-duplicate near-identical boxes
    boxes = sorted(set(boxes))
    dedup = []
    for b in boxes:
        if not any(abs(b[0]-d[0]) < 4 and abs(b[1]-d[1]) < 4 for d in dedup):
            dedup.append(b)
    boxes = dedup

    # group into rows by y (cells in the same comb row share y +/- tolerance)
    boxes.sort(key=lambda b: (b[1], b[0]))
    rows = []
    for b in boxes:
        placed = False
        for row in rows:
            if abs(row[0][1] - b[1]) < 8:
                row.append(b)
                placed = True
                break
        if not placed:
            rows.append([b])
    for row in rows:
        row.sort(key=lambda b: b[0])

    rows.sort(key=lambda row: row[0][1])

    out_rows = []
    for row in rows:
        cells = []
        for (x, y, w, h) in row:
            # PDF point coords: x same, y flipped (PDF origin bottom-left, but FPDI/FPDF use top-left too, so just scale)
            px = round(x * PT_PER_PX, 1)
            py = round(y * PT_PER_PX, 1)
            pw = round(w * PT_PER_PX, 1)
            ph = round(h * PT_PER_PX, 1)
            cells.append({"x": px, "y": py, "w": pw, "h": ph})
        out_rows.append({
            "y_px": row[0][1],
            "cell_count": len(row),
            "cells": cells
        })

    return out_rows

if __name__ == "__main__":
    for pnum in (1, 2, 3):
        rows = detect(f"pdf/page-{pnum}.png", pnum)
        with open(f"boxes_page{pnum}.json", "w") as f:
            json.dump(rows, f, indent=2)
        print(f"page {pnum}: {len(rows)} comb-rows detected, cell counts:",
              [r["cell_count"] for r in rows])