import cv2
import numpy as np
import math

# Load the image
image_path = "output_folder/20241027_092609_tp00004/frame_0000.jpg"
original_image = cv2.imread(image_path)

if original_image is None:
    print("Error: Image not found or unable to load.")
    exit()

# Constants for the toolbar
toolbar_height = 60
button_width = 120
button_height = 50
button_margin = 10

# Scale related variables
scale_mode = False
ref_pairs = []  # List to store tuples of (point1, point2, real_distance)
scale_factors = []  # List to store scale factors from each reference pair
final_scale = None  # Final scale factor (real-world units per pixel)

# Homography related variables
homography_matrix = None
ref_points_image = []
ref_points_real = []  # User should input real-world coordinates for these points
reference_mode = False

# Define button properties: (label, color)
buttons = [
    {"label": "Add", "color": (200, 200, 200)},
    {"label": "Delete", "color": (200, 200, 200)},
    {"label": "Move", "color": (200, 200, 200)},
    {"label": "Line", "color": (200, 200, 200)},
    {"label": "Remove Line", "color": (200, 200, 200)},
    {"label": "Set Scale", "color": (200, 200, 200)},  # New Button
    {"label": "Clear", "color": (200, 200, 200)},
    {"label": "Quit", "color": (200, 200, 200)}
]

# Calculate window size to include toolbar
toolbar = np.ones((toolbar_height, original_image.shape[1], 3), dtype=np.uint8) * 50  # Dark gray toolbar
window_image = np.vstack((toolbar, original_image.copy()))

# Initialize variables
points = []
lines = []  # List to store lines with their endpoints and distance
selected_point = -1  # Index of the selected point, if any
mode = 'add'  # Default mode; can be 'add', 'delete', 'move', 'line', 'remove_line'
point_radius = 7  # Increased radius for better hit detection
line_mode_selected_points = []  # To store two points for line drawing
current_mouse_position = (-1, -1)  # To store current mouse position

# Function to draw buttons on the toolbar
def draw_toolbar(img):
    for idx, button in enumerate(buttons):
        x_start = button_margin + idx * (button_width + button_margin)
        y_start = 5
        # Draw button rectangle
        cv2.rectangle(img, (x_start, y_start),
                      (x_start + button_width, y_start + button_height),
                      button["color"], -1)
        # Highlight the selected mode button
        if button["label"].lower().replace(" ", "_") == mode:
            cv2.rectangle(img, (x_start, y_start),
                          (x_start + button_width, y_start + button_height),
                          (0, 255, 0), 3)
        # Put button label
        text_size = cv2.getTextSize(button["label"], cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]
        text_x = x_start + (button_width - text_size[0]) // 2
        text_y = y_start + (button_height + text_size[1]) // 2
        cv2.putText(img, button["label"], (text_x, text_y),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 2)
    
    # If in scale mode, display instructions
    if scale_mode:
        num_pairs = len(ref_pairs)
        if num_pairs < 2:
            cv2.putText(img, f"Select pair {num_pairs + 1} of 2: Click first point.",
                        (10, toolbar_height - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 255), 2)
        else:
            cv2.putText(img, "Scale set. You can add more reference pairs or proceed.",
                        (10, toolbar_height - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 255), 2)
            

# Function to calculate distance between two points
def calculate_distance(pt1, pt2):
    return math.hypot(pt2[0] - pt1[0], pt2[1] - pt1[1])

# Function to check if a click is near a line
def is_near_line(click_pt, line_pt1, line_pt2, tolerance=10):
    # Calculate distance from click_pt to the line
    x0, y0 = click_pt
    x1, y1 = line_pt1
    x2, y2 = line_pt2
    numerator = abs((y2 - y1)*x0 - (x2 - x1)*y0 + x2*y1 - y2*x1)
    denominator = math.hypot(y2 - y1, x2 - x1)
    if denominator == 0:
        return False
    distance = numerator / denominator
    return distance <= tolerance

