#!/usr/bin/python3 #1
# coding=utf8 #2
# 第7章 ROS机器狗创意课程\3.AI自主追踪踢球(7.ROS Robot Creative Lesson\3.AI Auto Shooting) #3
import os #4
import sys #5
import cv2 #6
import time #7
import math #8
import threading #9
import numpy as np #10
from enum import Enum #11

from common import Misc #13

import rospy #15
from std_srvs.srv import * #16
from sensor_msgs.msg import Image #17
from ros_robot_controller.msg import RGBState, RGBsState #18
from object_tracking.srv import * #19
from puppy_control.msg import Velocity, Pose, Gait #20
from puppy_control.srv import SetRunActionName #21
from large_models.srv import SetString #22

ROS_NODE_NAME = 'kick_ball_demo' #24
is_shutdown = False #25

color_range_list = rospy.get_param('/lab_config_manager/color_range_list') #27

PuppyMove = {'x':0, 'y':0, 'yaw_rate':0} #29


if sys.version_info.major == 2: #32
    print('Please run this program with python3!') #33
    sys.exit(0) #34

lock = threading.Lock() #36
debug = False #37
__isRunning = True #38
haved_detect = False #39

class PuppyStatus(Enum): #41
    LOOKING_FOR = 0 #寻找 先低头寻找，没有就跳转到LOOKING_FOR_LEFT向左一步寻找，再没有的话就跳转到LOOKING_FOR_RIGHT，一直向右寻找(search, start by looking down,  If nothing is found, transition to LOOKING_FOR_LEFT to search one step to the left. If still nothing is found, transition to LOOKING_FOR_RIGHT and continue searching to the right) #42
    LOOKING_FOR_LEFT = 1 #43
    LOOKING_FOR_RIGHT = 2  #44
    FOUND_TARGET = 3 # 已经发现目标(target acquired) #45
    CLOSE_TO_TARGET = 4 # 靠近目标(approach target) #46
    CLOSE_TO_TARGET_FINE_TUNE = 5 # 细调(fine-tuning) #47
    KICK_BALL = 6 # 踢球(shooting) #48
    STOP = 10 #49
    END = 20             #50

puppyStatus = PuppyStatus.LOOKING_FOR #52
puppyStatusLast = PuppyStatus.END #53


expect_center = {'X':640/2,'Y':480/2} # #56
expect_center_kick_ball_left = {'X':150,'Y':480-150} # 左脚踢球，希望小球到达此坐标开始踢(left foot kicking the ball, aiming for it to reach this coordinate before starting) #57
expect_center_kick_ball_right = {'X':640-150,'Y':480-150}# 右脚踢球，希望小球到达此坐标开始踢(right foot kicking the ball, aiming for it to reach this coordinate before starting) #58
target_info = None # 小球中心点坐标(coordinates of the center point of the ball) #59

range_rgb = { #61
    'red': (0, 0, 255), #62
    'blue': (255, 0, 0), #63
    'green': (0, 255, 0), #64
    'black': (0, 0, 0), #65
    'white': (255, 255, 255), #66
} #67

color_list = [] #69
detect_color = 'None' #70
action_finish = True #71
draw_color = range_rgb["black"] #72
__target_color = ('red',) #73
thread_started = False  #74
start_time = 0 #75
elapsed_time = 0 #76
is_shutdown = False #77
enable_running = False #78
thread_started = True #79


# 找出面积最大的轮廓(find out the contour with the maximal area) #82
# 参数为要比较的轮廓的列表(the parameter is the list of contour to be compared) #83
def getAreaMaxContour(contours): #84
    contour_area_temp = 0 #85
    contour_area_max = 0 #86
    area_max_contour = None #87

    for c in contours:  # 历遍所有轮廓(iterate through all contours) #89
        contour_area_temp = math.fabs(cv2.contourArea(c))  # 计算轮廓面积(calculate the contour area) #90
        if contour_area_temp > contour_area_max: #91
            contour_area_max = contour_area_temp #92
            if contour_area_temp >= 5:  # 只有在面积大于300时，最大面积的轮廓才是有效的，以过滤干扰(only when the area is greater than 300, the contour with the largest area is considered valid to filter out interference) #93
                area_max_contour = c #94

    return area_max_contour, contour_area_max  # 返回最大的轮廓(return the maximal contour) #96


