import cv2
from IPython.display import display
from PIL import Image as PILImage
import numpy as np

def get_text_segment_bbox(img, full_text, segment_text, original_bbox, occurrence=1, color=(0, 0, 255), thickness=2, font=cv2.FONT_HERSHEY_SIMPLEX):
    """
    Calculate the bbox for a segment of text within a larger text string, accounting for word spacing.

    Args:
        img: Input image (numpy array)
        full_text: Complete text string
        segment_text: The portion of text to focus on
        original_bbox: Original bounding box coordinates [x1, y1, x2, y2]
        occurrence: The occurrence number of the segment_text in full_text (1-based)
        color: Text and box color in BGR format
        thickness: Line thickness
        font: OpenCV font type
    """
    if segment_text not in full_text:
        raise ValueError("Segment text must be part of the full text")
    
    # Find the start index of the specified occurrence
    split_text = full_text.split(segment_text)
    if len(split_text) <= occurrence:
        raise ValueError(f"Segment text does not occur {occurrence} times in the full text.")
    
    # Calculate the number of characters before the target segment
    prefix = segment_text.join(split_text[:-1])
    prefix_length = len(prefix)
    
    # Calculate total width of the prefix
    # Use binary search to find the font scale that matches original_bbox width
    original_width = original_bbox[2] - original_bbox[0]
    
    def find_font_scale(target_width, text, font, thickness, tolerance=1):
        low = 0.1
        high = 5.0
        best_scale = low
        while high - low > 1e-3:
            mid = (low + high) / 2
            (w, h), _ = cv2.getTextSize(text, font, mid, thickness)
            if w < target_width:
                best_scale = mid
                low = mid
            else:
                high = mid
        return best_scale
    
    # Find font scale for full text
    font_scale = find_font_scale(original_width, full_text, font, thickness)
    
    # Calculate space width
    (space_width, _), _ = cv2.getTextSize(" ", font, font_scale, thickness)
    
    # Calculate the width up to the prefix
    prefix_width = 0
    for char in prefix:
        if char == ' ':
            cw = space_width
        else:
            (cw, _) = cv2.getTextSize(char, font, font_scale, thickness)
        prefix_width += cw
    
    # Calculate the width of the segment text
    (segment_width, segment_height), baseline = cv2.getTextSize(segment_text, font, font_scale, thickness)
    
    # Calculate new bbox coordinates
    new_x1 = original_bbox[0] + prefix_width
    new_x2 = new_x1 + segment_width
    new_y1 = original_bbox[1]
    new_y2 = original_bbox[3]
    
    # Ensure bbox is within image boundaries
    img_height, img_width = img.shape[:2]
    new_x1 = max(new_x1, 0)
    new_x2 = min(new_x2, img_width)
    new_y1 = max(new_y1, 0)
    new_y2 = min(new_y2, img_height)
    
    new_bbox = [int(new_x1), int(new_y1), int(new_x2), int(new_y2)]
    
    # Draw the new bounding box
    cv2.rectangle(img, (new_bbox[0], new_bbox[1]), (new_bbox[2], new_bbox[3]), color=color, thickness=thickness)
    
    # Calculate position to place text above bbox
    text_height = segment_height + baseline
    y_position = new_bbox[1] - 5  # 5 pixels above the bbox
    
    # Put the text only if y_position is positive
    if y_position - text_height > 0:
        cv2.putText(img, segment_text, (new_bbox[0], y_position), 
                    fontFace=font, 
                    fontScale=font_scale, 
                    color=color, 
                    thickness=thickness)
    else:
        # If there's no space above, put the text inside the bbox
        cv2.putText(img, segment_text, (new_bbox[0], new_bbox[3] - baseline), 
                    fontFace=font, 
                    fontScale=font_scale, 
                    color=color, 
                    thickness=thickness)
    
    return img, new_bbox

def display_image_with_cropped_text(img_path, full_text, segment_text, original_bbox, occurrence=1):
    """
    Load an image, calculate the bounding box for a text segment, and display the image with the new bbox.

    Args:
        img_path: Path to the input image
        full_text: Complete text string in the image
        segment_text: The portion of text to focus on
        original_bbox: Original bounding box coordinates [x1, y1, x2, y2]
        occurrence: The occurrence number of the segment_text in full_text (1-based)
    """
    # Load the image
    img = cv2.imread(img_path)
    if img is None:
        raise FileNotFoundError(f"Image not found at path: {img_path}")
    
    # Apply text fitting with cropped bbox
    result, new_bbox = get_text_segment_bbox(img, full_text, segment_text, original_bbox, occurrence=occurrence)
    
    # Convert to RGB for display
    img_rgb = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
    pil_img = PILImage.fromarray(img_rgb)
    
    # Display the image
    display(pil_img)
    
    return result, new_bbox

# Usage example:

bbox = [962, 99, 1268, 122]
full_text = "AaBbc AaB AaBbCc AaBbccD"
segment_text = "AaBbc"
occurrence = 1  # Change if you want to target a different occurrence

try:
    result, new_bbox = display_image_with_cropped_text(img_path, full_text, segment_text, bbox, occurrence=occurrence)
    print(f"New bounding box coordinates: {new_bbox}")
except Exception as e:
    print(f"Error: {e}")


Error: unsupported operand type(s) for +=: 'int' and 'tuple'