import cv2
import numpy as np
import os

pdf_path = r"C:\Users\Certes\Desktop\guitar_score\output\shintakarajima_perfect.pdf"
out_dir = r"C:\Users\Certes\Desktop\guitar_score\scripts\debug"

img_path = os.path.join(out_dir, "verify_chunk_0.jpg") # from previous sessions
if not os.path.exists(img_path):
    print("verify_chunk_0.jpg not found. extracting a page from PDF...")
    import fitz
    doc = fitz.open(pdf_path)
    page = doc.load_page(0)
    pix = page.get_pixmap(dpi=150)
    pix.save(os.path.join(out_dir, "pdf_test_page.png"))
    img_path = os.path.join(out_dir, "pdf_test_page.png")

img = cv2.imread(img_path)
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
_, bin_inv = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY_INV)

h, w = img.shape[:2]
staff_y_top = int(h * 0.3)
row_sums = np.sum(bin_inv[:, :min(w, 1000)], axis=1) / 255.0
staff_rows = np.where(row_sums > min(w, 1000) * 0.4)[0]
if len(staff_rows) >= 6: staff_y_top = staff_rows[0]

# locate measure bars
col_sums = np.sum(bin_inv[staff_y_top:staff_y_top+100, :], axis=0) / 255.0
bar_xs = np.where(col_sums > 30)[0]
bars = []
if len(bar_xs) > 0:
    curr = [bar_xs[0]]
    for x in bar_xs[1:]:
        if x - curr[-1] < 10: curr.append(x)
        else:
            bars.append(int(np.mean(curr)))
            curr = [x]
    bars.append(int(np.mean(curr)))

# Crop first 5 measure numbers
for i, x in enumerate(bars[:5]):
    box_y1 = max(0, staff_y_top - 40)
    box_y2 = staff_y_top
    box_x1 = x
    box_x2 = min(w, x + 60)
    crop = img[box_y1:box_y2, box_x1:box_x2]
    out_file = os.path.join(out_dir, f"measure_num_{i}.png")
    cv2.imwrite(out_file, crop)
    print(f"Saved measure number crop to {out_file}")
