working_yolo_pipeline.py

import json
import argparse
import os
import re
import torch
import torch.nn as nn
from TorchCRF import CRF
from transformers import LayoutLMv3TokenizerFast, LayoutLMv3Model, LayoutLMv3Config
from typing import List, Dict, Any, Optional, Union, Tuple
import fitz  # PyMuPDF
import numpy as np
import cv2
from ultralytics import YOLO
import glob
import pytesseract
from PIL import Image
from scipy.signal import find_peaks
from scipy.ndimage import gaussian_filter1d
import sys
import io
import base64
import tempfile
import time
import shutil
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.metrics.pairwise import cosine_similarity

# ============================================================================
# --- CONFIGURATION AND CONSTANTS ---
# ============================================================================

# NOTE: Update these paths to match your environment before running!
WEIGHTS_PATH = 'YOLO_MATH/yolo_split_data/runs/detect/math_figure_detector_v3/weights/best.pt'
DEFAULT_LAYOUTLMV3_MODEL_PATH = "97.pth"

# DIRECTORY CONFIGURATION
OCR_JSON_OUTPUT_DIR = './ocr_json_output_final'
FIGURE_EXTRACTION_DIR = './figure_extraction'
TEMP_IMAGE_DIR = './temp_pdf_images'

# Detection parameters
CONF_THRESHOLD = 0.2
TARGET_CLASSES = ['figure', 'equation']
IOU_MERGE_THRESHOLD = 0.4
IOA_SUPPRESSION_THRESHOLD = 0.7
LINE_TOLERANCE = 15


#Similarity
SIMILARITY_THRESHOLD = 0.10
RESOLUTION_MARGIN = 0.05

# Global counters for sequential numbering across the entire PDF
GLOBAL_FIGURE_COUNT = 0
GLOBAL_EQUATION_COUNT = 0

# LayoutLMv3 Labels
ID_TO_LABEL = {
    0: "O",
    1: "B-QUESTION", 2: "I-QUESTION",
    3: "B-OPTION", 4: "I-OPTION",
    5: "B-ANSWER", 6: "I-ANSWER",
    7: "B-SECTION_HEADING", 8: "I-SECTION_HEADING",
    9: "B-PASSAGE", 10: "I-PASSAGE"
}
NUM_LABELS = len(ID_TO_LABEL)


# ============================================================================
# --- PERFORMANCE OPTIMIZATION: OCR CACHE ---
# ============================================================================

class OCRCache:
    """Caches OCR results per page to avoid redundant Tesseract runs."""

    def __init__(self):
        self.cache = {}

    def get_key(self, pdf_path: str, page_num: int) -> str:
        return f"{pdf_path}:{page_num}"

    def has_ocr(self, pdf_path: str, page_num: int) -> bool:
        return self.get_key(pdf_path, page_num) in self.cache

    def get_ocr(self, pdf_path: str, page_num: int) -> Optional[list]:
        return self.cache.get(self.get_key(pdf_path, page_num))

    def set_ocr(self, pdf_path: str, page_num: int, ocr_data: list):
        self.cache[self.get_key(pdf_path, page_num)] = ocr_data

    def clear(self):
        self.cache.clear()


# Global OCR cache instance
_ocr_cache = OCRCache()


# ============================================================================
# --- PHASE 1: YOLO/OCR PREPROCESSING FUNCTIONS ---
# ============================================================================

def calculate_iou(box1, box2):
    x1_a, y1_a, x2_a, y2_a = box1
    x1_b, y1_b, x2_b, y2_b = box2
    x_left = max(x1_a, x1_b)
    y_top = max(y1_a, y1_b)
    x_right = min(x2_a, x2_b)
    y_bottom = min(y2_a, y2_b)
    intersection_area = max(0, x_right - x_left) * max(0, y_bottom - y_top)
    box_a_area = (x2_a - x1_a) * (y2_a - y1_a)
    box_b_area = (x2_b - x1_b) * (y2_b - y1_b)
    union_area = float(box_a_area + box_b_area - intersection_area)
    return intersection_area / union_area if union_area > 0 else 0


def calculate_ioa(box1, box2):
    x1_a, y1_a, x2_a, y2_a = box1
    x1_b, y1_b, x2_b, y2_b = box2
    x_left = max(x1_a, x1_b)
    y_top = max(y1_a, y1_b)
    x_right = min(x2_a, x2_b)
    y_bottom = min(y2_a, y2_b)
    intersection_area = max(0, x_right - x_left) * max(0, y_bottom - y_top)
    box_a_area = (x2_a - x1_a) * (y2_a - y1_a)
    return intersection_area / box_a_area if box_a_area > 0 else 0


def filter_nested_boxes(detections, ioa_threshold=0.80):
    """
    Removes boxes that are inside larger boxes (Containment Check).
    Prioritizes keeping the LARGEST box (the 'parent' container).
    """
    if not detections: 
        return []

    # 1. Calculate Area for all detections
    for d in detections:
        x1, y1, x2, y2 = d['coords']
        d['area'] = (x2 - x1) * (y2 - y1)

    # 2. Sort by Area Descending (Largest to Smallest)
    # This ensures we process the 'container' first
    detections.sort(key=lambda x: x['area'], reverse=True)

    keep_indices = []
    is_suppressed = [False] * len(detections)

    for i in range(len(detections)):
        if is_suppressed[i]: continue

        keep_indices.append(i)
        box_a = detections[i]['coords']
        
        # Compare with all smaller boxes
        for j in range(i + 1, len(detections)):
            if is_suppressed[j]: continue

            box_b = detections[j]['coords']
            
            # Calculate Intersection
            x_left = max(box_a[0], box_b[0])
            y_top = max(box_a[1], box_b[1])
            x_right = min(box_a[2], box_b[2])
            y_bottom = min(box_a[3], box_b[3])

            if x_right < x_left or y_bottom < y_top:
                intersection = 0
            else:
                intersection = (x_right - x_left) * (y_bottom - y_top)

            # Calculate IoA (Intersection over Area of the SMALLER box)
            # Since we sorted by area, 'box_b' (detections[j]) is the smaller one.
            area_b = detections[j]['area']
            
            if area_b > 0:
                ioa_small = intersection / area_b
                
                # If the small box is > 90% inside the big box, suppress the small one.
                if ioa_small > ioa_threshold:
                    is_suppressed[j] = True
                    #print(f"    [Suppress] Removed nested object inside larger '{detections[i]['class']}'")

    return [detections[i] for i in keep_indices]    


def merge_overlapping_boxes(detections, iou_threshold):
    if not detections: return []
    detections.sort(key=lambda d: d['conf'], reverse=True)
    merged_detections = []
    is_merged = [False] * len(detections)
    for i in range(len(detections)):
        if is_merged[i]: continue
        current_box = detections[i]['coords']
        current_class = detections[i]['class']
        merged_x1, merged_y1, merged_x2, merged_y2 = current_box
        for j in range(i + 1, len(detections)):
            if is_merged[j] or detections[j]['class'] != current_class: continue
            other_box = detections[j]['coords']
            iou = calculate_iou(current_box, other_box)
            if iou > iou_threshold:
                merged_x1 = min(merged_x1, other_box[0])
                merged_y1 = min(merged_y1, other_box[1])
                merged_x2 = max(merged_x2, other_box[2])
                merged_y2 = max(merged_y2, other_box[3])
                is_merged[j] = True
        merged_detections.append({
            'coords': (merged_x1, merged_y1, merged_x2, merged_y2),
            'y1': merged_y1, 'class': current_class, 'conf': detections[i]['conf']
        })
    return merged_detections


