#!/usr/bin/env python3 #1
# coding=utf8 #2

import cv2 #4
import time #5
import math #6
import threading #7
import numpy as np #8

from enum import Enum #10
from arm_kinematics.ArmMoveIK import ArmIK #11
from common import Misc, yaml_handle #12
from servo_controller import setServoPulse #13
from action_group_control import runActionGroup #14
from HiwonderPuppy import HiwonderPuppy, PWMServoParams #15

class PuppyStatus(Enum): #17
    IDLE = 0          # Idle state, waiting for color setting or manual start （空闲状态，等待颜色设置或手动启动） #18
    START = 1         # Startup phase, initializing posture （启动阶段，初始化姿态） #19
    NORMAL = 2        # Moving forward normally （正常前进中） #20
    FOUND_TARGET = 3  # Target block detected （已经发现目标物块） #21
    ADJUSTING_YAW = 4 # Fine-tuning yaw direction （微调方向（yaw）状态） #22
    ADJUSTING_X = 5   # Fine-tuning X position （微调位置（x）状态） #23
    STOP = 10         # Stop state, some services disabled （停止状态，禁用部分服务） #24

class PuppyController: #26
    def __init__(self, debug=False): #27
        self.Debug = debug #28

        # Initialize robot dog parameters and related objects （初始化机器狗相关参数和对象） #30
        self.puppy = HiwonderPuppy(setServoPulse=setServoPulse, servoParams=PWMServoParams(), dof='8') #31

        self.stand_pose = { #33
            'roll': math.radians(0.0), #34
            'pitch': math.radians(0.0), #35
            'yaw': 0.0, #36
            'height': -10.0, #37
            'x_shift': -1.0, #38
            'stance_x': 0.0, #39
            'stance_y': 0.0 #40
        } #41
        self.bend_pose = { #42
            'roll': math.radians(0.0), #43
            'pitch': math.radians(-17.0), #44
            'yaw': 0.0, #45
            'height': -10.0, #46
            'x_shift': 0.5, #47
            'stance_x': 0.0, #48
            'stance_y': 0.0 #49
        } #50
        self.current_pose = self.stand_pose.copy() #51

        self.gait_config = { #53
            'overlap_time': 0.15, #54
            'swing_time': 0.15, #55
            'clearance_time': 0.0, #56
            'z_clearance': 3 #57
        } #58

        self.puppy_move = {'x': 0, 'y': 0, 'yaw_rate': 0} #60

        self.arm_ik = ArmIK() #62
        self.color_list = [] #63

        self.status = PuppyStatus.IDLE #65
        self.status_last = PuppyStatus.STOP #66

        self.img_centerx = 320 #68
        self.block_color = 'None' #69
        self.block_center_point = (0, 0) #70
        self.radius = 0 #71
        self.size = (640, 480) #72
        self.lab_data = None #73
        self.MIN_RADIUS = 129 #74
        self.MAX_RADIUS = 140 #75

        self.color_tracking_enabled = False #77
        self.target_color = 'None' #78

        self.adjusting_counter_yaw = 0 #80
        self.required_frames_yaw = 3 #81

        # Grabbing range parameters, adjust y_close_lower and y_close_upper if the target is far （夹取范围参数（根据打印的范围参考，如果离目标远就增加 y_close_lower 和 y_close_upper 的范围）） #83
        self.y_close_lower = 410 #84
        self.y_close_upper = 420 #85
        self.x_error_thresh = 30 #86

        self.status_lock = threading.Lock() #88

        self.running = True #90

        self.load_config() #92
        self.initialize_puppy() #93

        self.puppy.start() #95
        self.puppy.move_stop(servo_run_time=500) #96
        self.arm_ik.setPitchRangeMoving((8.51, 0, 3.3), 500) #97
        setServoPulse(9, 1500, 500) #98
        time.sleep(0.5) #99

        if self.Debug: #101
            self.current_pose = self.bend_pose.copy() #102
            self.puppy.stance_config( #103
                self.stance(self.current_pose['stance_x'], self.current_pose['stance_y'], #104
                            self.current_pose['height'], self.current_pose['x_shift']), #105
                self.current_pose['pitch'], self.current_pose['roll'] #106
            ) #107
            time.sleep(0.5) #108
        else: #109
            self.move_thread = threading.Thread(target=self.move) #110
            self.move_thread.daemon = True #111
            self.move_thread.start() #112

    def load_config(self): #114
        self.lab_data = yaml_handle.get_yaml_data(yaml_handle.lab_file_path)['color_range_list'] #115

    def stance(self, x=0.0, y=0.0, z=-11.0, x_shift=2.0): #117
        return np.array([ #118
            [x + x_shift, x + x_shift, -x + x_shift, -x + x_shift], #119
            [y, y, y, y], #120
            [z, z, z, z], #121
        ]) #122

    def initialize_puppy(self): #124
        self.puppy.stance_config( #125
            self.stance(self.current_pose['stance_x'], self.current_pose['stance_y'], #126
                        self.current_pose['height'], self.current_pose['x_shift']), #127
            self.current_pose['pitch'], self.current_pose['roll'] #128
        ) #129
        self.puppy.gait_config(**self.gait_config) #130

    def get_area_max_contour(self, contours): #132
        contour_area_max = 0 #133
        area_max_contour = None #134
        for c in contours: #135
            contour_area = abs(cv2.contourArea(c)) #136
            if contour_area > contour_area_max and contour_area >= 5: #137
                contour_area_max = contour_area #138
                area_max_contour = c #139
        return area_max_contour, contour_area_max #140

    def process_image(self, img): #142
        """ #143
        Image processing, color detection, target tracking, and drawing visual aids （图像处理，颜色检测，目标跟踪，画辅助信息） #144
        Return the annotated image （返回标记好的图像） #145
        """ #146
        if not self.color_tracking_enabled: #147
            return img #148

        img_copy = img.copy() #150
        img_h, img_w = img.shape[:2] #151

        frame_resize = cv2.resize(img_copy, self.size, interpolation=cv2.INTER_NEAREST) #153
        frame_gb = cv2.GaussianBlur(frame_resize, (3, 3), 3) #154
        frame_lab_all = cv2.cvtColor(frame_gb, cv2.COLOR_BGR2LAB) #155

        if self.target_color not in self.lab_data: #157
            if self.Debug: #158
                print(f"目标颜色 {self.target_color} 未配置") #159
            return img #160

        color_range = self.lab_data[self.target_color] #162
        frame_mask = cv2.inRange(frame_lab_all, #163
                                 tuple(color_range['min']), #164
                                 tuple(color_range['max'])) #165
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3)) #166
        eroded = cv2.erode(frame_mask, kernel) #167
        dilated = cv2.dilate(eroded, kernel) #168

        contours = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)[-2] #170
        max_contour, max_area = self.get_area_max_contour(contours) #171

        if max_contour is not None and max_area >= 5: #173
            ((centerX, centerY), radius_temp) = cv2.minEnclosingCircle(max_contour) #174
            centerX = int(Misc.map(centerX, 0, self.size[0], 0, img_w)) #175
            centerY = int(Misc.map(centerY, 0, self.size[1], 0, img_h)) #176
            radius = int(Misc.map(radius_temp, 0, self.size[0], 0, img_w)) #177
            print(f"[夹取范围检测] centerY={centerY}, range=[{self.y_close_lower}, {self.y_close_upper}], in_range={self.y_close_lower < centerY < self.y_close_upper}") #178


            self.block_center_point = (centerX, centerY) #181
            self.radius = radius #182

            cv2.circle(img, self.block_center_point, self.radius, (0, 255, 255), 2) #184

            color = 1 if self.MIN_RADIUS <= self.radius <= self.MAX_RADIUS else 0 #186
            self.color_list.append(color) #187

            if len(self.color_list) == 3: #189
                color_avg = int(round(np.mean(np.array(self.color_list)))) #190
                self.color_list = [] #191
                self.block_color = self.target_color if color_avg == 1 else 'None' #192

            # Multi-frame stability check, based on Y coordinate and X deviation （多帧稳定性判断，条件中 y 坐标和 x 偏差） #194
            if (self.y_close_lower < centerY < self.y_close_upper and #195
                self.block_color == self.target_color and #196
                abs(centerX - self.img_centerx) < self.x_error_thresh): #197
                self.adjusting_counter_yaw += 1 #198
                if self.adjusting_counter_yaw >= self.required_frames_yaw: #199
                    try: #200
                        self.puppy.move_stop(servo_run_time=500) #201
                    except Exception: #202
                        pass #203
                    self.status = PuppyStatus.ADJUSTING_YAW #204
                    self.adjusting_counter_yaw = 0 #205
                    time.sleep(0.5) #206
            else: #207
                self.adjusting_counter_yaw = 0 #208
        else: #209
            self.block_color = 'None' #210
            self.radius = 0 #211

        if self.status == PuppyStatus.STOP and self.block_color == self.target_color: #213
            with self.status_lock: #214
                self.status = PuppyStatus.START #215
                self.color_tracking_enabled = True #216

        # --- Draw visual aids （画辅助信息） --- #218
        # Blue rectangle indicating the grabbing range （夹取范围蓝色矩形框） #219
        cv2.rectangle(img, (0, self.y_close_lower), (img_w, self.y_close_upper), (255, 0, 0), 2) #220
        # Green dot indicating the target center （目标中心绿点） #221
        if self.block_center_point != (0, 0): #222
            cv2.circle(img, self.block_center_point, 5, (0, 255, 0), -1) #223
        # Vertical center line （垂直中心线） #224
        cv2.line(img, (self.img_centerx, 0), (self.img_centerx, img_h), (0, 255, 255), 2) #225
        # Red arrow for yaw adjustment （方向调整箭头（红）） #226
        if self.status == PuppyStatus.ADJUSTING_YAW: #227
            if self.puppy_move['yaw_rate'] > 0: #228
                cv2.arrowedLine(img, (self.img_centerx, self.block_center_point[1]), #229
                                (self.img_centerx + 30, self.block_center_point[1]), (0, 0, 255), 2) #230
            elif self.puppy_move['yaw_rate'] < 0: #231
                cv2.arrowedLine(img, (self.img_centerx, self.block_center_point[1]), #232
                                (self.img_centerx - 30, self.block_center_point[1]), (0, 0, 255), 2) #233
        # Green arrow for position adjustment （位置调整箭头（绿）） #234
        if self.status == PuppyStatus.ADJUSTING_X: #235
            if self.puppy_move['x'] > 0: #236
                cv2.arrowedLine(img, (self.block_center_point[0], self.block_center_point[1]), #237
                                (self.block_center_point[0], self.block_center_point[1] - 30), (0, 255, 0), 2) #238
            elif self.puppy_move['x'] < 0: #239
                cv2.arrowedLine(img, (self.block_center_point[0], self.block_center_point[1]), #240
                                (self.block_center_point[0], self.block_center_point[1] + 30), (0, 255, 0), 2) #241

        return img #243

    def is_within_clamp_range(self, y, x_error): #245
        y_close = self.y_close_lower < y <= self.y_close_upper #246
        x_close = abs(x_error) < 30 #247
        return y_close and x_close #248

    def is_within_range(self): #250
        y = self.block_center_point[1] #251
        x_error = self.block_center_point[0] - self.img_centerx #252
        y_in_range = self.y_close_lower < y < self.y_close_upper #253
        x_in_range = abs(x_error) < 30 #254
        color_correct = self.block_color == self.target_color #255
        return y_in_range and x_in_range and color_correct #256


    def move(self): #259
        while self.running: #260
            try: #261
                if self.status == PuppyStatus.START: #262
                    self.puppy.move_stop(servo_run_time=500) #263
                    self.current_pose = self.bend_pose.copy() #264
                    self.puppy.stance_config( #265
                        self.stance(self.current_pose['stance_x'], self.current_pose['stance_y'], #266
                                    self.current_pose['height'], self.current_pose['x_shift']), #267
                        self.current_pose['pitch'], self.current_pose['roll']) #268
                    time.sleep(0.5) #269
                    self.status = PuppyStatus.NORMAL #270

                elif self.status == PuppyStatus.NORMAL: #272
                    if self.color_tracking_enabled: #273
                        value = self.block_center_point[0] - self.img_centerx #274
                        y = self.block_center_point[1] #275

                        if self.is_within_clamp_range(y, value): #277
                            self.status = PuppyStatus.FOUND_TARGET #278
                            self.puppy.move_stop(servo_run_time=500) #279
                        else: #280
                            move_x = 0.0 #281
                            yaw_rate = 0.0 #282
                            if y < self.y_close_lower: #283
                                move_x = 8.0 #284
                            if abs(value) > 80: #285
                                yaw_rate = math.radians(-8.0 * np.sign(value)) #286
                            elif abs(value) > 50.0: #287
                                yaw_rate = math.radians(-6.0 * np.sign(value)) #288

                            self.puppy_move['x'] = move_x #290
                            self.puppy_move['yaw_rate'] = yaw_rate #291
                            self.puppy.move(x=move_x, y=0, yaw_rate=yaw_rate) #292

                            if self.is_within_range(): #294
                                self.status = PuppyStatus.FOUND_TARGET #295
                                self.puppy.move_stop(servo_run_time=500) #296
                            else: #297
                                if self.y_close_lower < y < self.y_close_upper and abs(value) < 30: #298
                                    self.adjusting_counter_yaw += 1 #299
                                    if self.adjusting_counter_yaw >= self.required_frames_yaw: #300
                                        self.status = PuppyStatus.ADJUSTING_YAW #301
                                        self.adjusting_counter_yaw = 0 #302
                    else: #303
                        self.puppy.move_stop(servo_run_time=100) #304
                    time.sleep(0.02) #305

                elif self.status == PuppyStatus.ADJUSTING_YAW: #307
                    error_x = self.block_center_point[0] - self.img_centerx #308
                    yaw_rate = 0.0 #309

                    if error_x < -10: #311
                        yaw_rate = math.radians(6.0)  # Turn right （向右转） #312
                    elif error_x > 10: #313
                        yaw_rate = math.radians(-6.0)  # Turn left （向左转） #314

                    self.puppy_move['yaw_rate'] = yaw_rate #316
                    self.puppy_move['x'] = 0 #317

                    if yaw_rate != 0.0: #319
                        try: #320
                            self.puppy.move(x=0, y=0, yaw_rate=yaw_rate) #321
                        except Exception: #322
                            pass #323

                        time.sleep(0.6) #325
                        self.puppy.move_stop(servo_run_time=510) #326
                        time.sleep(0.6) #327

                    # Check alignment, X deviation ≤ 10 （检查是否对齐（x 误差小于等于10）） #329
                    if abs(error_x) <= 10: #330
                        y = self.block_center_point[1] #331
                        if self.y_close_lower < y < self.y_close_upper: #332
                            self.status = PuppyStatus.ADJUSTING_X #333
                        else: #334
                            self.status = PuppyStatus.NORMAL #335
                    else: #336
                        # Continue fine-tuning state （继续保持微调状态） #337
                        pass #338

                elif self.status == PuppyStatus.ADJUSTING_X: #340
                    # Position adjustment: adjust in the Y direction (forward/backward) （位置微调：调整 y 方向（前后）） #341
                    error_y = self.block_center_point[1] - 365  # Target center preset: 365 （目标中心预设：365） #342

                    move_x = 0.0 #344
                    if error_y < -10: #345
                        move_x = 3.0  # Move forward （向前移动） #346
                    elif error_y > 10: #347
                        move_x = -3.0  # Move backward （向后移动） #348

                    y = self.block_center_point[1] #350
                    if self.y_close_lower < y < self.y_close_upper: #351
                        self.puppy_move['x'] = move_x #352
                        self.puppy_move['yaw_rate'] = 0.0 #353

                        if move_x != 0.0: #355
                            try: #356
                                self.puppy.move(x=move_x, y=0, yaw_rate=0) #357
                            except Exception: #358
                                pass #359
                            time.sleep(0.6) #360
                            self.puppy.move_stop(servo_run_time=510) #361
                            time.sleep(0.51) #362

                        if abs(error_y) <= 10: #364
                            self.status = PuppyStatus.FOUND_TARGET #365
                        else: #366
                            # Continue position adjustment （继续位置微调） #367
                            pass #368
                    else: #369
                        # Target not close enough, return to normal state to continue searching （目标物不够近，跳回正常状态继续搜索） #370
                        self.status = PuppyStatus.NORMAL #371

                elif self.status == PuppyStatus.FOUND_TARGET: #373
                    time.sleep(2)  # Reserve preparation time for the grabbing action （给夹取动作预留准备时间） #374
                    try: #375
                        runActionGroup('Clamping.d6a', True) #376
                        self.current_pose = self.stand_pose.copy() #377
                        self.puppy.stance_config( #378
                            self.stance(self.current_pose['stance_x'], self.current_pose['stance_y'], #379
                                        self.current_pose['height'], self.current_pose['x_shift'] + 1), #380
                            self.current_pose['pitch'], self.current_pose['roll']) #381
                        time.sleep(0.5) #382
                        runActionGroup('place1.d6a', True) #383
                        time.sleep(1) #384
                        runActionGroup('look_down.d6a', True) #385
                        time.sleep(1) #386

                        with self.status_lock: #388
                            self.status = PuppyStatus.STOP #389
                            self.color_tracking_enabled = False #390
                    except Exception: #391
                        with self.status_lock: #392
                            self.status = PuppyStatus.NORMAL #393

                elif self.status == PuppyStatus.STOP: #395
                    try: #396
                        self.puppy.move_stop(servo_run_time=500) #397
                    except Exception: #398
                        pass #399

                    self.current_pose = self.stand_pose.copy() #401
                    self.puppy.stance_config( #402
                        self.stance(self.current_pose['stance_x'], self.current_pose['stance_y'], #403
                                    self.current_pose['height'], self.current_pose['x_shift']), #404
                        self.current_pose['pitch'], self.current_pose['roll']) #405
                    time.sleep(0.1) #406

                else: #408
                    time.sleep(0.02) #409

                with self.status_lock: #411
                    if self.status_last != self.status: #412
                        if self.Debug: #413
                            print(f"状态由 {self.status_last} 转换为 {self.status}") #414
                        self.status_last = self.status #415
                time.sleep(0.02) #416

            except Exception as e: #418
                if self.Debug: #419
                    print(f"控制线程异常: {e}") #420
                with self.status_lock: #421
                    self.status = PuppyStatus.NORMAL #422

if __name__ == '__main__': #424
    print("请作为模块导入使用") #425
