#!/usr/bin/env python

#necessary imports
from __future__ import print_function
import math
import numpy as np
import cv2

#set B value
def set_B_value(value):
    blue_val = value
    cv2.setTrackbarPos('B', canvas_name, blue_val)

def set_G_value(value):
    green_val = value
    cv2.setTrackbarPos('G', canvas_name, green_val)

def set_R_value(value):
    red_val = value
    cv2.setTrackbarPos('R', canvas_name, red_val)


#instructions
print("Welcome to object painting!"
      "\n\nTo exit, press 'q' while the video display window is selected."
      "\n\nThis program works best with a solid color background."
      "\n\nKeep your hand inside the blue bounding box in the bottom left corner."
      "\n\nYou can change ink color using the trackbars above the canvas."
      "\n\nTo draw, hold up two fingers. To erase, hold up an open hand. To move the brush "
      "without interacting with the canvas, hold up a closed fist.")

#hsv colorspace ranges
#ranges found using script downloaded from: https://github.com/opencv/opencv/blob/3.4/samples/cpp/tutorial_code/ImgProc/Threshold_inRange.cpp
#tennis ball
H_range = (20, 60)
S_range = (40, 160)
V_range = (100, 240)
#skin
H_skin = (1, 20)
S_skin = (60, 170)
V_skin = (120, 255)

#canvas window name
canvas_name = 'canvas'
cv2.namedWindow(canvas_name)

#create trackbars to set ink color
cv2.createTrackbar('B', canvas_name, 0, 255, set_B_value)
cv2.createTrackbar('G', canvas_name, 0, 255, set_G_value)
cv2.createTrackbar('R', canvas_name, 0, 255, set_R_value)

#initialize video capture
live_cap = cv2.VideoCapture(0)

#window height & width
img_Height = int(1.5*live_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))        #rows
img_Width = int(1.5*live_cap.get(cv2.CAP_PROP_FRAME_WIDTH))          #columns

#create blank canvas
canvas = 255*np.ones((img_Height, img_Width, 3), dtype=np.uint8)
canvas_in = np.copy(canvas)                 #create canvas to be "painted"
canvas_brush = np.copy(canvas)              #create canvas to show brush