def move(): #99
    global detect_color, start_time, elapsed_time,is_shutdown,enable_running #100
    global puppyStatus, puppyStatusLast, haved_detect, action_finish, target_info, PuppyPose #101
    time.sleep(2) #102
    is_shutdown = True #103

    while is_shutdown: #105
        time.sleep(0.01) #106
        if not is_shutdown:break #107
        
        if enable_running: #109
            time.sleep(0.02) #110
            while(puppyStatus == PuppyStatus.LOOKING_FOR) : #111
                if haved_detect: #112
                    puppyStatus = PuppyStatus.FOUND_TARGET #113
                    start_time = 0 #114
                    break #115
                with lock: #116
                    if start_time == 0: #117
                        start_time == time.time() #118
                    PuppyPose = PP['LookDown_10deg'].copy() #119
                    PuppyPosePub.publish(stance_x=PuppyPose['stance_x'], stance_y=PuppyPose['stance_y'], x_shift=PuppyPose['x_shift'] #120
                        ,height=PuppyPose['height'], roll=PuppyPose['roll'], pitch=PuppyPose['pitch'], yaw=PuppyPose['yaw']) #121
                    time.sleep(0.2) #122
                time.sleep(0.8) #123
                puppyStatus = PuppyStatus.LOOKING_FOR_LEFT #124
                break #125
            while(puppyStatus == PuppyStatus.LOOKING_FOR_LEFT) : #126
                if haved_detect: #127
                    puppyStatus = PuppyStatus.FOUND_TARGET #128
                    break #129
                with lock: #130
                    PuppyVelocityPub.publish(x=0, y=0, yaw_rate = math.radians(10)) #131
                    time.sleep(3) #132
                    PuppyVelocityPub.publish(x=0, y=0, yaw_rate = math.radians(0)) #133
                    time.sleep(0.3) #134
                time.sleep(0.8) #135
                puppyStatus = PuppyStatus.LOOKING_FOR_RIGHT #136
                break #137
            while(puppyStatus == PuppyStatus.LOOKING_FOR_RIGHT) : #138
                if haved_detect: #139
                    start_time = 0 #140
                    puppyStatus = PuppyStatus.FOUND_TARGET #141
                    break #142
    
                PuppyPose = PP['LookDown_10deg'].copy() #144
                PuppyPosePub.publish(stance_x=PuppyPose['stance_x'], stance_y=PuppyPose['stance_y'], x_shift=PuppyPose['x_shift'] #145
                        ,height=PuppyPose['height'], roll=PuppyPose['roll'], pitch=PuppyPose['pitch'], yaw=PuppyPose['yaw']) #146
                time.sleep(0.2) #147
                PuppyVelocityPub.publish(x=2, y=0, yaw_rate = math.radians(-12)) #148
                elapsed_time = time.time() - start_time #149
                if elapsed_time > 3: #150
                    start_time = 0 #151
                    elapsed_time = 0 #152
                    PuppyVelocityPub.publish(x=0, y=0, yaw_rate = math.radians(0)) #153
                    puppyStatus == PuppyStatus.LOOKING_FOR #154
                    haved_detect = False                     #155
                    enable_running = False #156
                    is_shutdown = False #157
                    break               #158
                break #159
            while(puppyStatus == PuppyStatus.FOUND_TARGET) : #160
                # with lock: #161
                if target_info['centerY'] > 380: #162
                    puppyStatus = PuppyStatus.CLOSE_TO_TARGET #163
                    PuppyPose = PP['LookDown_20deg'].copy() #164
                    PuppyPosePub.publish(stance_x=PuppyPose['stance_x'], stance_y=PuppyPose['stance_y'], x_shift=PuppyPose['x_shift'] #165
                        ,height=PuppyPose['height'], roll=PuppyPose['roll'], pitch=PuppyPose['pitch'], yaw=PuppyPose['yaw']) #166
                    time.sleep(0.2) #167
                    break #168
                if expect_center['X'] - target_info['centerX'] < -80: #169
                    PuppyVelocityPub.publish(x=3, y=0, yaw_rate = math.radians(-12)) #170
                    time.sleep(0.2) #171
                elif expect_center['X'] - target_info['centerX'] > 80: #172
                    PuppyVelocityPub.publish(x=3, y=0, yaw_rate = math.radians(12)) #173
                    time.sleep(0.2) #174
                else: #175
                    PuppyVelocityPub.publish(x=10, y=0, yaw_rate = math.radians(0)) #176
                    time.sleep(0.2) #177
                break #178
            while(puppyStatus == PuppyStatus.CLOSE_TO_TARGET) : #179
                # with lock: #180
                if target_info['centerY'] > 380: #181
                    PuppyMove['x'] = 0 #182
                    PuppyMove['yaw_rate'] = math.radians(0) #183
                    PuppyVelocityPub.publish(x=0, y=0, yaw_rate = math.radians(0)) #184
                    puppyStatus = PuppyStatus.CLOSE_TO_TARGET_FINE_TUNE #185
                    print('expect_center[X] , target_info[centerX]',expect_center['X'] , target_info['centerX']) #186
                    if expect_center['X'] > target_info['centerX']: #187
                        which_foot_kick_ball = 'left' #188
                    else: #189
                        which_foot_kick_ball = 'right' #190
                    print(which_foot_kick_ball) #191
                    break #192
                if expect_center['X'] - target_info['centerX'] < -50: #193
                    PuppyVelocityPub.publish(x=3, y=0, yaw_rate = math.radians(-10)) #194
                    time.sleep(0.2) #195
                elif expect_center['X'] - target_info['centerX'] > 50: #196
                    PuppyVelocityPub.publish(x=3, y=0, yaw_rate = math.radians(10)) #197
                    time.sleep(0.2) #198
                else: #199
                    PuppyVelocityPub.publish(x=8, y=0, yaw_rate = math.radians(0)) #200
                    time.sleep(0.2) #201
                # print(target_info) #202
                break #203
            
            while(puppyStatus == PuppyStatus.CLOSE_TO_TARGET_FINE_TUNE) : #205
    
                if target_info['centerY'] < expect_center_kick_ball_left['Y']: #207
                    PuppyVelocityPub.publish(x=4, y=0, yaw_rate = math.radians(0)) #208
                    time.sleep(0.1) #209
                elif which_foot_kick_ball == 'left' and target_info['centerX'] > expect_center_kick_ball_left['X']: #210
                    PuppyVelocityPub.publish(x=0, y=0, yaw_rate = math.radians(-8)) #211
                    time.sleep(0.1) #212
                elif which_foot_kick_ball == 'right' and target_info['centerX'] < expect_center_kick_ball_right['X']: #213
                    PuppyVelocityPub.publish(x=0, y=0, yaw_rate = math.radians(8)) #214
                    time.sleep(0.1) #215
                else:# 最后一次微调(the final fine-tuning) #216
                    if which_foot_kick_ball == 'left': #217
                        PuppyVelocityPub.publish(x=5, y=0, yaw_rate = math.radians(-10)) #218
                    else: #219
                        PuppyVelocityPub.publish(x=5, y=0, yaw_rate = math.radians(10)) #220
                    time.sleep(1.8) #221
                    PuppyVelocityPub.publish(x=0, y=0, yaw_rate = math.radians(0)) #222
                    puppyStatus = PuppyStatus.KICK_BALL #223
                # PuppyVelocityPub.publish(x=0, y=0, yaw_rate = math.radians(0)) #224
                # time.sleep(0.1)#停下来需要稳定的时间(the time required to come to a stable stop) #225
                time.sleep(0.1) #226
                break #227
            while(puppyStatus == PuppyStatus.KICK_BALL) : #228
                with lock: #229
                    PuppyVelocityPub.publish(x=0, y=0, yaw_rate = math.radians(0)) #230
                    time.sleep(0.2) #231
                    if which_foot_kick_ball == 'left': #232
                        runActionGroup_srv('kick_ball_left.d6ac',True) #233
                        #os.system('rosnode kill kick_ball_demo') #234
                    else: #235
                        runActionGroup_srv('kick_ball_right.d6ac',True) #236
                        #os.system('rosnode kill kick_ball_demo')                   #237
                    haved_detect = False                     #238
                    enable_running = False #239
                    is_shutdown = False #240
                    time.sleep(1) #241
                    go_home_service()  #242
                    time.sleep(1) #243
                    puppyStatus = PuppyStatus.LOOKING_FOR #244
    
            if puppyStatus == PuppyStatus.STOP: #246
                PuppyMove['x'] = 0 #247
                PuppyMove['yaw_rate'] = math.radians(0) #248
                PuppyVelocityPub.publish(x=0, y=0, yaw_rate = math.radians(0)) #249
    
            
            if puppyStatusLast != puppyStatus: #252
                print('puppyStatus',puppyStatus) #253
            puppyStatusLast = puppyStatus #254
        else: #255
            time.sleep(0.02)             #256
    


