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

def get_text_segment_bbox(img, full_text, segment_text, original_bbox, color=(0, 0, 255), thickness=1, 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
        full_text: Complete text string
        segment_text: The portion of text to focus on
        original_bbox: Original bounding box coordinates [x1, y1, x2, y2]
        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")
    
    # Split the text into words and find the target word
    words = full_text.split()
    target_word_index = words.index(segment_text)
    
    # Calculate the font scale that was used for the original bbox
    original_width = original_bbox[2] - original_bbox[0]
    font_scale = 0.1
    max_scale = 100
    tolerance = 2
    
    while font_scale < max_scale:
        (text_width, text_height), baseline = cv2.getTextSize(full_text, font, font_scale, thickness)
        
        if abs(text_width - original_width) < tolerance:
            break
        elif text_width < original_width:
            font_scale *= 1.1
        else:
            font_scale *= 0.9
    
    # Calculate space width
    (space_width, _), _ = cv2.getTextSize(" ", font, font_scale, thickness)
    
    # Calculate width up to our target word
    prefix_width = 0
    for i in range(target_word_index):
        (word_width, _), _ = cv2.getTextSize(words[i], font, font_scale, thickness)
        prefix_width += word_width + space_width
    
    # Get the width of our segment
    (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]
    
    new_bbox = [int(new_x1), new_y1, int(new_x2), 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] - text_height + baseline - 5
    
    # Put the text
    cv2.putText(img, segment_text, (new_bbox[0], y_position), 
                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):
    # Load the image
    img = cv2.imread(img_path)
    
    # Apply text fitting with cropped bbox
    result, new_bbox = get_text_segment_bbox(img, full_text, segment_text, original_bbox)
    
    # 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 = "AaB"
result, new_bbox = display_image_with_cropped_text(img_path, full_text, segment_text, bbox)
print(f"New bounding box coordinates: {new_bbox}")

space between  AaB and AaBbc also will be counted