def merge_yolo_into_word_data(raw_word_data: list, yolo_detections: list, scale_factor: float) -> list:
    """
    Filters out raw words that are inside YOLO boxes and replaces them with 
    a single solid 'placeholder' block for the column detector.
    """
    if not yolo_detections:
        return raw_word_data

    # 1. Convert YOLO boxes (Pixels) to PDF Coordinates (Points)
    pdf_space_boxes = []
    for det in yolo_detections:
        x1, y1, x2, y2 = det['coords']
        pdf_box = (
            x1 / scale_factor, 
            y1 / scale_factor, 
            x2 / scale_factor, 
            y2 / scale_factor
        )
        pdf_space_boxes.append(pdf_box)

    # 2. Filter out raw words that are inside YOLO boxes
    cleaned_word_data = []
    for word_tuple in raw_word_data:
        wx1, wy1, wx2, wy2 = word_tuple[1], word_tuple[2], word_tuple[3], word_tuple[4]
        w_center_x = (wx1 + wx2) / 2
        w_center_y = (wy1 + wy2) / 2
        
        is_inside_yolo = False
        for px1, py1, px2, py2 in pdf_space_boxes:
            if px1 <= w_center_x <= px2 and py1 <= w_center_y <= py2:
                is_inside_yolo = True
                break
        
        if not is_inside_yolo:
            cleaned_word_data.append(word_tuple)

    # 3. Add the YOLO boxes themselves as "Solid Words"
    for i, (px1, py1, px2, py2) in enumerate(pdf_space_boxes):
        dummy_entry = (f"BLOCK_{i}", px1, py1, px2, py2)
        cleaned_word_data.append(dummy_entry)

    return cleaned_word_data



# ============================================================================
# --- MISSING HELPER FUNCTION ---
# ============================================================================

def preprocess_image_for_ocr(img_np):
    """
    Converts image to grayscale and applies Otsu's Binarization 
    to separate text from background clearly.
    """
    # 1. Convert to Grayscale if needed
    if len(img_np.shape) == 3:
        gray = cv2.cvtColor(img_np, cv2.COLOR_BGR2GRAY)
    else:
        gray = img_np
        
    # 2. Apply Otsu's Thresholding (Automatic binary threshold)
    # This makes text solid black and background solid white
    _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    
    return thresh


def calculate_vertical_gap_coverage(word_data: list, sep_x: int, page_height: float, gutter_width: int = 10) -> float:
    """
    Calculates what percentage of the page's vertical text span is 'cleanly split' by the separator.
    A valid column split should split > 65% of the page verticality.
    """
    if not word_data:
        return 0.0

    # Determine the vertical span of the actual text content
    y_coords = [w[2] for w in word_data] + [w[4] for w in word_data] # y1 and y2
    min_y, max_y = min(y_coords), max(y_coords)
    total_text_height = max_y - min_y
    
    if total_text_height <= 0:
        return 0.0

    # Create a boolean array representing the Y-axis (1 pixel per unit)
    gap_open_mask = np.ones(int(total_text_height) + 1, dtype=bool)
    
    zone_left = sep_x - (gutter_width / 2)
    zone_right = sep_x + (gutter_width / 2)
    offset_y = int(min_y)
    
    for _, x1, y1, x2, y2 in word_data:
        # Check if this word horizontally interferes with the separator
        if x2 > zone_left and x1 < zone_right:
            y_start_idx = max(0, int(y1) - offset_y)
            y_end_idx = min(len(gap_open_mask), int(y2) - offset_y)
            if y_end_idx > y_start_idx:
                gap_open_mask[y_start_idx:y_end_idx] = False

    open_pixels = np.sum(gap_open_mask)
    coverage_ratio = open_pixels / len(gap_open_mask)
    
    return coverage_ratio




def calculate_x_gutters(word_data: list, params: Dict, page_height: float) -> List[int]:
    """
    Calculates X-axis histogram and validates using BRIDGING DENSITY and Vertical Coverage.
    """
    if not word_data: return []

    x_points = []
    # Use only word_data elements 1 (x1) and 3 (x2)
    for item in word_data:
        x_points.extend([item[1], item[3]])

    if not x_points: return []
    max_x = max(x_points)

    # 1. Determine total text height for ratio calculation
    y_coords = [item[2] for item in word_data] + [item[4] for item in word_data]
    min_y, max_y = min(y_coords), max(y_coords)
    total_text_height = max_y - min_y
    if total_text_height <= 0: return []

    # Histogram Setup
    bin_size = params.get('cluster_bin_size', 5)
    smoothing = params.get('cluster_smoothing', 1)
    min_width = params.get('cluster_min_width', 20)
    threshold_percentile = params.get('cluster_threshold_percentile', 85)

    num_bins = int(np.ceil(max_x / bin_size))
    hist, bin_edges = np.histogram(x_points, bins=num_bins, range=(0, max_x))
    smoothed_hist = gaussian_filter1d(hist.astype(float), sigma=smoothing)
    inverted_signal = np.max(smoothed_hist) - smoothed_hist

    peaks, properties = find_peaks(
        inverted_signal,
        height=np.max(inverted_signal) - np.percentile(smoothed_hist, threshold_percentile),
        distance=min_width / bin_size
    )

    if not peaks.size: return []
    separator_x_coords = [int(bin_edges[p]) for p in peaks]
    final_separators = []
    
    for x_coord in separator_x_coords:
        # --- CHECK 1: BRIDGING DENSITY (The "Cut Through" Check) ---
        # Calculate the total vertical height of words that physically cross this line.
        bridging_height = 0
        bridging_count = 0
        
        for item in word_data:
            wx1, wy1, wx2, wy2 = item[1], item[2], item[3], item[4]
            
            # Check if this word physically sits on top of the separator line
            if wx1 < x_coord and wx2 > x_coord:
                word_h = wy2 - wy1
                bridging_height += word_h
                bridging_count += 1
        
        # Calculate Ratio: How much of the page's text height is blocked by these crossing words?
        bridging_ratio = bridging_height / total_text_height
        
        # THRESHOLD: If bridging blocks > 8% of page height, REJECT.
        # This allows for page numbers or headers (usually < 5%) to cross, but NOT paragraphs.
        if bridging_ratio > 0.08: 
            print(f"      ❌ Separator X={x_coord} REJECTED: Bridging Ratio {bridging_ratio:.1%} (>15%) cuts through text.")
            continue

        # --- CHECK 2: VERTICAL GAP COVERAGE (The "Clean Split" Check) ---
        # The gap must exist cleanly for > 65% of the text height.
        coverage = calculate_vertical_gap_coverage(word_data, x_coord, page_height, gutter_width=min_width)
        
        if coverage >= 0.80:
            final_separators.append(x_coord)
            print(f"      -> Separator X={x_coord} ACCEPTED (Coverage: {coverage:.1%}, Bridging: {bridging_ratio:.1%})")
        else:
            print(f"      ❌ Separator X={x_coord} REJECTED (Coverage: {coverage:.1%}, Bridging: {bridging_ratio:.1%})")

    return sorted(final_separators)





