RoboRally/remote_control/class_control_joystick.py

import socket
import pygame
import pygame.joystick
import rospy
import pygame.freetype
from marker_pos_angle.msg import id_pos_angle
from shapely.geometry import Polygon
import matplotlib.pyplot as plt
from shapely.geometry import Point

from copy import deepcopy

import time




# for debugging:
# rosrun image_view image_view image:=/fiducial_images

class race_track:
    """
    This class existing for a checking position of robots on the track
    We have four aruco markers for outside rectangle and four for inside small rectangle
    If one of the robots stay inside of the big rectangle we can to control its, but in another case works
    only keyboard control - we use it for came back of robots in there
    """

    def __init__(self, robot_1, robot_2):

        marker_sub = rospy.Subscriber("/marker_id_pos_angle", id_pos_angle, self.measurement_callback)

        self.robot_1 = robot_1
        self.robot_2 = robot_2

        self.rectangle_pos = None
        self.rectangle_pos_2 = None
        self.rectangle_pos_3 = None
        self.rectangle_pos_4 = None
        self.rectangle_pos_5 = None
        self.rectangle_pos_6 = None
        self.rectangle_pos_7 = None
        self.rectangle_pos_8 = None
        self.robot_pos_1 = None
        self.robot_pos_2 = None
        self.start_line = None
        self.checkpoint_line = None
        self.check_number_markers = True


        self.debug_plot = False

        self.checkpoint_count = 0
        self.finish_line_count = 0

        self.checkpoint_count_2 = 0
        self.finish_line_count_2 = 0


    def measurement_callback(self, data):
        # Aruco marker number 14 it's robot 102; Aruco 1,5,7,8 are drawn small rectangle Polygon
        # Aruco 2,4,0,11 are assembled big rectangle Polygon

        if data.id == 14:
            self.robot_pos_1 = (data.x, data.y)
            # print ("{}->ID MARKER, {}->POSITION".format(data.id, self.robot_pos_1))
        elif data.id == 15:
            self.robot_pos_2 = (data.x, data.y)
            # print ("{}->ID MARKER, {}->POSITION".format(data.id, self.robot_pos_1))
        elif data.id == 9:
            self.checkpoint_line = (data.x, data.y)
            # print ("{}->ID MARKER, {}->POSITION".format(data.id, self.start_line))
        elif data.id == 0:
            self.start_line = (data.x, data.y)
            # print ("{}->ID MARKER, {}->POSITION".format(data.id, self.start_line))

        elif data.id == 1:
            self.rectangle_pos = (data.x, data.y)
            # print ("{}->ID MARKER, {}->POSITION".format(data.id, self.rectangle_pos))

        elif data.id == 2:
            self.rectangle_pos_2 = (data.x, data.y)
            # print ("{}->ID MARKER, {}->POSITION".format(data.id, self.rectangle_pos_2))

        elif data.id == 3:
            self.rectangle_pos_3 = (data.x, data.y)
            # print ("{}->ID MARKER, {}->POSITION".format(data.id, start_marker_pos_3))

        elif data.id == 4:
            self.rectangle_pos_4 = (data.x, data.y)
            # print ("{}->ID MARKER, {}->POSITION".format(data.id,start_marker_pos_4))
        elif data.id == 5:
            self.rectangle_pos_5 = (data.x, data.y)
            # print ("{}->ID MARKER, {}->POSITION".format(data.id, self.rectangle_pos_5))

        elif data.id == 6:
            self.rectangle_pos_6 = (data.x, data.y)
            # print ("{}->ID MARKER, {}->POSITION".format(data.id,self.rectangle_pos_6))

        elif data.id == 7:
            self.rectangle_pos_7 = (data.x, data.y)
            # print ("{}->ID MARKER, {}->POSITION".format(data.id, self.rectangle_pos_7))

        elif data.id == 8:
            self.rectangle_pos_8 = (data.x, data.y)
            # print ("{}->ID MARKER, {}->POSITION".format(data.id,self.rectangle_pos_8))

    def check_number_marker(self):
        if self.rectangle_pos is None:
            self.check_number_markers = False
        elif self.rectangle_pos_2 is None:
            self.check_number_markers = False
        elif self.rectangle_pos_3 is None:
            self.check_number_markers = False
        elif self.rectangle_pos_4 is None:
            self.check_number_markers = False
        elif self.rectangle_pos_5 is None:
            self.check_number_markers = False
        elif self.rectangle_pos_6 is None:
            self.check_number_markers = False
        elif self.rectangle_pos_7 is None:
            self.check_number_markers = False
        elif self.rectangle_pos_8 is None:
            self.check_number_markers = False
        elif self.start_line is None:
            self.check_number_markers = False
        elif self.checkpoint_line is None:
            self.check_number_markers = False
        else:
            self.check_number_markers = True

    def polygon_rectangle(self):
        if self.check_number_markers:
            self.rectangle_track = Polygon([self.rectangle_pos, self.rectangle_pos_2,
                                            self.rectangle_pos_3, self.rectangle_pos_4])

            self.rectangle_track_2 = Polygon([self.rectangle_pos_5, self.rectangle_pos_6,
                                              self.rectangle_pos_7, self.rectangle_pos_8])

            self.start_strength_line = Polygon([(self.rectangle_pos_7[0]-0.08, self.rectangle_pos_7[1]),
                                                (self.start_line[0]-0.08, self.start_line[1]),
                                                (self.start_line[0]+0.08, self.start_line[1]),
                                                (self.rectangle_pos_7[0]+0.08, self.rectangle_pos_7[1])])

            self.checkpoint_strength_line = Polygon([self.checkpoint_line, self.rectangle_pos_5,
                                                (self.rectangle_pos_5[0]+0.1, self.rectangle_pos_5[1]),
                                                (self.checkpoint_line[0] + 0.1, self.checkpoint_line[1])])
            self.robot_pos_1_point = Point(self.robot_pos_1)
            self.robot_pos_2_point = Point(self.robot_pos_2)

            if self.debug_plot:
                x, y = (self.checkpoint_strength_line.exterior.xy)
                plt.plot(x, y)
                plt.show()



            if self.rectangle_track.contains(self.robot_pos_1_point) and not \
                    self.rectangle_track_2.contains(self.robot_pos_1_point):

                self.robot_1.controlling_allowed = True
            else:
                #print("robot outside track")
                self.robot_1.controlling_allowed = False

            if self.rectangle_track.contains(self.robot_pos_2_point) and not \
                    self.rectangle_track_2.contains(self.robot_pos_2_point):

                self.robot_2.controlling_allowed = True
            else:
                #print("robot_2 outside track")
                self.robot_2.controlling_allowed = False
        else:
            print("couldn't find the all ARUCO markers")

    def check_starting_position(self): # when you pressing B robot must staying on the start line for beginning of race
        button_robot_1 = self.robot_1.joystick.get_button(1)
        button_robot_2 = self.robot_2.joystick.get_button(1)
        if button_robot_1:
            if self.start_strength_line.contains(self.robot_pos_1_point):
                print("player 1 is on the start line")
                self.robot_1.controlling_allowed_2 = True
            else:
                print("put the player 1 on the start line")
                self.robot_1.controlling_allowed_2 = False
        if button_robot_2:
            if self.start_strength_line.contains(self.robot_pos_2_point):
                print("player 2 is on the start line")
                self.robot_2.controlling_allowed_2 = True
            else:
                print("put the player 2 on the start line")
                self.robot_2.controlling_allowed_2 = False
        else:
            print("no button pressed")

    def checkpoint(self):
        if self.checkpoint_strength_line.contains(self.robot_pos_1_point):
            if self.finish_line_count == self.checkpoint_count:
                self.checkpoint_count += 1
                print("Player 1 reached checkpoint, checkpoint count = {}".format(self.checkpoint_count))
        if self.checkpoint_strength_line.contains(self.robot_pos_2_point):
            if self.finish_line_count_2 == self.checkpoint_count_2:
                self.checkpoint_count_2 += 1
                print("Player 2 reached checkpoint, checkpoint count = {}".format(self.checkpoint_count_2))

    def finish_line(self):
        if self.start_strength_line.contains(self.robot_pos_1_point):
            if self.finish_line_count < self.checkpoint_count:
                self.finish_line_count += 1
                print("reached finish line, finish line count = {}".format(self.finish_line_count))

            elif self.finish_line_count == self.checkpoint_count and self.finish_line_count == 3:
                print("Player1 finished")
                self.robot_1.controlling_allowed = False

        if self.start_strength_line.contains(self.robot_pos_2_point):
            if self.finish_line_count_2 < self.checkpoint_count_2:
                self.finish_line_count_2 += 1
                print("Player 2 reached finish line, finish line count = {}".format(self.finish_line_count_2))

            elif self.finish_line_count_2 == self.checkpoint_count_2 and self.finish_line_count_2 == 3:
                self.robot_2.controlling_allowed = False
                self.player_2 = False
                if not self.player_2:
                    print("Player_2 win")
                    self.player_2 = True


