import cv2
from time import sleep
from pistreaming.server import WIDTH

#           (hue, satur, value)
mins_blue = (90, 150, 0)
maxs_blue = (125, 255, 255)

mins_red = (0, 100, 0)
maxs_red = (25, 255, 255)

mins_gray = (110, 50, 50)
maxs_gray = (130, 255, 255)

AREA_THRESHOLD = 20
CENTER_PADDING = 150

def capture2(lock):
    lock.acquire()
    capture = cv2.imread("images/capture.jpg")
    lock.release()
    return capture

def threshold_image(image, mins, maxs):
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    thresholded = cv2.inRange(hsv, mins, maxs)
    return thresholded

def crop_to_target_half(image, direction):
    left = image[:, :WIDTH//2+CENTER_PADDING]
    right = image[:, WIDTH//2-CENTER_PADDING:]
    if direction:
        return left
    else:
        return right

def get_bounding_box(thresholded, original):
    _, contours, _ = cv2.findContours(thresholded, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if len(contours) == 0:
        return None, None
    contour = max(contours, key=cv2.contourArea)
    bounding_box = cv2.boundingRect(contour)
    x, y, w, h = bounding_box
    boxed = cv2.rectangle(original, (x, y), (x+w, y+h), (0, 255, 0), 3)
    return bounding_box, boxed

def get_bounding_box_area(bounding_box):
    _, _, w, h = bounding_box
    return w * h

# True => left, False => right
def get_direction(image, bounding_box):
    col_start, row_start, cols, rows = bounding_box
    row_end = row_start + rows
    col_halfway = col_start + cols // 2
    col_end = col_start + cols
    for row in range(row_start, row_end):
        for col in range(col_start, col_halfway):
            if image[row, col] != 0:
                return True
        for col in range(col_halfway, col_end):
            if image[row, col] != 0:
                return False
    return False

def get_ball_x_center(bounding_box):
    x, _, w, _ = bounding_box
    x_center = x + w//2
    return x_center

def get_arrow_direction(lock):
    capture = capture2(lock)
    cv2.imwrite("images/arrow/capture.jpg", capture)
    thresholded = threshold_image(capture, mins_blue, maxs_blue)
    cv2.imwrite("images/arrow/thresholded.jpg", thresholded)
    bounding_box, boxed = get_bounding_box(thresholded, capture)
    cv2.imwrite("images/arrow/boxed.jpg", boxed)
    direction = get_direction(thresholded, bounding_box)
    return direction

def ball_finder(lock, direction, cropped=True):
    capture = capture2(lock)
    cv2.imwrite("images/ball/capture.jpg", capture)
    if cropped:
        target_half = crop_to_target_half(capture, direction)
        cv2.imwrite("images/ball/target_half.jpg", target_half)
    else:
        target_half = capture
    thresholded = threshold_image(target_half, mins_red, maxs_red)
    cv2.imwrite("images/ball/thresholded.jpg", thresholded)
    bounding_box, boxed = get_bounding_box(thresholded, capture)
    if bounding_box is None:
        return None
    cv2.imwrite("images/ball/boxed.jpg", boxed)
    ball_x_center = get_ball_x_center(bounding_box)
    if not direction and cropped: # Target on right and image is cropped
        ball_x_center += WIDTH//2 - CENTER_PADDING
    return ball_x_center

def ball_is_in_frame(lock):
    capture = capture2(lock)
    cv2.imwrite("images/ball/capture.jpg", capture)
    thresholded = threshold_image(capture, mins_red, maxs_red)
    cv2.imwrite("images/ball/thresholded.jpg", thresholded)
    bounding_box, boxed = get_bounding_box(thresholded, capture)
    if bounding_box is None:
        return False
    cv2.imwrite("images/ball/boxed.jpg", boxed)
    area = get_bounding_box_area(bounding_box)
    return area >= AREA_THRESHOLD