def get_word_data_for_detection(page: fitz.Page, pdf_path: str, page_num: int,
                                top_margin_percent=0.10, bottom_margin_percent=0.10) -> list:
    """Extract word data with OCR caching to avoid redundant Tesseract runs."""
    word_data = page.get_text("words")

    if len(word_data) > 0:
        word_data = [(w[4], w[0], w[1], w[2], w[3]) for w in word_data]
    else:
        if _ocr_cache.has_ocr(pdf_path, page_num):
            word_data = _ocr_cache.get_ocr(pdf_path, page_num)
        else:
            try:
                # --- OPTIMIZATION START ---
                # 1. Render at Higher Resolution (Zoom 4.0 = ~300 DPI)
                zoom_level = 4.0
                pix = page.get_pixmap(matrix=fitz.Matrix(zoom_level, zoom_level))
                
                # 2. Convert directly to OpenCV format (Faster than PIL)
                img_np = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, pix.n)
                if pix.n == 3: img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
                elif pix.n == 4: img_np = cv2.cvtColor(img_np, cv2.COLOR_RGBA2BGR)

                # 3. Apply Preprocessing (Thresholding)
                processed_img = preprocess_image_for_ocr(img_np)

                # 4. Optimized Tesseract Config
                # --psm 6: Assume a single uniform block of text (Great for columns/questions)
                # --oem 3: Default engine (LSTM)
                custom_config = r'--oem 3 --psm 6'
                
                data = pytesseract.image_to_data(processed_img, output_type=pytesseract.Output.DICT, config=custom_config)
                
                full_word_data = []
                for i in range(len(data['level'])):
                    text = data['text'][i].strip()
                    if text:
                        # Scale coordinates back to PDF points
                        x1 = data['left'][i] / zoom_level
                        y1 = data['top'][i] / zoom_level
                        x2 = (data['left'][i] + data['width'][i]) / zoom_level
                        y2 = (data['top'][i] + data['height'][i]) / zoom_level
                        full_word_data.append((text, x1, y1, x2, y2))
                
                word_data = full_word_data
                _ocr_cache.set_ocr(pdf_path, page_num, word_data)
                # --- OPTIMIZATION END ---
            except Exception as e:
                print(f"  ❌ OCR Error in detection phase: {e}")
                return []

    # Apply margin filtering
    page_height = page.rect.height
    y_min = page_height * top_margin_percent
    y_max = page_height * (1 - bottom_margin_percent)
    return [d for d in word_data if d[2] >= y_min and d[4] <= y_max]
    


def pixmap_to_numpy(pix: fitz.Pixmap) -> np.ndarray:
    img_data = pix.samples
    img = np.frombuffer(img_data, dtype=np.uint8).reshape(pix.height, pix.width, pix.n)
    if pix.n == 4: img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGR)
    elif pix.n == 3: img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    return img


def extract_native_words_and_convert(fitz_page, scale_factor: float = 2.0) -> list:
    raw_word_data = fitz_page.get_text("words")
    converted_ocr_output = []
    DEFAULT_CONFIDENCE = 99.0

    for x1, y1, x2, y2, word, *rest in raw_word_data:
        if not word.strip(): continue
        x1_pix = int(x1 * scale_factor)
        y1_pix = int(y1 * scale_factor)
        x2_pix = int(x2 * scale_factor)
        y2_pix = int(y2 * scale_factor)
        converted_ocr_output.append({
            'type': 'text',
            'word': word,
            'confidence': DEFAULT_CONFIDENCE,
            'bbox': [x1_pix, y1_pix, x2_pix, y2_pix],
            'y0': y1_pix, 'x0': x1_pix
        })
    return converted_ocr_output





def preprocess_and_ocr_page(original_img: np.ndarray, model, pdf_path: str,
                            page_num: int, fitz_page: fitz.Page,
                            pdf_name: str) -> Tuple[List[Dict[str, Any]], Optional[int]]:
    """
    OPTIMIZED FLOW: 
    1. Run YOLO to find Equations/Tables.
    2. Mask raw text with YOLO boxes.
    3. Run Column Detection on the MASKED data.
    4. Proceed with OCR (Native or High-Res Tesseract Fallback) and Output.
    """
    global GLOBAL_FIGURE_COUNT, GLOBAL_EQUATION_COUNT

    start_time_total = time.time()

    if original_img is None:
        print(f"  ❌ Invalid image for page {page_num}.")
        return None, None

    # ====================================================================
    # --- STEP 1: YOLO DETECTION ---
    # ====================================================================
    start_time_yolo = time.time()
    results = model.predict(source=original_img, conf=CONF_THRESHOLD, imgsz=640, verbose=False)
    
    relevant_detections = []
    if results and results[0].boxes:
        for box in results[0].boxes:
            class_id = int(box.cls[0])
            class_name = model.names[class_id]
            if class_name in TARGET_CLASSES:
                x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
                relevant_detections.append(
                    {'coords': (x1, y1, x2, y2), 'y1': y1, 'class': class_name, 'conf': float(box.conf[0])}
                )

    merged_detections = merge_overlapping_boxes(relevant_detections, IOU_MERGE_THRESHOLD)
    print(f"    [LOG] YOLO found {len(merged_detections)} objects in {time.time() - start_time_yolo:.3f}s.")

    # ====================================================================
    # --- STEP 2: PREPARE DATA FOR COLUMN DETECTION (MASKING) ---
    # ====================================================================
    # Note: This uses the updated 'get_word_data_for_detection' which has its own optimizations
    raw_words_for_layout = get_word_data_for_detection(
        fitz_page, pdf_path, page_num,
        top_margin_percent=0.10, bottom_margin_percent=0.10
    )

    masked_word_data = merge_yolo_into_word_data(raw_words_for_layout, merged_detections, scale_factor=2.0)

    # ====================================================================
    # --- STEP 3: COLUMN DETECTION ---
    # ====================================================================
    page_width_pdf = fitz_page.rect.width
    page_height_pdf = fitz_page.rect.height 
    
    column_detection_params = {
        'cluster_bin_size': 2, 'cluster_smoothing': 2,
        'cluster_min_width': 10, 'cluster_threshold_percentile': 85,
    }

    separators = calculate_x_gutters(masked_word_data, column_detection_params, page_height_pdf)
    
    page_separator_x = None
    if separators:
        central_min = page_width_pdf * 0.35
        central_max = page_width_pdf * 0.65
        central_separators = [s for s in separators if central_min <= s <= central_max]
        
        if central_separators:
            center_x = page_width_pdf / 2
            page_separator_x = min(central_separators, key=lambda x: abs(x - center_x))
            print(f"      ✅ Column Split Confirmed at X={page_separator_x:.1f}")
        else:
            print("      ⚠️ Gutter found off-center. Ignoring.")
    else:
        print("      -> Single Column Layout Confirmed.")

    # ====================================================================
    # --- STEP 4: COMPONENT EXTRACTION (Save Images) ---
    # ====================================================================
    start_time_components = time.time()
    component_metadata = []
    fig_count_page = 0
    eq_count_page = 0

    for detection in merged_detections:
        x1, y1, x2, y2 = detection['coords']
        class_name = detection['class']

        if class_name == 'figure':
            GLOBAL_FIGURE_COUNT += 1
            counter = GLOBAL_FIGURE_COUNT
            component_word = f"FIGURE{counter}"
            fig_count_page += 1
        elif class_name == 'equation':
            GLOBAL_EQUATION_COUNT += 1
            counter = GLOBAL_EQUATION_COUNT
            component_word = f"EQUATION{counter}"
            eq_count_page += 1
        else:
            continue

        component_crop = original_img[y1:y2, x1:x2]
        component_filename = f"{pdf_name}_page{page_num}_{class_name}{counter}.png"
        cv2.imwrite(os.path.join(FIGURE_EXTRACTION_DIR, component_filename), component_crop)

        y_midpoint = (y1 + y2) // 2
        component_metadata.append({
            'type': class_name, 'word': component_word,
            'bbox': [int(x1), int(y1), int(x2), int(y2)],
            'y0': int(y_midpoint), 'x0': int(x1)
        })

    # ====================================================================
    # --- STEP 5: HYBRID OCR (Native Text + Cached Tesseract Fallback) ---
    # ====================================================================
    raw_ocr_output = []
    scale_factor = 2.0 # Pipeline standard scale

    try:
        # Try getting native text first
        raw_ocr_output = extract_native_words_and_convert(fitz_page, scale_factor=scale_factor)
    except Exception as e:
        print(f"  ❌ Native text extraction failed: {e}")

    # If native text is missing, fall back to OCR
    if not raw_ocr_output:
        if _ocr_cache.has_ocr(pdf_path, page_num):
            print(f"  ⚡ Using cached Tesseract OCR for page {page_num}")
            cached_word_data = _ocr_cache.get_ocr(pdf_path, page_num)
            for word_tuple in cached_word_data:
                word_text, x1, y1, x2, y2 = word_tuple
                
                # Scale from PDF points to Pipeline Pixels (2.0)
                x1_pix = int(x1 * scale_factor)
                y1_pix = int(y1 * scale_factor)
                x2_pix = int(x2 * scale_factor)
                y2_pix = int(y2 * scale_factor)
                
                raw_ocr_output.append({
                    'type': 'text', 'word': word_text, 'confidence': 95.0,
                    'bbox': [x1_pix, y1_pix, x2_pix, y2_pix],
                    'y0': y1_pix, 'x0': x1_pix
                })
        else:
            # === START OF OPTIMIZED OCR BLOCK ===
            try:
                # 1. Re-render Page at High Resolution (Zoom 4.0 = ~300 DPI)
                # We do this specifically for OCR accuracy, separate from the pipeline image
                ocr_zoom = 4.0
                pix_ocr = fitz_page.get_pixmap(matrix=fitz.Matrix(ocr_zoom, ocr_zoom))
                
                # Convert PyMuPDF Pixmap to OpenCV format
                img_ocr_np = np.frombuffer(pix_ocr.samples, dtype=np.uint8).reshape(pix_ocr.height, pix_ocr.width, pix_ocr.n)
                if pix_ocr.n == 3: img_ocr_np = cv2.cvtColor(img_ocr_np, cv2.COLOR_RGB2BGR)
                elif pix_ocr.n == 4: img_ocr_np = cv2.cvtColor(img_ocr_np, cv2.COLOR_RGBA2BGR)

                # 2. Preprocess (Binarization)
                # Ensure 'preprocess_image_for_ocr' is defined at top of file!
                processed_img = preprocess_image_for_ocr(img_ocr_np)
                
                # 3. Run Tesseract with Optimized Configuration
                # --oem 3: Default LSTM engine
                # --psm 6: Assume a single uniform block of text (Critical for lists/questions)
                custom_config = r'--oem 3 --psm 6'
                
                hocr_data = pytesseract.image_to_data(
                    processed_img, 
                    output_type=pytesseract.Output.DICT, 
                    config=custom_config
                )
                
                for i in range(len(hocr_data['level'])):
                    text = hocr_data['text'][i].strip()
                    if text and hocr_data['conf'][i] > -1:
                        
                        # 4. Coordinate Mapping
                        # We scanned at Zoom 4.0, but our pipeline expects Zoom 2.0.
                        # Scale Factor = (Target 2.0) / (Source 4.0) = 0.5
                        scale_adjustment = scale_factor / ocr_zoom 
                        
                        x1 = int(hocr_data['left'][i] * scale_adjustment)
                        y1 = int(hocr_data['top'][i] * scale_adjustment)
                        w = int(hocr_data['width'][i] * scale_adjustment)
                        h = int(hocr_data['height'][i] * scale_adjustment)
                        x2 = x1 + w
                        y2 = y1 + h
                        
                        raw_ocr_output.append({
                            'type': 'text', 
                            'word': text, 
                            'confidence': float(hocr_data['conf'][i]),
                            'bbox': [x1, y1, x2, y2], 
                            'y0': y1, 
                            'x0': x1
                        })
            except Exception as e:
                print(f"  ❌ Tesseract OCR Error: {e}")
            # === END OF OPTIMIZED OCR BLOCK ===

    # ====================================================================
    # --- STEP 6: OCR CLEANING AND MERGING ---
    # ====================================================================
    items_to_sort = []
    
    for ocr_word in raw_ocr_output:
        is_suppressed = False
        for component in component_metadata:
            # Do not include words that are inside figure/equation boxes
            ioa = calculate_ioa(ocr_word['bbox'], component['bbox'])
            if ioa > IOA_SUPPRESSION_THRESHOLD:
                is_suppressed = True
                break
        if not is_suppressed:
            items_to_sort.append(ocr_word)

    # Add figures/equations back into the flow as "words"
    items_to_sort.extend(component_metadata)

    # ====================================================================
    # --- STEP 7: LINE-BASED SORTING ---
    # ====================================================================
    items_to_sort.sort(key=lambda x: (x['y0'], x['x0']))
    lines = []

    for item in items_to_sort:
        placed = False
        for line in lines:
            y_ref = min(it['y0'] for it in line)
            if abs(y_ref - item['y0']) < LINE_TOLERANCE:
                line.append(item)
                placed = True
                break
        if not placed and item['type'] in ['equation', 'figure']:
            for line in lines:
                y_ref = min(it['y0'] for it in line)
                if abs(y_ref - item['y0']) < 20:
                    line.append(item)
                    placed = True
                    break
        if not placed:
            lines.append([item])

    for line in lines:
        line.sort(key=lambda x: x['x0'])

    final_output = []
    for line in lines:
        for item in line:
            data_item = {"word": item["word"], "bbox": item["bbox"], "type": item["type"]}
            if 'tag' in item: data_item['tag'] = item['tag']
            final_output.append(data_item)

    return final_output, page_separator_x



