'''
This is a demo to combine OpenMV with Rotrics DexArm to recognize objects' color and sort color blocks. 

The color threshold depends on the environment light and object color. You need to test and adjust it according to your project. 

Camera: OpenMv4 H7
Robot Arm: Rotrics DexArm V1.0

Pin Connection:
	USB Type-C to DexArm
	DexArm's TX to OpenMv's Pin 5(RX)
	DexArm's RX to OpenMv's Pin 4(TX)

Date: 20/05/2020
Author: Baochun Feng
Contact: fengbaochun@rotrics.com
'''


import sensor, image, time
from pyb import UART

clock = time.clock()

#Serial setting
#uart_log = UART(3, 9600)
uart_log = UART(3, 115200)

#Output switch
#PRINT_SW = False   #control robot arm
PRINT_SW = True     #Terminal display

#camera recognize area [X1, Y1, X2, Y2]
img_roi = [10,10,310,310]

#Green threshold
green_td = [(60, 100, -126, -29, -93, 91)]
green_list = []

#Red threshold
red_td=[(45, 94, 21, 104, -1, 125)]
red_list = []

#Blue threshold
blue_td=[(54, 95, -128, 7, -19, -128)]
blue_list = []

#Robot arm origin position

#ARM_Origin_x = 330
#ARM_Origin_y = 74

ARM_Origin_x = 286
ARM_Origin_y = 45

#Physical Object size/mm
obj_fact_size = 28

#Image Object size/mm
obj_img_size = 64

#Calibrate value(Physical object size / Imgage object size)
multiple = obj_img_size / obj_fact_size

#Gcode x y list(default is empty)
Gcode_XY_list = []

#Gcode Object height
Gcode_Z = -35

#Target position for placing
move_pos_target = [(100,-210),(182,-210),(200,130)]


'''
Function： Calculate ARM XY coordinate value
Input： The coordinate value of obejct in the image
Return： ARM XY coordinate value
Notice：Please make sure the IMAGE coordinate is the same as the ARM coordinate. If not, fix it according to the formula.
		When the IMAGE coordinate is the same as the ARM coordinate: 

        Gcode_x = ARM_Origin_x - obj_img_x/multiple
        Gcode_y = ARM_Origin_y - obj_img_y/multiple

        When the IMAGE coordinate is NOT the same as the ARM coordinate:：
        Fix it with the formula below:
        Gcode_x = ARM_Origin_x - obj_img_y/multiple
        Gcode_y = ARM_Origin_y - obj_img_x/multiple
'''
def get_ARM_pos(obj_img_x,obj_img_y):
    Gcode_x = ARM_Origin_x - obj_img_y/multiple
    Gcode_y = ARM_Origin_y - obj_img_x/multiple

    return [Gcode_x,Gcode_y]
    pass


def chose_print( temp_str ):

    #read_flag = False
    read_flag = True

    #  print(temp_str)
    uart_log.write(temp_str)

    while (read_flag):
        read_data = uart_log.read()
        if read_data:
            read_flag = False
        else:
            read_flag = True

    pass


'''
Connect X Y Z value into a G-code and send it to the ARM via Serial port
'''
def send_gcode( x , y ,z ):

    Gcode_str = 'G0 X'+ str(int(x)) + ' Y' +str(int(y)) + ' Z' + str(z) +'\r\n'
    chose_print(Gcode_str)

    pass

def send_gcode_Z( z ):

    Gcode_str = 'G0' + ' Z' + str(z) +'\r\n'
    chose_print(Gcode_str)

    pass

def send_gcode_XY( x , y ):

    Gcode_str = 'G0 X'+ str(int(x)) + ' Y' +str(int(y)) +'\r\n'
    chose_print(Gcode_str)

    pass


def init_config():

    sensor.reset()
    sensor.set_pixformat(sensor.RGB565)
    sensor.set_framesize(sensor.VGA)
    sensor.set_windowing((420, 420))
    sensor.skip_frames(10)
    sensor.set_auto_whitebal(False)

    pass


#Recognize color block and draw it
def find_color_block(color_td,img_data,temp_list):
    #Settings to recognize color block 
    block = img_data.find_blobs(color_td,roi=img_roi,x_stride=20,y_stride=20,pixels_threshold=2200)
    #temp_list = []
    flag = True
    num = 0
    if block and flag:
        for bl in block:
            if bl.density()>0.4:                #Decide color block density 

                img_data.draw_rectangle(bl[0:4])  #Draw a rectangle according the data of tuple b[0] b[1] b[2] b[3]
                print(bl[2],bl[3])
                img_data.draw_cross(bl[5], bl[6])  #Draw a cross

                c=img_data.get_pixel(bl[5], bl[6])     #Return the RGB888 pixel tuple of (x,y)

                temp_list.append((bl[5],bl[6]))         #add coordinate value to the list

                num = num+1                             #Count color blocks
                flag = False

    return num      #Return Object Number
    pass

'''

Pick and place object
obj_num : Object number
obj_pos_list : Object coordinate list

'''
def move_obj(obj_num,obj_pos_list,target_x,target_y):

    index = 0
    Z_val = 30 #Object height
    print("*********************\n")
    while (index < obj_num) :

        Gcode_XY_list = get_ARM_pos(obj_pos_list[index][0],obj_pos_list[index][1])        #Map IMAGE coordinate to ARM coordinate.
        #print(Gcode_XY_list)

        send_gcode( Gcode_XY_list[0] , Gcode_XY_list[1] , -20 ) #Move to target position
        #Move Down
        send_gcode_Z( Gcode_Z )
        #Pick
        chose_print("M1000\r\n")
        #Move Up
        send_gcode_Z( Gcode_Z + 30)
        #Move to 120mm above the object
        send_gcode( target_x, target_y , Gcode_Z + 120 )

        #Move Down
        send_gcode_Z( Gcode_Z + 2 + Z_val*index)
        #Place
        chose_print("M1002\r\n")
        #Move up a little bit before next movement.
        send_gcode_Z( Gcode_Z + 2 + Z_val*index + 10)

        print("--------------------\n")

        index = index + 1

    print("*********************\n")

    pass



#main function
def main():
    uart_log.write("M1111\r\n")
    time.sleep(1)
    uart_log.write("G0X180Y120Z0F3000 \r\n")

    flag = True

    while flag :

        img=sensor.snapshot()                   #Take a picture
        img.draw_string(0, 0, "(0,0)",size=10)  #Draw coordinate
        img.draw_rectangle(img_roi[0:4])        #Set recognize area


        #Blue recognize and picking
        num_b = find_color_block(blue_td,img,blue_list)
        print(num_b,blue_list)   #Print object number and coordinate value
        move_obj(num_b,blue_list,move_pos_target[0][0],move_pos_target[0][1])
        blue_list.clear()      #Clear buffer

        #Red recognize and picking
        num_r = find_color_block(red_td,img,red_list)
        print(num_r,red_list)   #Print object number and coordinate value
        move_obj(num_r,red_list,move_pos_target[2][0],move_pos_target[2][1])
        red_list.clear()      #Clear buffer

        #Green recognize and picking
        num_g = find_color_block(green_td,img,green_list)
        print(num_g,green_list)   #Print object number and coordinate value
        move_obj(num_g,green_list,move_pos_target[1][0],move_pos_target[1][1])
        green_list.clear()      #Clear buffer


    pass



#Main loop
if __name__ == '__main__':
    init_config()
    main()
