import sensor, image, time
from pyb import UART
from image import SEARCH_EX

usart3 = UART(3, 115200)
usart3.init(115200, bits=8, parity=None, stop=1)

sensor.reset()
sensor.set_pixformat(sensor.RGB565)
sensor.set_framesize(sensor.QQVGA)
sensor.set_vflip(True)
sensor.set_hmirror(True)
sensor.skip_frames(time=2000)
print("硬件初始化完毕")

# 引入模板图片的路径
templates1 = ["/01.1.pgm", "/01.2.pgm", "/01.3.pgm", "/01.4.pgm", "/01.5.pgm"]
templates2 = ["/02.1.pgm", "/02.2.pgm", "/02.3.pgm", "/02.4.pgm", "/02.5.pgm"]

# 将图片路径列表置换为图片列表
for i in range(len(templates1)):
    templates1[i] = image.Image(templates1[i])
for i in range(len(templates2)):
    templates2[i] = image.Image(templates2[i])

# 循迹红色色块阈值
red = [24, 69, 0, 43, 7, 33]

# 循迹识别区块
areas = [(44, 56, 14, 8), (58, 56, 14, 8), (72, 56, 14, 8), (86, 56, 14, 8), (100, 56, 14, 8)]

# 数字识别标志
recognized_num = None
last_recognized_num = 0

# 得到循迹值函数
def getTrace(img):
    Trace = [0, 0, 0, 0, 0]
    for area in areas:
        statistics = img.get_statistics(roi=area)
        if red[0] < statistics.l_mean() < red[1] and red[2] < statistics.a_mean() < red[3] and red[4] < statistics.b_mean() < red[5]:
            Trace[areas.index(area)] = 1
    return Trace

# 标记函数
def Mark(img, Trace):
    for i in range(len(Trace)):
        area = areas[i]
        img.draw_rectangle(area)
        if Trace[i] == 1:
            img.draw_string(area[0] + 4, area[1], "1")
        else:
            img.draw_string(area[0] + 4, area[1], "0")

# 串口发送数据函数
def usart_send(Trace, recognized_num):
    Trace_num = 0
    if Trace:
        Trace_num = Trace[4] * 1 + Trace[3] * 2 + Trace[2] * 4 + Trace[1] * 8 + Trace[0] * 16
    data = bytearray([0xa3, 0xb3, Trace_num, recognized_num, 0xc3])
    usart3.write(data)
    print("Data sent:", [hex(b) for b in data])  # 打印发送的数据以调试

# 终端打印函数
def Print(Trace, recognized_num):
    print("Trace=", Trace, "Recognized Num=", recognized_num)

# 找数字函数：（返回数字的外框）
def findNum(img, templates):
    img_gray = img.to_grayscale()
    best_score = 0.0
    best_rect = None
    for t in templates:
        result = img_gray.find_template(t, 0.4, step=4, search=SEARCH_EX)
        if result is not None:
            print(f"Found template {t} at {result}")
            if best_rect is None or result[1] > best_score:
                best_score = result[1]
                best_rect = result[0]
    return best_rect

###########################################  主函数  ################################################

for i in range(50):
    img = sensor.snapshot()

data = bytearray([0x6B, 0x6B, 0x6B])
usart3.write(data)

while True:
    img = sensor.snapshot()

    if recognized_num is None:
        Rect2 = findNum(img, templates2)
        Rect1 = findNum(img, templates1)
        if Rect2 is not None:
            recognized_num = 2
        elif Rect1 is not None:
            recognized_num = 1
        if recognized_num is not None:
            last_recognized_num = recognized_num
            usart_send([], recognized_num=recognized_num)
            print("Recognized num:", recognized_num)

    Trace = getTrace(img)
    Mark(img, Trace)
    usart_send(Trace, last_recognized_num)
    Print(Trace, last_recognized_num)