def run_single_pdf_preprocessing(pdf_path: str, preprocessed_json_path: str) -> Optional[str]:
    global GLOBAL_FIGURE_COUNT, GLOBAL_EQUATION_COUNT

    GLOBAL_FIGURE_COUNT = 0
    GLOBAL_EQUATION_COUNT = 0
    _ocr_cache.clear()

    print("\n" + "=" * 80)
    print("--- 1. STARTING OPTIMIZED YOLO/OCR PREPROCESSING PIPELINE ---")
    print("=" * 80)

    if not os.path.exists(pdf_path):
        print(f"❌ FATAL ERROR: Input PDF not found at {pdf_path}.")
        return None

    os.makedirs(os.path.dirname(preprocessed_json_path), exist_ok=True)
    os.makedirs(FIGURE_EXTRACTION_DIR, exist_ok=True)

    model = YOLO(WEIGHTS_PATH)
    pdf_name = os.path.splitext(os.path.basename(pdf_path))[0]

    try:
        doc = fitz.open(pdf_path)
        print(f"✅ Opened PDF: {pdf_name} ({doc.page_count} pages)")
    except Exception as e:
        print(f"❌ ERROR loading PDF file: {e}")
        return None

    all_pages_data = []
    total_pages_processed = 0
    mat = fitz.Matrix(2.0, 2.0)

    print("\n[STEP 1.2: ITERATING PAGES - IN-MEMORY PROCESSING]")

    for page_num_0_based in range(doc.page_count):
        page_num = page_num_0_based + 1
        print(f"  -> Processing Page {page_num}/{doc.page_count}...")

        fitz_page = doc.load_page(page_num_0_based)

        try:
            pix = fitz_page.get_pixmap(matrix=mat)
            original_img = pixmap_to_numpy(pix)
        except Exception as e:
            print(f"  ❌ Error converting page {page_num} to image: {e}")
            continue

        final_output, page_separator_x = preprocess_and_ocr_page(
            original_img,
            model,
            pdf_path,
            page_num,
            fitz_page,
            pdf_name
        )

        if final_output is not None:
            page_data = {
                "page_number": page_num,
                "data": final_output,
                "column_separator_x": page_separator_x
            }
            all_pages_data.append(page_data)
            total_pages_processed += 1
        else:
            print(f"  ❌ Skipped page {page_num} due to processing error.")

    doc.close()

    if all_pages_data:
        try:
            with open(preprocessed_json_path, 'w') as f:
                json.dump(all_pages_data, f, indent=4)
            print(f"\n  ✅ Combined structured OCR JSON saved to: {os.path.basename(preprocessed_json_path)}")
        except Exception as e:
            print(f"❌ ERROR saving combined JSON output: {e}")
            return None
    else:
        print("❌ WARNING: No page data generated. Halting pipeline.")
        return None

    print("\n" + "=" * 80)
    print(f"--- YOLO/OCR PREPROCESSING COMPLETE ({total_pages_processed} pages processed) ---")
    print("=" * 80)

    return preprocessed_json_path