#run as a loop
while True:
    #designate exit command - press q to quit
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

    #read captured image for display
    ret, frame_in = live_cap.read()
    #flip image horizontally
    frame_in = cv2.flip(frame_in, 1)
    #upscale image
    new_dims = (img_Width, img_Height)
    frame_in = cv2.resize(frame_in, new_dims, cv2.INTER_AREA)

    #bounding box for hand
    ROI_hand = np.copy(frame_in[int(img_Height/2):int(img_Height), 0:300, :])

    #place ROI rectangle on brush image
    img_brush = cv2.rectangle(frame_in, (0, img_Height), (300, img_Height/2), (255, 0, 0), 2)

    #convert input from bgr to hsv
    frame_hsv = cv2.cvtColor(frame_in, cv2.COLOR_BGR2HSV)
    ROI_hsv = cv2.cvtColor(ROI_hand, cv2.COLOR_BGR2HSV)

    #thresholding
    frame_threshold = cv2.inRange(frame_hsv, (H_range[0], S_range[0], V_range[0]),
                                  (H_range[1], S_range[1], V_range[1]))
    ROI_threshold = cv2.inRange(ROI_hsv, (H_skin[0], S_skin[0], V_skin[0]),
                                (H_skin[1], S_skin[1], V_skin[1]))

    #gaussian blur
    frame_blur = cv2.GaussianBlur(frame_threshold, (7, 7), 0)
    ROI_blur = cv2.GaussianBlur(ROI_threshold, (7, 7), 0)

    #erosion/dilation to remove noise
    frame_erosion = cv2.erode(frame_blur, np.ones((3, 3)), iterations=2)
    frame_dilation = cv2.dilate(frame_erosion, np.ones((5, 5)), iterations=2)
    ROI_erosion = cv2.erode(ROI_blur, np.ones((3, 3)), iterations=2)
    ROI_dilation = cv2.dilate(ROI_erosion, np.ones((5, 5)), iterations=3)

    #find contours
    _, contours, _ = cv2.findContours(frame_dilation, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    _, ROI_contours, _ = cv2.findContours(ROI_dilation, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)

    #determine number of fingers held up
    num_gaps = 0
    if len(ROI_contours) == 0:          #do nothing if no contours in ROI
        num_fingers = 0
    elif len(ROI_contours) > 0:         #if contours exist
        ROI_big_cont = max(ROI_contours, key=cv2.contourArea)       #find biggest contour (hand)
        convex_hull = cv2.convexHull(ROI_big_cont, returnPoints=False)  #find outline of contour
        convex_defects = cv2.convexityDefects(ROI_big_cont, convex_hull)    #find convex defects in outline
        if isinstance(convex_defects, np.ndarray):          #if defects exist
            for i in range(convex_defects.shape[0]):        #for all defects
                s, e, f, d = convex_defects[i, 0]           #get important defect points
                start = tuple(ROI_big_cont[s][0])               #start
                end = tuple(ROI_big_cont[e][0])                 #end
                far = tuple(ROI_big_cont[f][0])                 #farthest point
                a = math.sqrt((end[0] - start[0])**2 + (end[1] - start[1])**2)  #side 1
                b = math.sqrt((far[0] - start[0])**2 + (far[1] - start[1])**2)  #side 2
                c = math.sqrt((end[0] - far[0])**2 + (end[1] - far[1])**2)      #side 3
                angle = math.acos((b**2 + c**2 - a**2)/(2*b*c)) #find angle at defect using law of cosines
                if angle <= math.pi/2:              #if angle less than 90 degrees
                    num_gaps += 1                   #count defect as gap in fingers
                    cv2.circle(ROI_hand, far, 5, (0, 0, 255), 2)

        if num_gaps == 0:                   #if no gaps detected
            num_fingers = 0                 #no fingers raised
        elif num_gaps > 0:                  #if >0 gaps
            num_fingers = num_gaps + 1      #extra finger raised

    #drawing conditions
    if num_fingers == 0:                #do nothing if no fingers raised
        state = 'hold'
    elif num_fingers == 2:              #draw if two fingers raised
        state = 'draw'
    elif num_fingers == 5:              #erase if open hand
        state = 'erase'
    else:                               #do nothing if other gesture
        state = 'hold'

    #tennis ball
    if len(contours) > 0:                                       #if contours exist
        biggest_contour = max(contours, key=cv2.contourArea)    #find biggest contour
    elif len(contours) == 0:                                    #if no contours
        cv2.imshow('brush spot', img_brush)
        continue

    #get ink color
    B_val = cv2.getTrackbarPos('B', canvas_name)
    G_val = cv2.getTrackbarPos('G', canvas_name)
    R_val = cv2.getTrackbarPos('R', canvas_name)

    if cv2.contourArea(biggest_contour) >= 4000:                        #if ball present
        #find smallest circle enclosing biggest contour
        (x, y), radius = cv2.minEnclosingCircle(biggest_contour)        #smallest enclosing circle
        centerX = int(x)                #brush center
        centerY = int(y)
        img_brush = cv2.circle(img_brush, (centerX, centerY),
                               3, (0, 0, 255), 2)                       #draw brush on captured image
        canvas_brush = cv2.circle(canvas_brush, (centerX, centerY),
                                  6, (0, 0, 255), 2)                    #draw brush on canvas
        if state == 'draw':                                             #if draw state
            canvas_in = cv2.circle(canvas_in, (centerX, centerY),
                                   2, (B_val, G_val, R_val), -1)        #paint on canvas
            canvas_brush = cv2.circle(canvas_brush, (centerX, centerY),
                                      2, (B_val, G_val, R_val), -1)     #paint on brush-canvas
        elif state == 'erase':                                          #if erase state
            canvas_in = cv2.circle(canvas_in, (centerX, centerY),
                                   2, (255, 255, 255), -1)              #turn canvas white
            canvas_brush = cv2.circle(canvas_brush, (centerX, centerY),
                                      2, (255, 255, 255), -1)                 #turn brush-canvas white
        elif state == 'hold':                                           #if hold state
            pass

    elif cv2.contourArea(biggest_contour) < 4000:                       #if ball not present
        img_brush = img_brush                                           #brush image is captured image
        canvas_brush = np.copy(canvas_in)                               #canvas does not show brush

    cv2.imshow('brush spot', img_brush)                                 #show video with brush
    cv2.imshow(canvas_name, canvas_brush)                               #show canvas with brush

    canvas_brush = np.copy(canvas_in)                                   #erase brush marker

#end video capture and close windows
live_cap.release()
cv2.destroyAllWindows()