size = (320, 240) #260
def run(img): #261
    global draw_color #262
    global color_list #263
    global detect_color #264
    global action_finish #265
    global haved_detect #266
    global target_info  #267
    global __target_color #268
    # img_copy = img.copy() #269
    img_h, img_w = img.shape[:2] #270

    if not __isRunning: #272
        return img #273

    frame_resize = cv2.resize(img, size, interpolation=cv2.INTER_NEAREST) #275
    frame_gb = cv2.GaussianBlur(frame_resize, (3, 3), 3)       #276
    frame_lab = cv2.cvtColor(frame_gb, cv2.COLOR_BGR2LAB)  # 将图像转换到LAB空间(convert the image to LAB space) #277

    max_area = 0 #279
    color_area_max = None     #280
    areaMaxContour_max = 0 #281
    
    if True:#action_finish #283
        for i in color_range_list: #284
            if i in __target_color: #285
                frame_mask = cv2.inRange(frame_lab, #286
                                             (color_range_list[i]['min'][0], #287
                                              color_range_list[i]['min'][1], #288
                                              color_range_list[i]['min'][2]), #289
                                             (color_range_list[i]['max'][0], #290
                                              color_range_list[i]['max'][1], #291
                                              color_range_list[i]['max'][2]))  #对原图像和掩模进行位运算(perform bitwise operation to original image and mask) #292
                eroded = cv2.erode(frame_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3)))  #腐蚀(corrosion) #293
                dilated = cv2.dilate(eroded, cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))) #膨胀(dilation) #294
                if debug: #295
                    cv2.imshow(i, dilated) #296
                contours = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)[-2]  #找出轮廓(find out the contour) #297
                areaMaxContour, area_max = getAreaMaxContour(contours)  #找出最大轮廓(find out the maximal contour) #298
                if areaMaxContour is not None: #299
                    if area_max > max_area:#找最大面积(find the maximal area) #300
                        max_area = area_max #301
                        color_area_max = i #302
                        areaMaxContour_max = areaMaxContour #303
                       
        if max_area > 200:  # 200   #305

            rect = cv2.minAreaRect(areaMaxContour_max)#最小外接矩形(the minimal bounding rectangle) #307
            
            box = np.intp(cv2.boxPoints(rect))#最小外接矩形的四个顶点(the four vertices of the minimum bounding rectangle) #309
            centerX = int(Misc.map(rect[0][0], 0, size[0], 0, img_w)) #310
            centerY = int(Misc.map(rect[0][1], 0, size[1], 0, img_h)) #311
            sideX = int(Misc.map(rect[1][0], 0, size[0], 0, img_w)) #312
            sideY = int(Misc.map(rect[1][1], 0, size[1], 0, img_h)) #313
            angle = rect[2] #314
            for i in range(4): #315
                box[i, 1] = int(Misc.map(box[i, 1], 0, size[1], 0, img_h)) #316
            for i in range(4):                 #317
                box[i, 0] = int(Misc.map(box[i, 0], 0, size[0], 0, img_w)) #318
            cv2.drawContours(img, [box], -1, (0,0,255,255), 2)#画出四个点组成的矩形(draw a rectangle formed by connecting the four points) #319


            if color_area_max == 'red':  #红色最大(red is the maximal area) #322
                color = 1 #323
            elif color_area_max == 'green':  #绿色最大(green is the maximal area) #324
                color = 2 #325
            elif color_area_max == 'blue':  #蓝色最大(blue is the maximal area) #326
                color = 3 #327
            else: #328
                color = 0 #329
            color_list.append(color) #330

            if len(color_list) == 3:  #多次判断(multiple judgement) #332
                # 取平均值(take the average value) #333
                color = int(round(np.mean(np.array(color_list)))) #334
                color_list = [] #335
                if color == 1: #336
                    detect_color = 'red' #337
                    draw_color = range_rgb["red"] #338
                elif color == 2: #339
                    detect_color = 'green' #340
                    draw_color = range_rgb["green"] #341
                elif color == 3: #342
                    detect_color = 'blue' #343
                    draw_color = range_rgb["blue"] #344
                else: #345
                    detect_color = 'None' #346
                    draw_color = range_rgb["black"]                #347
        else: #348
            detect_color = 'None' #349
            draw_color = range_rgb["black"] #350
        if detect_color == __target_color[0]: #351
            haved_detect = True #352
            if sideX > sideY: #353
                target_info = {'centerX':centerX, 'centerY':centerY, 'sideX':sideX, 'sideY':sideY, 'scale':sideX/sideY, 'angle':angle} #354
            else: #355
                target_info = {'centerX':centerX, 'centerY':centerY, 'sideX':sideX, 'sideY':sideY, 'scale':sideY/sideX, 'angle':angle} #356
        else: #357
            haved_detect = False #358
    cv2.putText(img, "Color: " + detect_color, (10, img.shape[0] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.65, draw_color, 2) #359
    
    return img #361


def image_callback(ros_image): #364
    global thread_started, enable_running #365
    
    image = np.ndarray(shape=(ros_image.height, ros_image.width, 3), dtype=np.uint8, #367
                       buffer=ros_image.data)  # 将自定义图像消息转化为图像(convert the customize image information to image) #368
    cv2_img = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) #369
    
    frame = cv2_img.copy() #371
    frame_result = frame #372
    
    if not thread_started: #374
        thread = threading.Thread(target=move) #375
        thread.start() #376
        thread_started = True     #377
    else: #378
        time.sleep(0.02) #379
        
    
    if enable_running: #382
        frame_result = run(frame) #383
        cv2.imshow('Frame', frame_result) #384
        key = cv2.waitKey(1) #385
    else:             #386
        cv2.destroyAllWindows() #387
        