# ============================================================================
# --- PHASE 2: LAYOUTLMV3 INFERENCE FUNCTIONS ---
# ============================================================================

class LayoutLMv3ForTokenClassification(nn.Module):
    def __init__(self, num_labels: int = NUM_LABELS):
        super().__init__()
        self.num_labels = num_labels
        config = LayoutLMv3Config.from_pretrained("microsoft/layoutlmv3-base", num_labels=num_labels)
        self.layoutlmv3 = LayoutLMv3Model.from_pretrained("microsoft/layoutlmv3-base", config=config)
        self.classifier = nn.Linear(config.hidden_size, num_labels)
        self.crf = CRF(num_labels)
        self.init_weights()

    def init_weights(self):
        nn.init.xavier_uniform_(self.classifier.weight)
        if self.classifier.bias is not None: nn.init.zeros_(self.classifier.bias)

    def forward(self, input_ids: torch.Tensor, bbox: torch.Tensor, attention_mask: torch.Tensor, labels: Optional[torch.Tensor] = None):
        outputs = self.layoutlmv3(input_ids=input_ids, bbox=bbox, attention_mask=attention_mask, return_dict=True)
        sequence_output = outputs.last_hidden_state
        emissions = self.classifier(sequence_output)
        mask = attention_mask.bool()
        if labels is not None:
            loss = -self.crf(emissions, labels, mask=mask).mean()
            return loss
        else:
            return self.crf.viterbi_decode(emissions, mask=mask)

def _merge_integrity(all_token_data: List[Dict[str, Any]],
                      column_separator_x: Optional[int]) -> List[List[Dict[str, Any]]]:
    """Splits the token data objects into column chunks based on a separator."""
    if column_separator_x is None:
        print("    -> No column separator. Treating as one chunk.")
        return [all_token_data]

    left_column_tokens, right_column_tokens = [], []
    for token_data in all_token_data:
        bbox_raw = token_data['bbox_raw_pdf_space']
        center_x = (bbox_raw[0] + bbox_raw[2]) / 2
        if center_x < column_separator_x:
            left_column_tokens.append(token_data)
        else:
            right_column_tokens.append(token_data)

    chunks = [c for c in [left_column_tokens, right_column_tokens] if c]
    print(f"    -> Data split into {len(chunks)} column chunk(s) using separator X={column_separator_x}.")
    return chunks

def run_inference_and_get_raw_words(pdf_path: str, model_path: str,
                                    preprocessed_json_path: str,
                                    column_detection_params: Optional[Dict] = None) -> List[Dict[str, Any]]:
    print("\n" + "=" * 80)
    print("--- 2. STARTING LAYOUTLMV3 INFERENCE PIPELINE (Raw Word Output) ---")
    print("=" * 80)

    tokenizer = LayoutLMv3TokenizerFast.from_pretrained("microsoft/layoutlmv3-base")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"  -> Using device: {device}")

    try:
        model = LayoutLMv3ForTokenClassification(num_labels=NUM_LABELS)
        checkpoint = torch.load(model_path, map_location=device)
        model_state = checkpoint.get('model_state_dict', checkpoint)
        fixed_state_dict = {key.replace('layoutlm.', 'layoutlmv3.'): value for key, value in model_state.items()}
        model.load_state_dict(fixed_state_dict)
        model.to(device)
        model.eval()
        print(f"✅ LayoutLMv3 Model loaded successfully from {os.path.basename(model_path)}.")
    except Exception as e:
        print(f"❌ FATAL ERROR during LayoutLMv3 model loading: {e}")
        return []

    try:
        with open(preprocessed_json_path, 'r', encoding='utf-8') as f:
            preprocessed_data = json.load(f)
        print(f"✅ Loaded preprocessed data with {len(preprocessed_data)} pages.")
    except Exception:
        print("❌ Error loading preprocessed JSON.")
        return []

    try:
        doc = fitz.open(pdf_path)
    except Exception:
        print("❌ Error loading PDF.")
        return []

    final_page_predictions = []
    CHUNK_SIZE = 500

    for page_data in preprocessed_data:
        page_num_1_based = page_data['page_number']
        page_num_0_based = page_num_1_based - 1
        page_raw_predictions = []
        print(f"\n  *** Processing Page {page_num_1_based} ({len(page_data['data'])} raw tokens) ***")

        fitz_page = doc.load_page(page_num_0_based)
        page_width, page_height = fitz_page.rect.width, fitz_page.rect.height
        print(f"    -> Page dimensions: {page_width:.0f}x{page_height:.0f} (PDF points).")

        all_token_data = []
        scale_factor = 2.0

        for item in page_data['data']:
            raw_yolo_bbox = item['bbox']
            bbox_pdf = [
                int(raw_yolo_bbox[0] / scale_factor), int(raw_yolo_bbox[1] / scale_factor),
                int(raw_yolo_bbox[2] / scale_factor), int(raw_yolo_bbox[3] / scale_factor)
            ]
            normalized_bbox = [
                max(0, min(1000, int(1000 * bbox_pdf[0] / page_width))),
                max(0, min(1000, int(1000 * bbox_pdf[1] / page_height))),
                max(0, min(1000, int(1000 * bbox_pdf[2] / page_width))),
                max(0, min(1000, int(1000 * bbox_pdf[3] / page_height)))
            ]
            all_token_data.append({
                "word": item['word'],
                "bbox_raw_pdf_space": bbox_pdf,
                "bbox_normalized": normalized_bbox,
                "item_original_data": item
            })

        if not all_token_data: continue

        column_separator_x = page_data.get('column_separator_x', None)
        if column_separator_x is not None:
            print(f"    -> Using SAVED column separator: X={column_separator_x}")
        else:
            print("    -> No column separator found. Assuming single chunk.")

        token_chunks = _merge_integrity(all_token_data, column_separator_x)
        total_chunks = len(token_chunks)

        for chunk_idx, chunk_tokens in enumerate(token_chunks):
            if not chunk_tokens: continue

            chunk_words = [t['word'] for t in chunk_tokens]
            chunk_normalized_bboxes = [t['bbox_normalized'] for t in chunk_tokens]

            total_sub_chunks = (len(chunk_words) + CHUNK_SIZE - 1) // CHUNK_SIZE
            for i in range(0, len(chunk_words), CHUNK_SIZE):
                sub_chunk_idx = i // CHUNK_SIZE + 1
                sub_words = chunk_words[i:i + CHUNK_SIZE]
                sub_bboxes = chunk_normalized_bboxes[i:i + CHUNK_SIZE]
                sub_tokens_data = chunk_tokens[i:i + CHUNK_SIZE]

                print(f"      -> Chunk {chunk_idx + 1}/{total_chunks}, Sub-chunk {sub_chunk_idx}/{total_sub_chunks}: {len(sub_words)} words. Running Inference...")

                encoded_input = tokenizer(
                    sub_words, boxes=sub_bboxes, truncation=True, padding="max_length",
                    max_length=512, return_tensors="pt"
                )
                input_ids = encoded_input['input_ids'].to(device)
                bbox = encoded_input['bbox'].to(device)
                attention_mask = encoded_input['attention_mask'].to(device)

                with torch.no_grad():
                    predictions_int_list = model(input_ids, bbox, attention_mask)

                if not predictions_int_list: continue
                predictions_int = predictions_int_list[0]
                word_ids = encoded_input.word_ids()
                word_idx_to_pred_id = {}

                for token_idx, word_idx in enumerate(word_ids):
                    if word_idx is not None and word_idx < len(sub_words):
                        if word_idx not in word_idx_to_pred_id:
                            word_idx_to_pred_id[word_idx] = predictions_int[token_idx]

                for current_word_idx in range(len(sub_words)):
                    pred_id_or_tensor = word_idx_to_pred_id.get(current_word_idx, 0)
                    pred_id = pred_id_or_tensor.item() if torch.is_tensor(pred_id_or_tensor) else pred_id_or_tensor
                    predicted_label = ID_TO_LABEL[pred_id]
                    original_token = sub_tokens_data[current_word_idx]
                    page_raw_predictions.append({
                        "word": original_token['word'],
                        "bbox": original_token['bbox_raw_pdf_space'],
                        "predicted_label": predicted_label,
                        "page_number": page_num_1_based
                    })

        if page_raw_predictions:
            final_page_predictions.append({
                "page_number": page_num_1_based,
                "data": page_raw_predictions
            })
            print(f"  *** Page {page_num_1_based} Finalized: {len(page_raw_predictions)} labeled words. ***")

    doc.close()
    print("\n" + "=" * 80)
    print("--- LAYOUTLMV3 INFERENCE COMPLETE ---")
    print("=" * 80)
    return final_page_predictions


