import cv2
import pickle
import numpy as np
import easyocr
import time
import re

reader = easyocr.Reader(['en'], gpu=False)

with open('unique_pages.pkl', 'rb') as f:
    unique_pages = pickle.load(f)

print(f"Loaded {len(unique_pages)} chunks. Running OCR on jump-cut boundaries...")

def extract_measure_number(page_bgr):
    # Same logic as before to find the first measure box
    cw = min(page_bgr.shape[1], 1000)
    page_gray = cv2.cvtColor(page_bgr[:, :cw], cv2.COLOR_BGR2GRAY)
    _, bin_inv = cv2.threshold(page_gray, 200, 255, cv2.THRESH_BINARY_INV)
    
    row_sums = np.sum(bin_inv, axis=1) / 255.0
    staff_rows = np.where(row_sums > cw * 0.4)[0]
    
    if len(staff_rows) >= 6:
        staff_y_top, staff_y_bottom = staff_rows[0], staff_rows[-1]
        for r in staff_rows:
            if r - staff_y_top > 100: break
            staff_y_bottom = r
    else:
        return -1
        
    expected_h = max(10, staff_y_bottom - staff_y_top + 1)
    staff_region = bin_inv[staff_y_top:staff_y_bottom+1, :]
    col_sums = np.sum(staff_region, axis=0) / 255.0
    bar_xs = np.where(col_sums >= expected_h * 0.8)[0]
    
    if len(bar_xs) == 0: return -1
    x_bar = bar_xs[0]
    
    box_y1 = max(0, staff_y_top - 25)
    box_y2 = staff_y_top
    box_x1 = x_bar
    box_x2 = min(page_gray.shape[1], x_bar + 35)
    
    num_box = page_gray[box_y1:box_y2, box_x1:box_x2]
    
    # Preprocess for OCR
    _, num_inv = cv2.threshold(num_box, 200, 255, cv2.THRESH_BINARY_INV)
    
    # Must pass white background with black text to EasyOCR! (Since it reads printed text)
    num_for_ocr = cv2.bitwise_not(num_inv) 
    
    upscaled = cv2.resize(num_for_ocr, None, fx=4, fy=4, interpolation=cv2.INTER_CUBIC)
    padded = cv2.copyMakeBorder(upscaled, 20, 20, 20, 20, cv2.BORDER_CONSTANT, value=[255, 255, 255])
    
    results = reader.readtext(padded, allowlist="0123456789")
    if not results: return -1
    
    text = results[0][1]
    
    digits = re.findall(r'\d+', text)
    if digits:
        return int(digits[0])
    return -1

results = []
for i, page in enumerate(unique_pages):
    t0 = time.time()
    num = extract_measure_number(page)
    tf = time.time()
    print(f"Page {i:02d}: {num} (took {tf-t0:.2f}s)")
    results.append(num)

print(f"Sequential Detections: {results}")