def init_pose(): #389
    GaitConfig = {'overlap_time':0.15, 'swing_time':0.15, 'clearance_time':0.0, 'z_clearance':3} #390
    rospy.sleep(0.3) #391
    PuppyPosePub.publish(stance_x=PuppyPose['stance_x'], stance_y=PuppyPose['stance_y'], x_shift=PuppyPose['x_shift'] #392
            ,height=PuppyPose['height'], roll=PuppyPose['roll'], pitch=PuppyPose['pitch'], yaw=PuppyPose['yaw'], run_time = 500) #393
    
    rospy.sleep(0.2) #395
    PuppyGaitConfigPub.publish(overlap_time = GaitConfig['overlap_time'], swing_time = GaitConfig['swing_time'] #396
                    , clearance_time = GaitConfig['clearance_time'], z_clearance = GaitConfig['z_clearance']) #397
    rospy.sleep(0.2) #398

def cleanup(): #400
    global is_shutdown,__target_color #401
    is_shutdown = True #402
    PuppyVelocityPub.publish(x=0, y=0, yaw_rate=0) #403
    print('is_shutdown') #404
   
def enter_func(msg): #406
    global thread_started #407
    
    rospy.loginfo("enter kick ball") #409
    thread_started = False #410
    rospy.sleep(0.02) #411
    init_pose() #412
    
    response = TriggerResponse() #414
    response.success = True #415
    response.message = "start" #416
    return response #417
   
    