def create_label_studio_span(page_results, start_idx, end_idx, label):
    entity_words = [page_results[i]['word'] for i in range(start_idx, end_idx + 1)]
    entity_bboxes = [page_results[i]['bbox'] for i in range(start_idx, end_idx + 1)]
    x0 = min(bbox[0] for bbox in entity_bboxes)
    y0 = min(bbox[1] for bbox in entity_bboxes)
    x1 = max(bbox[2] for bbox in entity_bboxes)
    y1 = max(bbox[3] for bbox in entity_bboxes)
    all_words_on_page = [r['word'] for r in page_results]
    start_char = len(" ".join(all_words_on_page[:start_idx]))
    if start_idx != 0: start_char += 1
    end_char = start_char + len(" ".join(entity_words))
    span_text = " ".join(entity_words)
    return {
        "from_name": "label", "to_name": "text", "type": "labels",
        "value": {
            "start": start_char, "end": end_char, "text": span_text,
            "labels": [label],
            "bbox": {"x": x0, "y": y0, "width": x1 - x0, "height": y1 - y0}
        }, "score": 0.99
    }

def convert_raw_predictions_to_label_studio(page_data_list, output_path: str):
    final_tasks = []
    print("\n[PHASE: LABEL STUDIO CONVERSION]")
    for page_data in page_data_list:
        page_num = page_data['page_number']
        page_results = page_data['data']
        if not page_results: continue
        original_words = [r['word'] for r in page_results]
        text_string = " ".join(original_words)
        results = []
        current_entity_label = None
        current_entity_start_word_index = None

        for i, pred_item in enumerate(page_results):
            label = pred_item['predicted_label']
            tag_only = label.split('-', 1)[-1] if '-' in label else label
            if label.startswith('B-'):
                if current_entity_label:
                    results.append(create_label_studio_span(page_results, current_entity_start_word_index, i - 1, current_entity_label))
                current_entity_label = tag_only
                current_entity_start_word_index = i
            elif label.startswith('I-') and current_entity_label == tag_only:
                continue
            else:
                if current_entity_label:
                    results.append(create_label_studio_span(page_results, current_entity_start_word_index, i - 1, current_entity_label))
                current_entity_label = None
                current_entity_start_word_index = None
        if current_entity_label:
            results.append(create_label_studio_span(page_results, current_entity_start_word_index, len(page_results) - 1, current_entity_label))

        final_tasks.append({
            "data": {
                "text": text_string, "original_words": original_words,
                "original_bboxes": [r['bbox'] for r in page_results]
            },
            "annotations": [{"result": results}],
            "meta": {"page_number": page_num}
        })
    with open(output_path, "w", encoding='utf-8') as f:
        json.dump(final_tasks, f, indent=2, ensure_ascii=False)
    print(f"\n✅ Label Studio tasks saved to {output_path}.")


# ============================================================================
# --- PHASE 3: BIO TO STRUCTURED JSON DECODER ---
# ============================================================================



def convert_bio_to_structured_json_relaxed(input_path: str, output_path: str) -> Optional[List[Dict[str, Any]]]:
    print("\n" + "=" * 80)
    print("--- 3. STARTING BIO TO STRUCTURED JSON DECODING ---")
    print("=" * 80)
    try:
        with open(input_path, 'r', encoding='utf-8') as f:
            predictions_by_page = json.load(f)
    except Exception as e:
        print(f"❌ Error loading raw prediction file: {e}")
        return None

    predictions = []
    for page_item in predictions_by_page:
        if isinstance(page_item, dict) and 'data' in page_item:
            predictions.extend(page_item['data'])

    structured_data = []
    current_item = None
    current_option_key = None
    current_passage_buffer = []
    current_text_buffer = []
    first_question_started = False
    last_entity_type = None
    just_finished_i_option = False
    is_in_new_passage = False

    def finalize_passage_to_item(item, passage_buffer):
        if passage_buffer:
            passage_text = re.sub(r'\s{2,}', ' ', ' '.join(passage_buffer)).strip()
            if item.get('passage'): item['passage'] += ' ' + passage_text
            else: item['passage'] = passage_text
        passage_buffer.clear()

    for item in predictions:
        word = item['word']
        label = item['predicted_label']
        entity_type = label[2:].strip() if label.startswith(('B-', 'I-')) else None
        current_text_buffer.append(word)
        previous_entity_type = last_entity_type
        is_passage_label = (entity_type == 'PASSAGE')

        if not first_question_started:
            if label != 'B-QUESTION' and not is_passage_label:
                just_finished_i_option = False
                is_in_new_passage = False
                continue
            if is_passage_label:
                current_passage_buffer.append(word)
                last_entity_type = 'PASSAGE'
                just_finished_i_option = False
                is_in_new_passage = False
                continue

        if label == 'B-QUESTION':
            if not first_question_started:
                header_text = ' '.join(current_text_buffer[:-1]).strip()
                if header_text or current_passage_buffer:
                    metadata_item = {'type': 'METADATA', 'passage': ''}
                    finalize_passage_to_item(metadata_item, current_passage_buffer)
                    if header_text: metadata_item['text'] = header_text
                    structured_data.append(metadata_item)
                first_question_started = True
                current_text_buffer = [word]

            if current_item is not None:
                finalize_passage_to_item(current_item, current_passage_buffer)
                current_item['text'] = ' '.join(current_text_buffer[:-1]).strip()
                structured_data.append(current_item)
                current_text_buffer = [word]

            current_item = {
                'question': word, 'options': {}, 'answer': '', 'passage': '', 'text': ''
            }
            current_option_key = None
            last_entity_type = 'QUESTION'
            just_finished_i_option = False
            is_in_new_passage = False
            continue

        if current_item is not None:
            if is_in_new_passage:
                # 🔑 Robust Initialization and Appending for 'new_passage'
                if 'new_passage' not in current_item:
                    current_item['new_passage'] = word
                else:
                    current_item['new_passage'] += f' {word}'

                if label.startswith('B-') or (label.startswith('I-') and entity_type != 'PASSAGE'):
                    is_in_new_passage = False
                if label.startswith(('B-', 'I-')): last_entity_type = entity_type
                continue
            is_in_new_passage = False

            if label.startswith('B-'):
                if entity_type in ['QUESTION', 'OPTION', 'ANSWER', 'SECTION_HEADING']:
                    finalize_passage_to_item(current_item, current_passage_buffer)
                    current_passage_buffer = []
                last_entity_type = entity_type
                if entity_type == 'PASSAGE':
                    if previous_entity_type == 'OPTION' and just_finished_i_option:
                        current_item['new_passage'] = word # Initialize the new passage start
                        is_in_new_passage = True
                    else:
                        current_passage_buffer.append(word)
                elif entity_type == 'OPTION':
                    current_option_key = word
                    current_item['options'][current_option_key] = word
                    just_finished_i_option = False
                elif entity_type == 'ANSWER':
                    current_item['answer'] = word
                    current_option_key = None
                    just_finished_i_option = False
                elif entity_type == 'QUESTION':
                    current_item['question'] += f' {word}'
                    just_finished_i_option = False

            elif label.startswith('I-'):
                if entity_type == 'QUESTION':
                    current_item['question'] += f' {word}'
                elif entity_type == 'PASSAGE':
                    if previous_entity_type == 'OPTION' and just_finished_i_option:
                        current_item['new_passage'] = word # Initialize the new passage start
                        is_in_new_passage = True
                    else:
                        if not current_passage_buffer: last_entity_type = 'PASSAGE'
                        current_passage_buffer.append(word)
                elif entity_type == 'OPTION' and current_option_key is not None:
                    current_item['options'][current_option_key] += f' {word}'
                    just_finished_i_option = True
                elif entity_type == 'ANSWER':
                    current_item['answer'] += f' {word}'
                just_finished_i_option = (entity_type == 'OPTION')

            elif label == 'O':
                if last_entity_type == 'QUESTION':
                    current_item['question'] += f' {word}'
                just_finished_i_option = False

    if current_item is not None:
        finalize_passage_to_item(current_item, current_passage_buffer)
        current_item['text'] = ' '.join(current_text_buffer).strip()
        structured_data.append(current_item)
    
    for item in structured_data:
        item['text'] = re.sub(r'\s{2,}', ' ', item['text']).strip()
        if 'new_passage' in item:
            item['new_passage'] = re.sub(r'\s{2,}', ' ', item['new_passage']).strip()

    try:
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(structured_data, f, indent=2, ensure_ascii=False)
    except Exception: pass

    return structured_data