class robot:  # we have 4 arg for this class, because joysticks get the same (value, axis) events
    def __init__(self, joy, ip, port):
        self.joy = joy
        self.ip = ip
        self.port = port
        self.robot_stopped = True
        self.rc_socket = socket.socket()

        self.update_control = False

        self.controlling_allowed = True
        self.controlling_allowed_2 = False
        self.joystick = pygame.joystick.Joystick(self.joy)

        self.u1 = 0
        self.u2 = 0
        try:
            self.rc_socket.connect((self.ip, self.port))
        except socket.error():
            print("couldn't connect to socket")

        self.check_joystick = False
    def joystick_init(self):  # Joystick's initialisation
        joystick_count = pygame.joystick.get_count()
        for count in range(joystick_count):
            if joystick_count == 2:
                self.joystick = pygame.joystick.Joystick(self.joy)
                print("{}-->joystick count".format(joystick_count))
                self.joystick.init()
                self.check_joystick = True
            elif joystick_count == 1:
                joystick = pygame.joystick.Joystick(0)
                joystick.init()
                print("connected only 1 joystick - {}".format(joystick))
                self.check_joystick = False
            elif not joystick_count:
                print("no joysticks connected")
                self.check_joystick = False





    def check_control(self):
        if not self.controlling_allowed:
            if self.robot_stopped:# and self.robot0_stopped_2 and self.robot0_stopped_3 and self.robot0_stopped_4:
                pass
                #print("check control -> already stopped")
            else:
                self.u1 = self.u2 = 0
                self.robot_stopped = True
                # self.robot0_stopped_2 = True
                # self.robot0_stopped_3 = True
                # self.robot0_stopped_4 = True
                self.rc_socket.send('({},{})\n'.format(self.u1, self.u2).encode())

                print("check control -> stop")




    def control(self, event):  # the control of two robots with joysticks
        if self.controlling_allowed_2:
            if self.controlling_allowed:
                joy = event.joy
                value = event.value*0.75
                axis = event.axis

                if joy == self.joy:
                    if axis == 1:
                        if abs(value) > 0.2:
                            u1 = u2 = -value
                            self.rc_socket.send('({},{})\n'.format(u1, u2).encode())
                            self.robot_stopped = False

                        elif not self.robot_stopped:
                            u1 = u2 = 0
                            self.rc_socket.send('({},{})\n'.format(u1, u2).encode())
                            self.robot_stopped = True
                    elif axis == 3:
                        if abs(value) > 0.2:
                            u1 = value
                            u2 = -value
                            self.rc_socket.send('({},{})\n'.format(u1, u2).encode())
                            self.robot0_stopped_2 = False

                        elif not self.robot0_stopped_2:
                            u1 = u2 = 0
                            self.rc_socket.send('({},{})\n'.format(u1, u2).encode())
                            self.robot0_stopped_2 = True
                    elif axis == 2:
                        if value > 0.2:
                            u1 = value/1.9
                            u2 = value/1.2
                            self.rc_socket.send('({},{})\n'.format(u1, u2).encode())
                            self.robot0_stopped_3 = False

                        elif not self.robot0_stopped_3:
                            u1 = u2 = 0
                            self.rc_socket.send('({},{})\n'.format(u1, u2).encode())
                            self.robot0_stopped_3 = True
                    elif axis == 5:
                        if value > 0.2:
                            u1 = value/1.2
                            u2 = value/1.9
                            self.rc_socket.send('({},{})\n'.format(u1, u2).encode())
                            self.robot0_stopped_4 = False

                        elif not self.robot0_stopped_4:
                            u1 = u2 = 0
                            self.rc_socket.send('({},{})\n'.format(u1, u2).encode())
                            self.robot0_stopped_4 = True

            else:  # not self.robot_stopped and not self.robot0_stopped_2
                if self.robot_stopped and self.robot0_stopped_2 and self.robot0_stopped_3 and self.robot0_stopped_4:
                    pass
                    print("check control -> already stopped")
                else:
                    u1 = u2 = 0
                    self.robot_stopped = True
                    self.robot0_stopped_2 = True
                    self.robot0_stopped_3 = True
                    self.robot0_stopped_4 = True
                    self.rc_socket.send('({},{})\n'.format(u1, u2).encode())

    def control_alt(self):
        if self.controlling_allowed and self.controlling_allowed_2:
            ax1 = self.joystick.get_axis(1) * 0.75
            ax2 = self.joystick.get_axis(2) * 0.75
            ax3 = self.joystick.get_axis(3) * 0.75
            ax5 = self.joystick.get_axis(5) * 0.75

            u1_old = deepcopy(self.u1)
            u2_old = deepcopy(self.u2)

            threshold = 0.2
            if abs(ax1) < threshold and abs(ax2) < threshold and abs(ax3) < threshold and abs(ax5) < threshold and not self.robot_stopped:
                self.robot_stopped = True
                #print("all axis in neutral -> stopping")
                self.u1 = self.u2 = 0
                #self.rc_socket.send('({},{})\n'.format(u1, u2).encode())
            else:
                self.robot_stopped = False
                if abs(ax1) > threshold and abs(ax3) > threshold:
                    self.u1 = (-ax1 + ax3)/2.0
                    self.u2 = (-ax1 - ax3)/2.0
                    #self.rc_socket.send('({},{})\n'.format(u1, u2).encode())
                elif abs(ax1) > threshold:
                    self.u1 = self.u2 = -ax1
                    #self.rc_socket.send('({},{})\n'.format(u1, u2).encode())
                elif abs(ax3) > threshold:
                    self.u1 = ax3
                    self.u2 = -ax3

            if u1_old != self.u1 or u2_old != self.u2:
                #print("update_control = True")
                self.update_control = True


    def send_control(self):
        if self.update_control:
            #print("sending..({},{})\n".format(self.u1, self.u2))
            self.rc_socket.send('({},{})\n'.format(self.u1, self.u2).encode())


    def control_keyboard(self, event):  # keyboard control for robot1

        u1_old = deepcopy(self.u1)
        u2_old = deepcopy(self.u2)

        if pressed[pygame.K_LEFT]:
            self.u1 = -1.0
            self.u2 = 1.0
            #self.rc_socket.send('({},{})\n'.format(u1, u2).encode())

        elif pressed[pygame.K_RIGHT]:
            self.u1 = 1.0
            self.u2 = -1.0
            #self.rc_socket.send('({},{})\n'.format(u1, u2).encode())

        elif pressed[pygame.K_UP]:
            self.u1 = -1.0
            self.u2 = -1.0
            #self.rc_socket.send('({},{})\n'.format(u1, u2).encode())

        elif pressed[pygame.K_DOWN]:
            self.u1 = 1.0
            self.u2 = 1.0
            #self.rc_socket.send('({},{})\n'.format(u1, u2).encode())

        elif event.type == pygame.KEYUP:
            self.u1 = 0.0
            self.u2 = 0.0
            #selfd.rc_socket.send('({},{})\n'.format(self.u1, self.u2).encode())

        if u1_old != self.u1 or u2_old != self.u2:
            print("update_control = True")
            self.update_control = True

    def control_keyboard_2(self, event):  # keyboard control for robot2

        u1_old = deepcopy(self.u1)
        u2_old = deepcopy(self.u2)

        if pressed[pygame.K_a]:
            self.u1 = -1.0
            self.u2 = 1.0
            #self.rc_socket.send('({},{})\n'.format(u1, u2).encode())

        elif pressed[pygame.K_d]:
            self.u1 = 1.0
            self.u2 = -1.0
            #self.rc_socket.send('({},{})\n'.format(u1, u2).encode())

        elif pressed[pygame.K_w]:
            self.u1 = -1.0
            self.u2 = -1.0
            #self.rc_socket.send('({},{})\n'.format(u1, u2).encode())

        elif pressed[pygame.K_s]:
            self.u1 = 1.0
            self.u2 = 1.0
            #self.rc_socket.send('({},{})\n'.format(u1, u2).encode())

        elif event.type == pygame.KEYUP:
            self.u1 = 0.0
            self.u2 = 0.0
            #self.rc_socket.send('({},{})\n'.format(self.u1, self.u2).encode())

        if u1_old != self.u1 or u2_old != self.u2:
            print("update_control = True")
            self.update_control = True


