import sensor, image, time, math
from pyb import UART, LED
from machine import I2C
from vl53l1x import VL53L1X

# 颜色阈值
mature_threshold = (5, 72, 108, 21, 94, -78)  # 成熟果实阈值
immature_threshold = (9, 76, -11, -65, 16, -47)  # 未成熟果实阈值

i2c = I2C(2)
distance = VL53L1X(i2c)

sensor.reset()  # 初始化摄像头
sensor.set_pixformat(sensor.RGB565)  # 图像格式
sensor.set_framesize(sensor.QQVGA)
sensor.skip_frames(time=2000)
sensor.set_auto_gain(False)      # 关闭自动增益以进行颜色追踪
sensor.set_auto_whitebal(False)  # 关闭自动白平衡以进行颜色追踪
clock = time.clock()  # 追踪帧率
sensor.set_vflip(True)  # 垂直方向翻转
sensor.set_hmirror(True)  # 水平方向翻转


def find_fruits():

    img = sensor.snapshot().lens_corr(strength=1.8,zoom=1.0)  # 畸变校正
    mature_blobs = img.find_blobs([mature_threshold],x_stride=15, y_stride=15, merge=True)
    immature_blobs = img.find_blobs([immature_threshold],x_stride=20, y_stride=20, merge=True)
    find_mature = False
    find_immature = False
    threshold_min = 100  # 过滤小于150像素的色块
    threshold_max = 6000  # 过滤大于600像素的色块


    for r in mature_blobs:
        if threshold_min < r[2] * r[3] < threshold_max:
            find_mature = True
            img.draw_rectangle(r.rect(), color=(255, 0, 0))
            img.draw_cross(r.cx(), r.cy(), size=2, color=(255, 0, 0))
            distance_mm = distance.read() - 90
            X, Y, Z = r.cx() - sensor.width() // 2, distance_mm , sensor.height() // 2 - r.cy()
            angles = calculate_joint_angles(X, Y, Z, 60,80,255,20,40 )
            for angle in angles:
                print("成熟果实坐标: X={:.2f}mm, Y={:.2f}mm, Z={:.2f}mm, 角度: j1={:.2f}, j2={:.2f}, j3={:.2f}, j4={:.2f}".format(X, Y, Z, *angle))

    for t in immature_blobs:
        if threshold_min < t[2] * t[3] < threshold_max and not find_mature:
            find_immature = True
            img.draw_rectangle(t.rect(), color=(0, 255, 0))
            img.draw_cross(t.cx(), t.cy(), size=2, color=(0, 255, 0))
            distance_mm = distance.read() - 90
            X, Y, Z = t.cx() - sensor.height()//2, distance_mm , sensor.width() // 2 - t.cy()
            angles = calculate_joint_angles(X, Y, Z,60,80,255,20,40)
            for angle in angles:
                print("未成熟果实坐标: X={:.2f}mm, Y={:.2f}mm, Z={:.2f}mm, 角度: j1={:.2f}, j2={:.2f}, j3={:.2f}, j4={:.2f}".format(X, Y, Z, *angle))

    if not mature_blobs and not immature_blobs:
        print("未找到果实")


def calculate_joint_angles(X, Y, Z, P,a1, a2, a3, a4):
    a1, a2, a3, a4 = 80,255,320,40
    P = 60

    if X == 0:
        j1 = 90.0
    else:
        j1 = math.atan((Y + P) / X) * 57.3

    n, m = 0, 0
    target_count = 180
    solution = []

    for i in range(target_count + 1):
        j_all = 3.1415927 * i / 180.0

        len_val = math.sqrt((Y + P) ** 2 + X ** 2)
        high_val = Z + 37

        L = len_val - a4 * math.sin(j_all)
        H = high_val - a4 * math.cos(j_all) - a1

        Cosj3 = (L**2 + H**2 - a2**2 - a3**2) / (2 * a2 * a3)
        if Cosj3 < -1 or Cosj3 > 1:
            continue  # 跳过不可能的配置
        Sinj3 = math.sqrt(1 - Cosj3**2)

        j3 = math.atan(Sinj3 / Cosj3) * 57.3

        K2 = a3 * math.sin(j3 / 57.3)
        K1 = a2 + a3 * math.cos(j3 / 57.3)

        Cosj2 = (K2 * L + K1 * H) / (K1**2 + K2**2)
        if Cosj2 < -1 or Cosj2 > 1:
            continue  # 跳过不可能的配置
        Sinj2 = math.sqrt(1 - Cosj2**2)

        j2 = math.atan(Sinj2 / Cosj2) * 57.3
        j4 = j_all * 57.3 - j2 - j3

        if 0 <= j2 <= 180 and 0 <= j3 <= 180 and -90 <= j4 <= 90:
            n += 1
            continue
    for i in range(target_count + 1):
        j_all = 3.1415927 * i / 180.0

        len_val = math.sqrt((Y + P) ** 2 + X ** 2)
        high_val = Z + 35.0

        L = len_val - a4 * math.sin(j_all)
        H = high_val - a4 * math.cos(j_all) - a1

        Cosj3 = (L**2 + H**2 - a2**2 - a3**2) / (2 * a2 * a3)
        if Cosj3 < -1 or Cosj3 > 1:
            continue  # 跳过不可能的配置
        Sinj3 = math.sqrt(1 - Cosj3**2)

        j3 = math.atan(Sinj3 / Cosj3) * 57.3

        K2 = a3 * math.sin(j3 / 57.3)
        K1 = a2 + a3 * math.cos(j3 / 57.3)
        if K1 ** 2 + K2 ** 2 == 0:
            continue
        Cosj2 = (K2 * L + K1 * H) / (K1**2 + K2**2)
        if Cosj2 < -1 or Cosj2 > 1:
            continue  # 跳过不可能的配置
        Sinj2 = math.sqrt(1 - Cosj2**2)

        j2 = math.atan(Sinj2 / Cosj2) * 57.3
        j4 = j_all * 57.3 - j2 - j3

        if 0 <= j2 <= 180 and 0 <= j3 <= 180 and -90 <= j4 <= 90:
            m += 1
            if m == n // 2 or m == (n + 1) // 2:
                solution.append((j1, j2, j3, j4))
                break

    return solution

def main():
    while True:
        find_fruits()
        time.sleep(0)  # 每隔5ms执行一次

if __name__ == "__main__":
    main()