def create_query_text(entry: Dict[str, Any]) -> str:
    """Combines question and options into a single string for similarity matching."""
    query_parts = []
    if entry.get("question"):
        query_parts.append(entry["question"])

    for key in ["options", "options_text"]:
        options = entry.get(key)
        if options and isinstance(options, dict):
            for value in options.values():
                if value and isinstance(value, str):
                    query_parts.append(value)
    return " ".join(query_parts)

    

def calculate_similarity(doc1: str, doc2: str) -> float:
    """Calculates Cosine Similarity between two text strings."""
    if not doc1 or not doc2:
        return 0.0

    def clean_text(text):
        return re.sub(r'^\s*[\(\d\w]+\.?\s*', '', text, flags=re.MULTILINE)

    clean_doc1 = clean_text(doc1)
    clean_doc2 = clean_text(doc2)
    corpus = [clean_doc1, clean_doc2]

    try:
        vectorizer = CountVectorizer(stop_words='english', lowercase=True, token_pattern=r'(?u)\b\w\w+\b')
        tfidf_matrix = vectorizer.fit_transform(corpus)
        if tfidf_matrix.shape[1] == 0:
            return 0.0
        vectors = tfidf_matrix.toarray()
        # Handle cases where vectors might be empty or too short
        if len(vectors) < 2:
             return 0.0
        score = cosine_similarity(vectors[0:1], vectors[1:2])[0][0]
        return score
    except Exception:
        return 0.0





