import socket
import pygame
import pygame.joystick


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.robot0_stopped_1 = True
        self.robot0_stopped_2 = True
        self.robot0_stopped_3 = True
        self.robot0_stopped_4 = True
        self.rc_socket = socket.socket()
        try:
            self.rc_socket.connect((self.ip, self.port))
        except socket.error():
            print("couldn't connect to socket")

    def joystick_init(self):  # Joystick's initialisation
        self.joystick = pygame.joystick.Joystick(self.joy)
        self.joystick.init()
        self.axes = self.joystick.get_numaxes()

    def control(self, event):  # the control of two robots with joysticks
        joy = event.joy
        value = event.value
        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.robot0_stopped_1 = False

                elif not self.robot0_stopped_1:
                    u1 = u2 = 0
                    self.rc_socket.send('({},{})\n'.format(u1, u2).encode())
                    self.robot0_stopped_1 = 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

    def control_keyboard(self, event):  # keyboard control for robot1
        command_received = False
        if event.key == pygame.K_LEFT:
            u1 = -1.0
            u2 = 1.0
            command_received = True
        elif event.key == pygame.K_RIGHT:
            u1 = 1.0
            u2 = -1.0
            command_received = True
        elif event.key == pygame.K_UP:
            u1 = -1.0
            u2 = -1.0
            command_received = True
        elif event.key == pygame.K_DOWN:
            u1 = 1.0
            u2 = 1.0
            command_received = True
        if command_received:
            self.rc_socket.send('({},{})\n'.format(u1, u2).encode())

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

        command_received = False
        if event.key == pygame.K_a:
            u1 = -1.0
            u2 = 1.0
            command_received = True
        elif event.key == pygame.K_d:
            u1 = 1.0
            u2 = -1.0
            command_received = True
        elif event.key == pygame.K_w:
            u1 = -1.0
            u2 = -1.0
            command_received = True
        elif event.key == pygame.K_s:
            u1 = 1.0
            u2 = 1.0
            command_received = True
        if command_received:
            self.rc_socket.send('({},{})\n'.format(u1, u2).encode())

    def control_keyboard_stop(self):  # stop for both robot

        u1 = 0
        u2 = 0
        print("key released, resetting: ({},{})".format(u1, u2))
        self.rc_socket.send('({},{})\n'.format(u1, u2).encode())

def main():
    pygame.init()
    pygame.display.set_mode((640, 480))

    robot_1 = robot(0, '192.168.1.102', 1234)
    robot_1.joystick_init()
    robot_2 = robot(1, '192.168.1.103', 1234)
    robot_2.joystick_init()

    while True:
        events = pygame.event.get()
        for event in events:
            if event.type == pygame.JOYAXISMOTION:
                robot_1.control(event)
                robot_2.control(event)
            elif event.type == pygame.KEYDOWN:
                robot_1.control_keyboard(event)
                robot_2.control_keyboard_2(event)
            elif event.type == pygame.KEYUP:
                robot_1.control_keyboard_stop()
                robot_2.control_keyboard_stop()


if __name__ == '__main__':
    main()