def main():
    pygame.init()
    screen= pygame.display.set_mode((640, 480))
    robot_font = pygame.freetype.SysFont("shut up and play", 32)
    robot_font_2 = pygame.freetype.SysFont("shut up and play", 22)
    clock = pygame.time.Clock()
    rospy.init_node('game_node', anonymous=True)
    robot_1 = robot(0, '192.168.1.102', 1234)
    robot_1.joystick_init()
    robot_2 = robot(1, '192.168.1.101', 1234)
    robot_2.joystick_init()

    robot_track = race_track(robot_1, robot_2)
    time.sleep(1.0)

    while True:
        robot_track.check_number_marker()
        robot_track.polygon_rectangle()
        robot_track.checkpoint()
        robot_track.finish_line()

        robot_1.check_control()
        robot_2.check_control()


        events = pygame.event.get()
        global pressed
        pressed = pygame.key.get_pressed()

        robot_1.update_control = False
        robot_2.update_control = False
        input_1_recorded = False
        input_2_recorded = False
        for event in events:
            if event.type == pygame.JOYBUTTONDOWN:
                robot_track.check_starting_position()
                #robot_track.check_starting_position()
            if event.type == pygame.JOYAXISMOTION:
                if event.joy == robot_2.joy:
                    robot_2.control_alt()
                    input_2_recorded = True

                if event.joy == robot_1.joy:
                    robot_1.control_alt()
                    input_1_recorded = True

            else:
                robot_1.control_keyboard(event)
                robot_2.control_keyboard_2(event)
                input_1_recorded = True
                input_2_recorded = True
        if input_2_recorded:
            robot_2.send_control()
        if input_1_recorded:
            robot_1.send_control()


        screen.fill((0, 0, 0))
        # You can use 'render' and then blit the text surface ...

        # or just 'render_to' the target surface.
        robot_font.render_to(screen, (80, 50), "Become your weapon and fight!", (255, 255, 255))
        robot_font_2.render_to(screen, (280, 300), "PRESS 'B' FOR THE BEGINNING", (255, 130, 70))
        pygame.display.flip()


if __name__ == '__main__':
    main()