def process_context_linking(data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """
    Links questions to passages based on 'passage' flow vs 'new_passage' priority.
    Includes 'Decay Logic': If 2 consecutive questions fail to match the active passage,
    the passage context is dropped to prevent false positives downstream.
    """
    print("\n" + "=" * 80)
    print("--- STARTING CONTEXT LINKING (WITH DECAY LOGIC) ---")
    print("=" * 80)
    
    if not data: return []

    # --- PHASE 1: IDENTIFY PASSAGE DEFINERS ---
    passage_definer_indices = []
    for i, entry in enumerate(data):
        if entry.get("passage") and entry["passage"].strip():
            passage_definer_indices.append(i)
        if entry.get("new_passage") and entry["new_passage"].strip():
            if i not in passage_definer_indices:
                passage_definer_indices.append(i)

    # --- PHASE 2: CONTEXT TRANSFER & LINKING ---
    current_passage_text = None
    current_new_passage_text = None
    
    # NEW: Counter to track consecutive linking failures
    consecutive_failures = 0 
    MAX_CONSECUTIVE_FAILURES = 2

    for i, entry in enumerate(data):
        item_type = entry.get("type", "Question")

        # A. UNCONDITIONALLY UPDATE CONTEXTS (And Reset Decay Counter)
        if entry.get("passage") and entry["passage"].strip():
            current_passage_text = entry["passage"]
            consecutive_failures = 0 # Reset because we have fresh explicit context
            # print(f"  [Flow] Updated Standard Context from Item {i}")
        
        if entry.get("new_passage") and entry["new_passage"].strip():
            current_new_passage_text = entry["new_passage"]
            # We don't necessarily reset standard failures here as this is a local override

        # B. QUESTION LINKING
        if entry.get("question") and item_type != "METADATA":
            combined_query = create_query_text(entry)
            
            # Skip if query is too short (noise)
            if len(combined_query.strip()) < 5: 
                continue

            # Calculate scores
            score_old = calculate_similarity(current_passage_text, combined_query) if current_passage_text else 0.0
            score_new = calculate_similarity(current_new_passage_text, combined_query) if current_new_passage_text else 0.0

            q_preview = entry['question'][:30] + '...'
            
            # RESOLUTION LOGIC
            linked = False
            
            # 1. Prefer New Passage if significantly better
            if current_new_passage_text and (score_new > score_old + RESOLUTION_MARGIN) and (score_new >= SIMILARITY_THRESHOLD):
                entry["passage"] = current_new_passage_text
                print(f"  [Linker] 🚀 Q{i} ('{q_preview}') -> NEW PASSAGE (Score: {score_new:.3f})")
                linked = True
                # Note: We do not reset 'consecutive_failures' for the standard passage here,
                # because we matched the *new* passage, not the standard one.
            
            # 2. Otherwise use Standard Passage if it meets threshold
            elif current_passage_text and (score_old >= SIMILARITY_THRESHOLD):
                entry["passage"] = current_passage_text
                print(f"  [Linker] ✅ Q{i} ('{q_preview}') -> STANDARD PASSAGE (Score: {score_old:.3f})")
                linked = True
                consecutive_failures = 0 # Success! Reset the kill switch.
            
            if not linked:
                # 3. DECAY LOGIC
                if current_passage_text:
                    consecutive_failures += 1
                    print(f"  [Linker] ⚠️ Q{i} NOT LINKED. (Failures: {consecutive_failures}/{MAX_CONSECUTIVE_FAILURES})")
                    
                    if consecutive_failures >= MAX_CONSECUTIVE_FAILURES:
                        print(f"  [Linker] 🗑️  Context dropped due to {consecutive_failures} consecutive misses.")
                        current_passage_text = None
                        consecutive_failures = 0
                else:
                    print(f"  [Linker] ⚠️ Q{i} NOT LINKED (No active context).")

    # --- PHASE 3: CLEANUP AND INTERPOLATION ---
    print("  [Linker] Running Cleanup & Interpolation...")
    
    # 3A. Self-Correction (Remove weak links)
    for i in passage_definer_indices:
        entry = data[i]
        if entry.get("question") and entry.get("type") != "METADATA":
            passage_to_check = entry.get("passage") or entry.get("new_passage")
            if passage_to_check:
                self_sim = calculate_similarity(passage_to_check, create_query_text(entry))
                if self_sim < SIMILARITY_THRESHOLD:
                    entry["passage"] = ""
                    if "new_passage" in entry: entry["new_passage"] = ""
                    print(f"  [Cleanup] Removed weak link for Q{i}")

    # 3B. Interpolation (Fill gaps)
    # We only interpolate if the gap is strictly 1 question wide to avoid undoing the decay logic
    for i in range(1, len(data) - 1):
        current_entry = data[i]
        is_gap = current_entry.get("question") and not current_entry.get("passage")
        if is_gap:
            prev_p = data[i - 1].get("passage")
            next_p = data[i + 1].get("passage")
            if prev_p and next_p and (prev_p == next_p) and prev_p.strip():
                current_entry["passage"] = prev_p
                print(f"  [Linker] 🥪 Q{i} Interpolated from neighbors.")

    return data


def correct_misaligned_options(structured_data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    print("\n" + "=" * 80)
    print("--- 5. STARTING POST-PROCESSING: OPTION ALIGNMENT CORRECTION ---")
    print("=" * 80)
    tag_pattern = re.compile(r'(EQUATION\d+|FIGURE\d+)')
    corrected_count = 0
    for item in structured_data:
        if item.get('type') in ['METADATA']: continue
        options = item.get('options')
        if not options or len(options) < 2: continue
        option_keys = list(options.keys())
        for i in range(len(option_keys) - 1):
            current_key = option_keys[i]
            next_key = option_keys[i + 1]
            current_value = options[current_key].strip()
            next_value = options[next_key].strip()
            is_current_empty = current_value == current_key
            content_in_next = next_value.replace(next_key, '', 1).strip()
            tags_in_next = tag_pattern.findall(content_in_next)
            has_two_tags = len(tags_in_next) == 2
            if is_current_empty and has_two_tags:
                tag_to_move = tags_in_next[0]
                options[current_key] = f"{current_key} {tag_to_move}".strip()
                options[next_key] = f"{next_key} {tags_in_next[1]}".strip()
                corrected_count += 1
    print(f"✅ Option alignment correction finished. Total corrections: {corrected_count}.")
    return structured_data




# ============================================================================
# --- PHASE 4: IMAGE EMBEDDING (Base64) ---
# ============================================================================

def get_base64_for_file(filepath: str) -> str:
    try:
        with open(filepath, 'rb') as f:
            return base64.b64encode(f.read()).decode('utf-8')
    except Exception as e:
        print(f"  ❌ Error encoding file {filepath}: {e}")
        return ""

def embed_images_as_base64_in_memory(structured_data: List[Dict[str, Any]], figure_extraction_dir: str) -> List[Dict[str, Any]]:
    print("\n" + "=" * 80)
    print("--- 4. STARTING IMAGE EMBEDDING (Base64) ---")
    print("=" * 80)
    if not structured_data: return []
    image_files = glob.glob(os.path.join(figure_extraction_dir, "*.png"))
    image_lookup = {}
    tag_regex = re.compile(r'(figure|equation)(\d+)', re.IGNORECASE)
    for filepath in image_files:
        filename = os.path.basename(filepath)
        match = re.search(r'_(figure|equation)(\d+)\.png$', filename, re.IGNORECASE)
        if match:
            key = f"{match.group(1).upper()}{match.group(2)}"
            image_lookup[key] = filepath
    print(f"  -> Found {len(image_lookup)} image components.")
    final_structured_data = []
    for item in structured_data:
        text_fields = [item.get('question', ''), item.get('passage', '')]
        if 'options' in item:
            for opt_val in item['options'].values(): text_fields.append(opt_val)
        if 'new_passage' in item: text_fields.append(item['new_passage'])
        unique_tags_to_embed = set()
        for text in text_fields:
            if not text: continue
            for match in tag_regex.finditer(text):
                tag = match.group(0).upper()
                if tag in image_lookup: unique_tags_to_embed.add(tag)
        for tag in sorted(list(unique_tags_to_embed)):
            filepath = image_lookup[tag]
            base64_code = get_base64_for_file(filepath)
            base_key = tag.replace(' ', '').lower()
            item[base_key] = base64_code
        final_structured_data.append(item)
    print(f"✅ Image embedding complete.")
    return final_structured_data

# ============================================================================
# --- MAIN FUNCTION ---
# ============================================================================

def run_document_pipeline(input_pdf_path: str, layoutlmv3_model_path: str, label_studio_output_path: str) -> Optional[List[Dict[str, Any]]]:
    if not os.path.exists(input_pdf_path): return None
    
    print("\n" + "#" * 80)
    print("### STARTING OPTIMIZED FULL DOCUMENT ANALYSIS PIPELINE ###")
    print("#" * 80)
    
    pdf_name = os.path.splitext(os.path.basename(input_pdf_path))[0]
    temp_pipeline_dir = os.path.join(tempfile.gettempdir(), f"pipeline_run_{pdf_name}_{os.getpid()}")
    os.makedirs(temp_pipeline_dir, exist_ok=True)
    
    preprocessed_json_path = os.path.join(temp_pipeline_dir, f"{pdf_name}_preprocessed.json")
    raw_output_path = os.path.join(temp_pipeline_dir, f"{pdf_name}_raw_predictions.json")
    structured_intermediate_output_path = os.path.join(temp_pipeline_dir, f"{pdf_name}_structured_intermediate.json")
    
    final_result = None
    try:
        # Phase 1: Preprocessing with YOLO First + Masking
        preprocessed_json_path_out = run_single_pdf_preprocessing(input_pdf_path, preprocessed_json_path)
        if not preprocessed_json_path_out: return None
        
        # Phase 2: Inference
        page_raw_predictions_list = run_inference_and_get_raw_words(
            input_pdf_path, layoutlmv3_model_path, preprocessed_json_path_out
        )
        if not page_raw_predictions_list: return None
        
        with open(raw_output_path, 'w', encoding='utf-8') as f:
            json.dump(page_raw_predictions_list, f, indent=4)
            
        # Phase 3: Decoding
        structured_data_list = convert_bio_to_structured_json_relaxed(
            raw_output_path, structured_intermediate_output_path
        )
        if not structured_data_list: return None
        structured_data_list = correct_misaligned_options(structured_data_list)
        structured_data_list = process_context_linking(structured_data_list)
        
        try:
            convert_raw_predictions_to_label_studio(page_raw_predictions_list, label_studio_output_path)
        except Exception as e:
            print(f"❌ Error during Label Studio conversion: {e}")
            
        # Phase 4: Embedding
        final_result = embed_images_as_base64_in_memory(structured_data_list, FIGURE_EXTRACTION_DIR)
        
    except Exception as e:
        print(f"❌ FATAL ERROR: {e}")
        import traceback
        traceback.print_exc()
        return None
        
    finally:
        try:
            for f in glob.glob(os.path.join(temp_pipeline_dir, '*')):
                os.remove(f)
            os.rmdir(temp_pipeline_dir)
        except Exception: pass
        
    print("\n" + "#" * 80)
    print("### OPTIMIZED PIPELINE EXECUTION COMPLETE ###")
    print("#" * 80)
    return final_result


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Complete Pipeline")
    parser.add_argument("--input_pdf", type=str, required=True, help="Input PDF")
    parser.add_argument("--layoutlmv3_model_path", type=str, default=DEFAULT_LAYOUTLMV3_MODEL_PATH, help="Model Path")
    parser.add_argument("--ls_output_path", type=str, default=None, help="Label Studio Output Path")
    args = parser.parse_args()
    
    pdf_name = os.path.splitext(os.path.basename(args.input_pdf))[0]
    final_output_path = os.path.abspath(f"{pdf_name}_final_output_embedded.json")
    ls_output_path = os.path.abspath(args.ls_output_path if args.ls_output_path else f"{pdf_name}_label_studio_tasks.json")
    
    final_json_data = run_document_pipeline(args.input_pdf, args.layoutlmv3_model_path, ls_output_path)
    
    if final_json_data:
        with open(final_output_path, 'w', encoding='utf-8') as f:
            json.dump(final_json_data, f, indent=2, ensure_ascii=False)
        print(f"\n✅ Final Data Saved: {final_output_path}")
    else:
        print("\n❌ Pipeline Failed.")
        sys.exit(1)