# Function to redraw points, lines, and toolbar on the image
def redraw_image():
    display_img = window_image.copy()
    draw_toolbar(display_img)

    # Display current mouse position
    if current_mouse_position != (-1, -1):
        cv2.putText(display_img, f"Mouse: {current_mouse_position}", (10, toolbar_height - 25),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1)

    # Draw all saved lines
    for line in lines:
        pt1, pt2, distance = line['points'][0], line['points'][1], line['distance']
        adjusted_pt1 = (pt1[0], pt1[1] + toolbar_height)
        adjusted_pt2 = (pt2[0], pt2[1] + toolbar_height)
        cv2.line(display_img, adjusted_pt1, adjusted_pt2, (255, 0, 0), 2)

        # Calculate midpoint for distance text
        mid_point = ((pt1[0] + pt2[0]) // 2, (pt1[1] + pt2[1]) // 2)
        adjusted_mid_point = (mid_point[0], mid_point[1] + toolbar_height)
        if final_scale is not None:
            cv2.putText(display_img, f"{distance:.2f} units", adjusted_mid_point,
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 255), 2)
        else:
            cv2.putText(display_img, f"{distance:.2f}px", adjusted_mid_point,
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 255), 2)

    # If in line mode and two points are selected, draw temporary line
    if mode == 'line' and len(line_mode_selected_points) == 2:
        temp_pt1 = (line_mode_selected_points[0][0],
                    line_mode_selected_points[0][1] + toolbar_height)
        temp_pt2 = (line_mode_selected_points[1][0],
                    line_mode_selected_points[1][1] + toolbar_height)
        cv2.line(display_img, temp_pt1, temp_pt2, (0, 0, 255), 2)

    # Draw points
    for idx, point in enumerate(points):
        adjusted_point = (point[0], point[1] + toolbar_height)
        if idx == selected_point:
            cv2.circle(display_img, adjusted_point, point_radius + 3, (0, 0, 255), -1)  # Highlight selected point
        cv2.circle(display_img, adjusted_point, point_radius, (0, 255, 0), -1)

    # Draw reference pairs if any
    for i, pair in enumerate(ref_pairs):
        pt1, pt2 = pair['points']
        adjusted_pt1 = (pt1[0], pt1[1] + toolbar_height)
        adjusted_pt2 = (pt2[0], pt2[1] + toolbar_height)
        cv2.line(display_img, adjusted_pt1, adjusted_pt2, (0, 0, 255), 2)
        cv2.putText(display_img, f"P{i+1}", adjusted_pt1,
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
        cv2.putText(display_img, f"P{i+1}", adjusted_pt2,
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)

    cv2.imshow("Image", display_img)

# Mouse callback function
def handle_mouse(event, x, y, flags, param):
    global selected_point, points, mode, line_mode_selected_points
    global current_mouse_position, lines
    global scale_mode, ref_pairs, scale_factors, final_scale

    # Update current mouse position
    if y > toolbar_height:
        img_x, img_y = x, y - toolbar_height
        current_mouse_position = (img_x, img_y)
    else:
        current_mouse_position = (-1, -1)

    if y <= toolbar_height:
        # Click is on the toolbar
        if event == cv2.EVENT_LBUTTONDOWN:
            for idx, button in enumerate(buttons):
                x_start = button_margin + idx * (button_width + button_margin)
                y_start = 5
                x_end = x_start + button_width
                y_end = y_start + button_height
                if x_start <= x <= x_end and y_start <= y <= y_end:
                    label = button["label"].lower().replace(" ", "_")
                    if label == 'add':
                        mode = 'add'
                        scale_mode = False
                    elif label == 'delete':
                        mode = 'delete'
                        scale_mode = False
                    elif label == 'move':
                        mode = 'move'
                        scale_mode = False
                    elif label == 'line':
                        mode = 'line'
                        line_mode_selected_points = []
                        scale_mode = False
                    elif label == 'remove_line':
                        mode = 'remove_line'
                        scale_mode = False
                    elif label == 'set_scale':
                        mode = 'set_scale'
                        scale_mode = True
                        ref_pairs = []
                        scale_factors = []
                        print("Scale mode activated. Please select pairs of points with known distances.")
                    elif label == 'clear':
                        points.clear()
                        lines.clear()
                        line_mode_selected_points = []
                        selected_point = -1
                        scale_factors = []
                        final_scale = None
                        ref_pairs = []
                        print("All points, lines, and scale references cleared.")
                        scale_mode = False
                    elif label == 'quit':
                        cv2.destroyAllWindows()
                        exit()
                    print(f"Mode: {mode.replace('_', ' ').capitalize()}")
                    redraw_image()
                    return  # Exit after handling button click
    else:
        # Click is on the image
        if img_x < 0 or img_y < 0 or img_y >= original_image.shape[0] or img_x >= original_image.shape[1]:
            return  # Click outside the image area

        if scale_mode:
            # Handle scale setting
            if event == cv2.EVENT_LBUTTONDOWN:
                if len(ref_pairs) < 2:  # Allow up to 2 reference pairs for better accuracy
                    if len(ref_pairs) == 0 or len(ref_pairs[-1]) == 2:
                        # Select first point of a new pair
                        ref_pairs.append({'points': [ (img_x, img_y) ]})
                        print(f"Reference Pair {len(ref_pairs)}: Click the second point for this pair.")
                    elif len(ref_pairs[-1]['points']) == 1:
                        # Select second point of the current pair
                        ref_pairs[-1]['points'].append((img_x, img_y))
                        try:
                            real_distance = float(input(f"Enter the real-world distance for reference pair {len(ref_pairs)} (in desired units, e.g., meters): "))
                            ref_pairs[-1]['distance'] = real_distance
                            # Calculate pixel distance
                            pt1, pt2 = ref_pairs[-1]['points']
                            pixel_distance = calculate_distance(pt1, pt2)
                            if pixel_distance == 0:
                                print("Selected the same point twice. Please reselect the reference pair.")
                                ref_pairs.pop()
                            else:
                                scale = real_distance / pixel_distance
                                scale_factors.append(scale)
                                print(f"Reference Pair {len(ref_pairs)}: Scale factor {scale:.6f} units/pixel added.")
                        except ValueError:
                            print("Invalid input. Please enter a numerical value for the distance.")
                            ref_pairs.pop()
                
                    # Check if enough pairs to compute scale
                    if len(ref_pairs) >= 2:
                        # Compute average scale factor
                        final_scale = sum(scale_factors) / len(scale_factors)
                        print(f"Final scale factor: {final_scale:.6f} units/pixel")
                        scale_mode = False
                else:
                    print("Maximum number of reference pairs reached.")
                redraw_image()
            return  # Exit after handling scale mode click

        # Proceed with regular modes
        if mode == 'add' and event == cv2.EVENT_LBUTTONDOWN:
            # Add a new point
            points.append((img_x, img_y))
            print(f"Point added at: ({img_x}, {img_y})")
            redraw_image()

        elif mode == 'delete' and event == cv2.EVENT_LBUTTONDOWN:
            # Delete point if close to it
            for i, point in enumerate(points):
                distance_sq = (img_x - point[0]) ** 2 + (img_y - point[1]) ** 2
                if distance_sq <= (point_radius * 2) ** 2:  # Increased area for better hit detection
                    print(f"Point deleted at: {points[i]}")
                    points.pop(i)
                    # Also remove any lines connected to this point
                    lines[:] = [line for line in lines if line['points'][0] != point and line['points'][1] != point]
                    break
            redraw_image()

        elif mode == 'move':
            if event == cv2.EVENT_LBUTTONDOWN:
                # Select the point if close to it
                for i, point in enumerate(points):
                    distance_sq = (img_x - point[0]) ** 2 + (img_y - point[1]) ** 2
                    if distance_sq <= (point_radius * 2) ** 2:
                        selected_point = i
                        print(f"Selected point {i + 1} for moving.")
                        break
            elif event == cv2.EVENT_MOUSEMOVE and selected_point != -1 and flags & cv2.EVENT_FLAG_LBUTTON:
                # Update the position of the selected point
                points[selected_point] = (img_x, img_y)
                redraw_image()
            elif event == cv2.EVENT_LBUTTONUP:
                # Release the point
                if selected_point != -1:
                    print(f"Point {selected_point + 1} moved to: {points[selected_point]}")
                selected_point = -1
                redraw_image()

        elif mode == 'line':
            if event == cv2.EVENT_LBUTTONDOWN:
                # Select two points to draw a line
                clicked_point = None
                for point in points:
                    distance_sq = (img_x - point[0]) ** 2 + (img_y - point[1]) ** 2
                    if distance_sq <= (point_radius * 2) ** 2:
                        clicked_point = point
                        break
                if clicked_point:
                    # Avoid selecting the same point twice
                    if clicked_point not in line_mode_selected_points:
                        line_mode_selected_points.append(clicked_point)
                        print(f"Point selected for line: {clicked_point}")
                    else:
                        print("Same point selected. Please select a different point.")
                    if len(line_mode_selected_points) == 2:
                        pt1, pt2 = line_mode_selected_points
                        pixel_distance = calculate_distance(pt1, pt2)
                        if final_scale is not None:
                            real_distance = pixel_distance * final_scale
                            lines.append({'points': (pt1, pt2), 'distance': real_distance})
                            print(f"Line drawn between {pt1} and {pt2} with real-world distance {real_distance:.2f} units")
                        else:
                            # Fallback to pixel distance if scale not set
                            lines.append({'points': (pt1, pt2), 'distance': pixel_distance})
                            print(f"Line drawn between {pt1} and {pt2} with pixel distance {pixel_distance:.2f}px")
                        line_mode_selected_points = []  # Reset after drawing
                        redraw_image()
                else:
                    print("No point near the click to select for line.")

        elif mode == 'remove_line' and event == cv2.EVENT_LBUTTONDOWN:
            # Attempt to remove a line near the click
            for idx, line in enumerate(lines):
                pt1, pt2 = line['points']
                if is_near_line((img_x, img_y), pt1, pt2):
                    if final_scale is not None:
                        print(f"Line removed between {pt1} and {pt2} with distance {line['distance']:.2f} units")
                    else:
                        print(f"Line removed between {pt1} and {pt2} with distance {line['distance']:.2f}px")
                    lines.pop(idx)
                    redraw_image()
                    break
            else:
                print("No line found near the click to remove.")
print("Interactive Point and Line Editor with OpenCV")
print("-------------------------------------------------")
print("Use the toolbar buttons to switch modes:")
print("Add           - Add a new point")
print("Delete        - Delete an existing point")
print("Move          - Move an existing point")
print("Line          - Draw a line between two points and show distance")
print("Remove Line   - Remove a line by clicking near it")
print("Set Scale     - Define reference pairs with known real-world distances")
print("Clear         - Clear all points, lines, and scale references")
print("Quit          - Exit the application")
print("-------------------------------------------------")
print("\n--- Scale Setting Instructions ---")
print("1. Click the 'Set Scale' button.")
print("2. Select two points by clicking on the image to form the first reference pair.")
print("3. After selecting the first point, click the second point of the pair.")
print("4. Enter the known real-world distance between these two points when prompted.")
print("5. Repeat steps 2-4 to add a second reference pair for improved accuracy.")
print("6. Once scale is set, you can draw lines and view real-world distances.")
print("-----------------------------------")

# Initial drawing
redraw_image()

# Set up OpenCV window and callback
cv2.namedWindow("Image")
cv2.setMouseCallback("Image", handle_mouse)

# Main loop
while True:
    key = cv2.waitKey(1) & 0xFF
    # Optionally, you can handle keyboard shortcuts here
    if key == ord('q'):
        break

# Close OpenCV windows
cv2.destroyAllWindows()


# After the main loop
print("\nFinal Selected Points (x, y):")
for i, point in enumerate(points):
    print(f"Point {i + 1}: {point}")

print("\nFinal Drawn Lines:")
for i, line in enumerate(lines):
    pt1, pt2, distance = line['points'][0], line['points'][1], line['distance']
    if final_scale is not None:
        print(f"Line {i + 1}: ({pt1[0]}, {pt1[1]}) to ({pt2[0]}, {pt2[1]}) - Distance: {distance:.2f} units")
    else:
        print(f"Line {i + 1}: ({pt1[0]}, {pt1[1]}) to ({pt2[0]}, {pt2[1]}) - Distance: {distance:.2f}px")

if final_scale is not None:
    print(f"\nScale Factor: {final_scale:.6f} units/pixel")

Mode: Set scale
Reference Pair 1: Click the second point for this pair.
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[2], line 264
    262         else:
    263             print("Maximum number of reference pairs reached.")
--> 264         redraw_image()
    265     return  # Exit after handling scale mode click
    267 # Proceed with regular modes

Cell In[2], line 151
    149 # Draw reference pairs if any
    150 for i, pair in enumerate(ref_pairs):
--> 151     pt1, pt2 = pair['points']
    152     adjusted_pt1 = (pt1[0], pt1[1] + toolbar_height)
    153     adjusted_pt2 = (pt2[0], pt2[1] + toolbar_height)

ValueError: not enough values to unpack (expected 2, got 1)
Reference Pair 1: Scale factor 0.560525 units/pixel added.
Reference Pair 2: Click the second point for this pair.
Final scale factor: 0.560525 units/pixel
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[2], line 264
    262         else:
    263             print("Maximum number of reference pairs reached.")
--> 264         redraw_image()
    265     return  # Exit after handling scale mode click
    267 # Proceed with regular modes

Cell In[2], line 151
    149 # Draw reference pairs if any
    150 for i, pair in enumerate(ref_pairs):
--> 151     pt1, pt2 = pair['points']
    152     adjusted_pt1 = (pt1[0], pt1[1] + toolbar_height)
    153     adjusted_pt2 = (pt2[0], pt2[1] + toolbar_height)

ValueError: not enough values to unpack (expected 2, got 1)

resolve the error