import cv2
import numpy as np


def nothing(a):
    pass


def mask_red(frame_hsv):
    red_mask_orange = cv2.inRange(frame_hsv, np.array([0, 85, 110]), np.array([15, 255, 255]))
    red_mask_violet = cv2.inRange(frame_hsv, np.array([165, 85, 110]), np.array([180, 255, 255]))
    return red_mask_orange + red_mask_violet


def mask_green(frame_hsv):
    return cv2.inRange(frame_hsv, np.array([40, 85, 110]), np.array([91, 255, 255]))


wnd = 'frame'
cap = cv2.VideoCapture(0)
cv2.namedWindow(wnd)
cv2.createTrackbar('Hul', wnd, 0, 180, nothing)
cv2.createTrackbar('Sal', wnd, 0, 255, nothing)
cv2.createTrackbar('Val', wnd, 0, 255, nothing)

cv2.createTrackbar('Huh', wnd, 0, 180, nothing)
cv2.createTrackbar('Sah', wnd, 0, 255, nothing)
cv2.createTrackbar('Vah', wnd, 0, 255, nothing)

while (cap.isOpened()):
    _, frame = cap.read()
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    hsv = cv2.GaussianBlur(hsv, (5, 5), 0)

    hul = cv2.getTrackbarPos('Hul', wnd)
    huh = cv2.getTrackbarPos('Huh', wnd)
    sal = cv2.getTrackbarPos('Sal', wnd)
    sah = cv2.getTrackbarPos('Sah', wnd)
    val = cv2.getTrackbarPos('Val', wnd)
    vah = cv2.getTrackbarPos('Vah', wnd)

    mask = cv2.inRange(hsv, np.array([hul, sal, val]), np.array([huh, sah, vah]))
    green_mask = mask_green(hsv)
    red_mask = mask_red(hsv)
    cv2.imshow('green_mask', green_mask)
    cv2.imshow('red_mask', red_mask)
    green_circles = cv2.HoughCircles(green_mask, cv2.HOUGH_GRADIENT, 1, hsv.shape[0]/8, param1=100, param2=30, minRadius=10, maxRadius=0)
    red_circles = cv2.HoughCircles(red_mask, cv2.HOUGH_GRADIENT, 1, hsv.shape[0]/8, param1=100, param2=30, minRadius=10, maxRadius=0)
    if green_circles is not None:
        print('green', green_circles)
        for circle in green_circles:
            cv2.circle(frame, (circle[0][0], circle[0][1]), circle[0][2], (0, 255, 0), 0)
    if red_circles is not None:
        print('red', red_circles)
        for circle in red_circles:
            cv2.circle(frame, (circle[0][0], circle[0][1]), circle[0][2], (0, 0, 255), 0)
    cv2.imshow(wnd, frame)
    cv2.imshow('mask', mask)
    if cv2.waitKey(1) == ord('q'):
        cv2.destroyAllWindows()
        break