import socket
import pygame
import json
from math import sin,cos,atan2,pi
from argparse import ArgumentParser

parser = ArgumentParser()
parser.add_argument('bot', metavar='bot', type=str, help='ip address of the controlled robot')
parser.add_argument('id', metavar='id', type=str, help='id of the controlled robot')
parser.add_argument('meas', metavar='meas', type=str, help='ip address of the measurement server')
args = parser.parse_args()

bot = args.bot
meas = args.meas


pygame.init()
pygame.display.set_mode((640, 480))

meas_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
rc_socket = socket.socket()
try:
    rc_socket.connect((bot, 1234)) # connect to robot
except socket.error:
    print("could not connect to bot socket")

meas_socket.sendto(f"{args.id}\n".encode(), (meas, 42424))

class Bot(object):
    FAR, CLOSE, ANGLE, REACHED = range(4)
    def __init__(self, control):
        self.sock = control
        self.state = Bot.REACHED
        self.tx = self.ty = None
        self.ta = None
    
    def move(self, x, y, angle = None):
        self.tx = x
        self.ty = y
        self.ta = angle
        self.pps = True
        self.pas = True

        self.state = Bot.FAR

    def rotate(self, angle):
        self.ta = angle
        self.state = Bot.ANGLE

    def stop(self):
        self.sock.send(f'(0,0)\n'.encode())
    def step(self, x, y, angle):
        umax = 1.0

        if self.tx != None:
            dp = ((self.tx - x)**2 + (self.ty - y)**2)**0.5
            da = atan2(self.ty - y, self.tx - x) - angle
        
        if da > pi:
            da -= 2*pi
        elif da < -pi:
            da += 2*pi

        # Position state: Determine if the angle to the target flips over +-pi/2, i.e. going past the target now
        def pst(angle):
            return abs(angle) > pi
        # Angle state: When 'close' to target angle (absolute value <= pi/2) use the sign, else 0. When the angle flips over 0, the difference is +- 2
        def ast(angle):
            return 0 if abs(angle) >= pi/2 else 1 if angle >= 0 else -1
        # Move with speed and change towards angle
        def control(speed, angle):
            speed = speed * cos(angle)
            ul, ur = speed, speed
            ul -= angle * 0.5
            ur += angle * 0.5
            vd = max(1, abs(ul)/umax, abs(ur)/umax)
            ul /= vd
            ur /= vd
            self.sock.send(f'({ul},{ur})\n'.encode())
            self.pps = pst(angle)
            self.pas = ast(angle)

        if self.state == Bot.FAR:
            if dp <= 0.3:
                # Close to the target
                self.state = Bot.NEAR
            else:
                # When far, just move towards the target
                control(umax, da)
        if self.state == Bot.NEAR:
            if pst(da) != self.pps:
                # The angle flips over +- pi/2, go for the angle setpoint
                self.state = Bot.ANGLE
            else:
                control(umax * dp / 0.3, da)
        if self.state == Bot.ANGLE:
            if self.ta == None:
                # No angle setpoint, we're done
                self.state = Bot.REACHED
            else:
                # Use angle setpoint for angle difference
                da = self.ta - angle
                # Difference between angle states:
                # +- 1 means flipping between 'closer than pi/2' and 'further than pi/2'
                # +- 2 means abs(angle) < pi/2 and the sign changes -> In position!
                if abs(ast(da) - self.pas) == 2:
                    # Done
                    self.state = Bot.REACHED
                else:
                    # Don't move, just rotate
                    control(0, da)

bot = Bot(rc_socket)
bot.move(0, 0, 0)

running = True
ml = 0.5
while running:
    received = json.loads(str(meas_socket.recv(1024), "utf-8").strip())
    print("Received: {}".format(received))

    # Extract position, angle
    x = received['x']
    y = received['y']
    a = received['angle']

    events = pygame.event.get()
    for event in events:
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_RIGHT:
                bot.move(ml, 0, pi)
            if event.key == pygame.K_LEFT:
                bot.move(-ml, 0, 0)
            if event.key == pygame.K_UP:
                bot.move(0, ml, pi*3/2)
            if event.key == pygame.K_DOWN:
                bot.move(0, -ml, pi/2)
            if event.key == pygame.K_ESCAPE:
                bot.stop()
                running = False
    if running:
        bot.step(x, y, a)