def set_running(msg): #420
    rospy.loginfo("enter kick ball") #421
    global enable_running #422
    
    if msg.data: #424
        enable_running = msg.data #425
    else: #426
        enable_running = False #427
    
    response = SetBoolResponse() #429
    response.success = True #430
    return response #431
    
    
def set_target(msg): #434
    global __target_color #435
    
    rospy.loginfo("%s", msg) #437
    __target_color = (msg.data, ) #438
    
    response = SetBoolResponse() #440
    response.success = True #441
        
    return [True, 'set_target'] #443
    
    
if __name__ == '__main__': #446
    rospy.init_node(ROS_NODE_NAME, log_level=rospy.DEBUG) #447
    

    param_name = '/puppy_control/PuppyPose' #450
    while not rospy.has_param(param_name): #451
        rospy.logwarn("parameter %s not available yet. Waiting....",param_name) #452
        rospy.sleep(1) #453
    
    PP = rospy.get_param('/puppy_control/PuppyPose') #455
    PuppyPose = PP['LookDown_10deg'].copy() #456
    
    __target_color = ("red",) #458
    
    GaitConfig = {'overlap_time':0.15, 'swing_time':0.15, 'clearance_time':0.0, 'z_clearance':3} #460

    image_sub = rospy.Subscriber('/usb_cam/image_raw', Image, image_callback) #462

    image_pub = rospy.Publisher('/%s/image_result'%ROS_NODE_NAME, Image, queue_size=1)  # register result image publisher #464


    PuppyGaitConfigPub = rospy.Publisher('/puppy_control/gait', Gait, queue_size=1) #467
    PuppyVelocityPub = rospy.Publisher('/puppy_control/velocity', Velocity, queue_size=1) #468
    PuppyPosePub = rospy.Publisher('/puppy_control/pose', Pose, queue_size=1) #469

    runActionGroup_srv = rospy.ServiceProxy('/puppy_control/runActionGroup', SetRunActionName) #471
    go_home_service = rospy.ServiceProxy('/puppy_control/go_home', Empty) #472
    
    enter_srv = rospy.Service('/%s/enter'%ROS_NODE_NAME, Trigger, enter_func) #474
    running_srv = rospy.Service('/%s/enable_running'%ROS_NODE_NAME, SetBool, set_running) #475
    set_target_srv = rospy.Service('/%s/set_color_target'%ROS_NODE_NAME, SetString, set_target) #476
    
    rospy.on_shutdown(cleanup) #478

    ''' #480
    rospy.sleep(0.3) #481
    PuppyPosePub.publish(stance_x=PuppyPose['stance_x'], stance_y=PuppyPose['stance_y'], x_shift=PuppyPose['x_shift'] #482
            ,height=PuppyPose['height'], roll=PuppyPose['roll'], pitch=PuppyPose['pitch'], yaw=PuppyPose['yaw'], run_time = 500) #483
    
    rospy.sleep(0.2) #485
    PuppyGaitConfigPub.publish(overlap_time = GaitConfig['overlap_time'], swing_time = GaitConfig['swing_time'] #486
                    , clearance_time = GaitConfig['clearance_time'], z_clearance = GaitConfig['z_clearance']) #487
    rospy.sleep(0.2) #488
    ''' #489
    


    try: #493
        rospy.spin() #494
    except KeyboardInterrupt: #495
        print("Shutting down") #496
    finally: #497
        cv2.destroyAllWindows() #498

