import cv2
import numpy as np
import sys
from pathlib import Path

def test_lines(img_path, out_path):
    img = cv2.imread(img_path)
    if img is None:
        print("Could not read image")
        return
        
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    
    # 1. Use standard binary thresholding instead of adaptive to reduce noise from text
    # Newspapers usually have very black lines on white background
    _, thresh = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY_INV)
    
    # We want lines that are at least 300px long (a typical column width)
    line_length = 300
    
    # 2. Detect horizontal lines (MUST be thin: e.g., height <= 5px)
    horizontal_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (line_length, 1))
    horizontal_lines_raw = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, horizontal_kernel, iterations=2)
    
    # 3. Detect vertical lines (MUST be thin: e.g., width <= 5px)
    vertical_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, line_length))
    vertical_lines_raw = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, vertical_kernel, iterations=2)
    
    # 4. Filter contours to ensure they are actually THIN lines, not thick blocks/images
    h_contours_raw, _ = cv2.findContours(horizontal_lines_raw, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    v_contours_raw, _ = cv2.findContours(vertical_lines_raw, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    h_contours = []
    for cnt in h_contours_raw:
        x, y, w, h = cv2.boundingRect(cnt)
        if h <= 15:  # Line must not be thicker than 15px
            h_contours.append(cnt)
            
    v_contours = []
    for cnt in v_contours_raw:
        x, y, w, h = cv2.boundingRect(cnt)
        if w <= 15:  # Line must not be thicker than 15px
            v_contours.append(cnt)
            
    print(f"Horizontal separators found: {len(h_contours)} (filtered from {len(h_contours_raw)})")
    print(f"Vertical separators found: {len(v_contours)} (filtered from {len(v_contours_raw)})")
    
    # 5. Create clean masks from only the filtered thin lines
    clean_h_mask = np.zeros_like(thresh)
    cv2.drawContours(clean_h_mask, h_contours, -1, 255, thickness=cv2.FILLED)
    
    clean_v_mask = np.zeros_like(thresh)
    cv2.drawContours(clean_v_mask, v_contours, -1, 255, thickness=cv2.FILLED)
    
    grid = cv2.add(clean_h_mask, clean_v_mask)
    rect_contours, _ = cv2.findContours(grid, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)
    
    box_count = 0
    debug_img = img.copy()
    
    for cnt in rect_contours:
        epsilon = 0.02 * cv2.arcLength(cnt, True)
        approx = cv2.approxPolyDP(cnt, epsilon, True)
        
        if len(approx) == 4:
            x, y, w, h = cv2.boundingRect(approx)
            if w > 300 and h > 200:
                box_count += 1
                cv2.rectangle(debug_img, (x, y), (x+w, y+h), (255, 0, 255), 8)
                
    print(f"Rectangular boxes found: {box_count}")
    
    cv2.drawContours(debug_img, h_contours, -1, (0, 0, 255), 3) # Red for horizontal
    cv2.drawContours(debug_img, v_contours, -1, (255, 0, 0), 3) # Blue for vertical
    
    cv2.imwrite(out_path, debug_img)
    print(f"Saved debug image to {out_path}")

if __name__ == "__main__":
    test_lines(sys.argv[1], sys.argv[2])