import sensor, image, time, math, pyb, lcd
from pyb import UART
#import m_lcd as lcd
sensor.reset()
sensor.set_pixformat(sensor.RGB565)
sensor.set_framesize(sensor.QQVGA2)
lcd.init()
sensor.skip_frames(time = 1000)
sensor.set_auto_gain(False)
sensor.set_auto_whitebal(False)
sensor.set_auto_whitebal(False)
clock = time.clock()
tag_families = 0
tag_families |= image.TAG16H5
tag_families |= image.TAG25H7
tag_families |= image.TAG25H9
tag_families |= image.TAG36H10
tag_families |= image.TAG36H11
tag_families |= image.ARTOOLKIT

def family_name(tag):
    if(tag.family() == image.TAG16H5):
        return "TAG16H5"
    if(tag.family() == image.TAG25H7):
        return "TAG25H7"
    if(tag.family() == image.TAG25H9):
        return "TAG25H9"
    if(tag.family() == image.TAG36H10):
        return "TAG36H10"
    if(tag.family() == image.TAG36H11):
        return "TAG36H11"
    if(tag.family() == image.ARTOOLKIT):
        return "ARTOOLKIT"

def Wait_Ok():
    global uartSendFlag
    if uart.any():
        cmd = uart.readline().decode().strip()
        if(int(cmd) == 'K'):
            uartSendFlag = True

led1 = pyb.LED(1)
led2 = pyb.LED(2)
uart = pyb.UART(3, 9600)



cCount = 0
def workMode1():
    global uartSendFlag
    global aCount
    global bCount
    global cCount
    clock.tick()
    img = sensor.snapshot()
    rect_tuple = (0,160-22,128,22)
    img.draw_rectangle(rect_tuple, color=(255,0,0))
    img.draw_line((64,160-22,64,160), color=(255,0,0))
    img.draw_string(5, 160-20, "id:", color=(0, 255, 0), scale=1.6)
    img.draw_string(65, 160-20, "C:", color=(0, 255, 0), scale=1.6)
    for tag in img.find_apriltags(families=tag_families):
        led1.on()
        img.draw_rectangle(tag.rect(), color = (255, 0, 0))
        img.draw_cross(tag.cx(), tag.cy(), color = (0, 255, 0))
        print_args = (family_name(tag), tag.id(), (180 * tag.rotation()) / math.pi)
    #    print("Tag Family %s, Tag ID %d, rotation %f (degrees)" % print_args)
        id = print_args[1]
        cls = ''
        if id == 1:
            cls = 'A'
            img.draw_string(85, 160-20, str(aCount), color=(0, 255, 0), scale=1.6)
        elif id == 2:
            cls = 'B'
            img.draw_string(85, 160-20, str(bCount), color=(0, 255, 0), scale=1.6)
        elif id == 3:
            cls = 'C'
            img.draw_string(85, 160-20, str(cCount), color=(0, 255, 0), scale=1.6)
        img.draw_string(40, 160-20, cls, color=(0, 255, 0), scale=1.6)

        Wait_Ok()
        if id <= 3 and uartSendFlag == True:
            uart.write(cls)
            uartSendFlag = False
            if cls == 'A':
                aCount += 1
            elif cls == 'B':
                bCount += 1
            elif cls == 'C':
                cCount += 1
        led1.off()
    lcd.display(img)
    print("fps:%f" %clock.fps())


green = (61, 23, -39, -17, 9, -19)
red = (47, 39, 36, 96, 24, 71);
yellow = (0, 63, -16, 4, 32, 79);
greenCount = 0
redCount = 0
yellowCount = 0



def doWork():
   # mode = GetWork_Mode()
    #while True:
     #   if mode == 1:
            workMode1()
      #  elif mode == 2:
      #      workMode2()
       # else:
            doWork()

while(True):
    doWork()