import cv2
 
# You definitely wanna to edit these
INPUT_FILE = '/path/to/screencap.mp4'
OUTPUT_FILE = '/path/to/pedal-data.csv'
# If something's not working, use this as
# the range of frames to see on screen for debugging
DEBUG_FRAME_RANGE = [-1, -1]
# Edit colors for debug views.
# Note colors are (b, g, r) for some reason
INPUT_EXPECTED_BOUNDS_COLOR = (255, 0, 0)
TEXT_COLOR = (0, 255, 0)
CONTOUR_COLOR = (0,255,0)
MEASURED_INPUT_BOUNDS_COLOR = (0, 0, 255)
# Misc. If you're editing these then you already
# know what you're doing
MAX_PEDAL_INPUT_VALUE_PIXELS = 83
PEDAL_INPUT_BAR_WIDTH = 10
PEDAL_INPUT_VERTICAL_BOUNDS = [867, 953]
THROTTLE_INPUT_HORIZ_BOUNDS = [1120, 1140]
BRAKE_INPUT_HORIZ_BOUNDS = [782, 797]
DEBUG = False
 
def extract_pedal_inputs(frame):
    # Init values
    brake_pct = 0
    throttle_pct = 0
    # Convert frame to grayscale
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    # Threshold the image to get any white parts of the screen
    _, thresh = cv2.threshold(gray, 220, 255, cv2.THRESH_BINARY)
    # Find contours in the thresholded image. These are essentially outlines of white parts of the screen
    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    throttle_found = False
    brake_found = False
    if DEBUG:
        # Draw expected input bounds
        cv2.rectangle(frame, (BRAKE_INPUT_HORIZ_BOUNDS[0], PEDAL_INPUT_VERTICAL_BOUNDS[0]), (BRAKE_INPUT_HORIZ_BOUNDS[0] + (BRAKE_INPUT_HORIZ_BOUNDS[1] - BRAKE_INPUT_HORIZ_BOUNDS[0]), PEDAL_INPUT_VERTICAL_BOUNDS[0] + (PEDAL_INPUT_VERTICAL_BOUNDS[1] - PEDAL_INPUT_VERTICAL_BOUNDS[0])), INPUT_EXPECTED_BOUNDS_COLOR, 2)
        cv2.rectangle(frame, (THROTTLE_INPUT_HORIZ_BOUNDS[0], PEDAL_INPUT_VERTICAL_BOUNDS[0]), (THROTTLE_INPUT_HORIZ_BOUNDS[0] + (THROTTLE_INPUT_HORIZ_BOUNDS[1] - THROTTLE_INPUT_HORIZ_BOUNDS[0]), PEDAL_INPUT_VERTICAL_BOUNDS[0] + (PEDAL_INPUT_VERTICAL_BOUNDS[1] - PEDAL_INPUT_VERTICAL_BOUNDS[0])), INPUT_EXPECTED_BOUNDS_COLOR, 2)
    for contour in contours:
        # Get the bounding box of the contour
        x, y, w, h = cv2.boundingRect(contour)
 
        is_input_bar = w > PEDAL_INPUT_BAR_WIDTH and (y > PEDAL_INPUT_VERTICAL_BOUNDS[0] and y < PEDAL_INPUT_VERTICAL_BOUNDS[1])
        is_throttle_bar = x > THROTTLE_INPUT_HORIZ_BOUNDS[0] and x < THROTTLE_INPUT_HORIZ_BOUNDS[1]
        is_brake_bar = x > BRAKE_INPUT_HORIZ_BOUNDS[0] and x < BRAKE_INPUT_HORIZ_BOUNDS[1]
 
        if is_input_bar:
            input_pct = int((h / MAX_PEDAL_INPUT_VALUE_PIXELS) * 100)
            if is_throttle_bar:
                throttle_pct = input_pct
                throttle_found = True
            elif is_brake_bar:
                brake_pct = input_pct
                brake_found = True
            else:
                # Keep searching for the input bars
                continue
            if DEBUG:
                # Show the computed contours
                cv2.drawContours(frame, [contour], -1, CONTOUR_COLOR, 3)
                # Draw a bounding box around the input bar
                cv2.rectangle(frame, (x, y), (x + w, y + h), MEASURED_INPUT_BOUNDS_COLOR, 2)
                if is_throttle_bar:
                    input_name = "Throttle"
                else:
                    input_name = "Brake"
                cv2.putText(frame, f'{input_name}: {input_pct}%', (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, TEXT_COLOR, 2)
            if throttle_found and brake_found:
                break
    # Display the frame with overlays
    if DEBUG:
        # cv2.drawContours(frame, contours, -1, (0,255,0), 3) # Draw all contours
        # Show image with overlays on top
        cv2.imshow('Throttle and Brake Percentages', frame)
        cv2.waitKey(0) # Wait for any key before closing window
        cv2.destroyAllWindows()
    return brake_pct, throttle_pct
 
# For testing just one image frame at a time
# if __name__ == "__main__":
#     image_path = '/path/to/bumper-view.png'
#     frame = cv2.imread(image_path)
#     brake_pct, throttle_pct = extract_pedal_inputs(frame)
#     print(brake_pct, throttle_pct)
 
if __name__ == "__main__":
    vid_capture = cv2.VideoCapture(INPUT_FILE)
    if not vid_capture.isOpened():
        print("Error: Couldn't open video")
        exit(1)
    with open(OUTPUT_FILE, 'w') as inputs_csv:
        inputs_csv.write('Frame,Brake,Throttle\n')
        frame_index = 0
        while True:
            if frame_index >= DEBUG_FRAME_RANGE[0] and frame_index <= DEBUG_FRAME_RANGE[1]:
                DEBUG = True
            else:
                DEBUG = False
                if DEBUG_FRAME_RANGE[0] > 0 and frame_index == DEBUG_FRAME_RANGE[0] - 1:
                    print("Beginning debug of frames. Take a close look")
                if DEBUG_FRAME_RANGE[1] > 0 and frame_index == DEBUG_FRAME_RANGE[1] + 1:
                    print("Debugging frames complete. Finishing write to file. This could take a while")
            ret, frame = vid_capture.read()
            if not ret:
                print("Error: Couldn't retrieve frame")
                break
            brake_pct, throttle_pct = extract_pedal_inputs(frame)
            # Note the + 1 below
            inputs_csv.write(f'{frame_index + 1},{brake_pct},{throttle_pct}\n')
            frame_index += 1
    vid_capture.release()