import rtde_control #pip install ur_rtde  (use python 3.9)
import rtde_receive
import time
import math
import socket
import numpy as np

import pygame

use_rtde=True

pygame.init()
width=500
height=500
screen = pygame.display.set_mode([width, height])
running = True

c_border = (100,100,100)
c_name=(200,200,200)
c_cursor = (255,255,255)
c_text=(255,255,255)

def mapFromTo(x,a,b,c,d):
   y=(x-a)/(b-a)*(d-c)+c
   return y

def text_to_screen(screen, text, x, y, size = 50,
    color = (200, 000, 000)):

    text = str(text)
    font = pygame.font.SysFont(None, size)
    text = font.render(text, True, color)
    screen.blit(text, (x, y))

    
def drawBar(screen,_topleft,_size,_pos,_name,_text):
    pygame.draw.rect(screen, c_border, (_topleft[0],_topleft[1],_size[0],_size[1]))
    text_to_screen(screen,_name,_topleft[0],_topleft[1],_size[1],c_name)
    pygame.draw.line(screen, c_cursor, (_topleft[0]+_pos*_size[0],_topleft[1]),(_topleft[0]+_pos*_size[0],_topleft[1]+_size[1]))
    text_to_screen(screen,_text,_size[0]+_topleft[0],_topleft[1],_size[1],c_text)



# Socket stuff
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
host =""
port =30002
s.bind((host,port))
s.listen(1) # Number of connections
s.setblocking(False)
client = None


# Robot Control Stuff
if use_rtde:
    rtde_c = rtde_control.RTDEControlInterface("192.168.1.101")
    rtde_r = rtde_receive.RTDEReceiveInterface("192.168.1.101")

def toRad(d):
    return d/360*2*math.pi

def toDeg(r):
    return r*360/2/math.pi


def RPYtoVec(roll,pitch,yaw):
    yawMatrix = np.matrix([
    [math.cos(yaw), -math.sin(yaw), 0],
    [math.sin(yaw), math.cos(yaw), 0],
    [0, 0, 1]
    ])

    pitchMatrix = np.matrix([
    [math.cos(pitch), 0, math.sin(pitch)],
    [0, 1, 0],
    [-math.sin(pitch), 0, math.cos(pitch)]
    ])

    rollMatrix = np.matrix([
    [1, 0, 0],
    [0, math.cos(roll), -math.sin(roll)],
    [0, math.sin(roll), math.cos(roll)]
    ])

    R = yawMatrix * pitchMatrix * rollMatrix

    theta = math.acos(((R[0, 0] + R[1, 1] + R[2, 2]) - 1) / 2)
    multi = 1 / (2 * math.sin(theta))

    rx = multi * (R[2, 1] - R[1, 2]) * theta
    ry = multi * (R[0, 2] - R[2, 0]) * theta
    rz = multi * (R[1, 0] - R[0, 1]) * theta

    return rx, ry, rz

paused=True


if use_rtde:
    actual_q = rtde_r.getActualQ() #current joint rotations in radians.
    print("Actual Q:"+str(actual_q))
    #fw_actual_q = rtde_c.getForwardKinematics(actual_q)
    fw_actual_q = rtde_c.getForwardKinematics()
    print("forward Actual Q:"+str(fw_actual_q))
    joint_q=rtde_c.getInverseKinematics(fw_actual_q)
    print("Final Q:"+str(joint_q))

winkellimit=45

maxbasedeviation=60 # in degrees, maximum rotation deviation from 0°
basespeed=1.5
distfromcenter=0.9
basedeviation=0
centerpos=[math.cos(toRad(180))*distfromcenter,math.sin(toRad(180))*distfromcenter,0.91]
#centerpos=[-0.895,-0.177,0.91]
centerrotRPY=[-1.572,0,1.572] #roll, pitch, yaw

pan=0 #-=left, +=right, degrees
tilt=0 #-=down, +=up, degrees

joint_q_init=[0,toRad(-90),toRad(90),toRad(-180),toRad(-90),toRad(0)]   #start position
joint_q=joint_q_init


#rtde_c.moveJ([0,toRad(-90),toRad(90-tilt/2),toRad(-180-tilt/2),toRad(-90-pan),toRad(0)],0.1,0.1)
if use_rtde:
    rtde_c.moveJ(joint_q,0.2,0.1) #move to initial position


# Parameters


velocity = 0.2   #0.5
acceleration = 0.1   #0.5
frequency= 100
dt = 1.0/frequency  # 2ms
lookahead_time = 0.1
gain = 300


speed_1=[0.2,0.2,0.2,0.2,0.2,0.2]
speed_2=[0.5,0.5,0.5,1,1,1]
speed_3=[1,1,1,3,3,3]
speed_4=[2,2,2,4,4,4]

speed=speed_1


joint_min=[toRad(-90),toRad(-90-40),toRad(90-90),toRad(-180-80),toRad(-90-90),toRad(0-75)]
joint_max=[toRad(90),toRad(-90+40),toRad(90+40),toRad(-180+110),toRad(-90+90),toRad(0+75)]


def constrain(v,_min,_max):
    return min(_max,max(_min,v))


receivetimearray=np.zeros(100)
receivetimearray_pos=0
time_last_data_received=time.time()
while running:
    
    start = time.time()



    for event in pygame.event.get():
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_p and paused==False:
                print("paused")
                paused=True
            if event.key == pygame.K_s and paused==True:
                print("resumed")
                paused=False
            if event.key == pygame.K_q:
                print("stop")
                running=False
            if event.key == pygame.K_1:
                speed=speed_1
            if event.key == pygame.K_2:
                speed=speed_2
            if event.key == pygame.K_3:
                speed=speed_3
            if event.key == pygame.K_4:
                speed=speed_4
                
        if event.type == pygame.QUIT:
            running = False

    # Receive Socket stuff
    try:
        client, address = s.accept()
        print("Connected to", address)

    except socket.error:
        pass

    if client is not None:
        try:
            data = client.recv( 1024 ).decode( 'utf-8' )
            if data:
                lastdata=data.split('\n')
                _timediff=time.time()-time_last_data_received
                receivetimearray[receivetimearray_pos]=_timediff
                receivetimearray_pos+=1
                receivetimearray_pos%=len(receivetimearray)
                time_last_data_received=time.time()
                #print("Received :", repr(data))
                if (len(lastdata)>1):
                    splitdata=lastdata[-2].split(",")
                    if len(splitdata)==2:
                        try:
                            _pan=float(splitdata[0])
                            _tilt=float(splitdata[1])

                            
                            if _pan>=-winkellimit and _pan <=winkellimit and _tilt>=-winkellimit and _tilt <=winkellimit:
                                pan=_pan
                                tilt=_tilt

                            #print("Pan="+str(pan)+" Tilt="+str(tilt))
                        except ValueError:
                            print("Not a float")
        except socket.error:
            pass


    # Robot Control stuff
    if use_rtde:
        if not paused:
            basedeviation=min(maxbasedeviation,max(-maxbasedeviation,basedeviation+pan*basespeed/frequency))
            basedeviation*=-winkellimit*basespeed/frequency/maxbasedeviation+1
            basedeviation=min(maxbasedeviation,max(-maxbasedeviation,basedeviation))
            

        baserotation=toRad(180-basedeviation) #positive is ccw
        centerpos=[math.cos(baserotation)*distfromcenter,math.sin(baserotation)*distfromcenter,0.91]
        _rot=RPYtoVec(centerrotRPY[0]+toRad(tilt),centerrotRPY[1],centerrotRPY[2]+toRad(-pan))
        pos_q=[centerpos[0],centerpos[1],centerpos[2],_rot[0],_rot[1],_rot[2]]
        joint_q_aim=rtde_c.getInverseKinematics(pos_q,joint_q_init)

        #joint_q_aim = rtde_c.getInverseKinematics([-0.895,-0.173,0.911,2.4,2.4,2.4])
        # inverse kinematics: rx+ =look up
        # ry+ = look right
        # rz+ = rotate CW

        
        reachable=rtde_c.isPoseWithinSafetyLimits(pos_q)
        if not reachable:
            print("Reachable:"+str(reachable))


        insideJointConstraints=True
        for i in range(6):
            if joint_q_aim[i]<joint_min[i] or joint_q_aim[i]>joint_max[i]:
                print("Joint "+str(i)+" Outside constraint with "+str(joint_q_aim[i])+" ("+str(toDeg(joint_q_aim[i]))+")")
                insideJointConstraints=False


        if reachable and insideJointConstraints and not paused:
            for i in range(6):
                joint_q[i] += constrain(joint_q_aim[i]-joint_q[i], -speed[i]/frequency,speed[i]/frequency) #constrain joint speeds
            
        
        #print(joint_q)

        
        if not paused:
            rtde_c.servoJ(joint_q, velocity, acceleration, dt, lookahead_time, gain)

        

    screen.fill((0,0,0))
    #pygame.draw.circle(screen, (0, 0, 255), (250, 250), 75)
    jointnames=["Base","Shoulder","Elbow","Wrist1","Wrist2","Wrist3"]

    
    drawBar(screen,(10,10+30*0),(100,20),mapFromTo(pan,-winkellimit,winkellimit,0.0,1.0),"Pan",str(round(pan,2))+"°")
    drawBar(screen,(10,10+30*1),(100,20),mapFromTo(tilt,-winkellimit,winkellimit,0.0,1.0),"Tilt",str(round(tilt,2))+"°")

    
    for i in range(6):
        drawBar(screen,(10,10+30*(3+i)),(100,20),mapFromTo(joint_q[i],joint_min[i],joint_max[i],0.0,1.0),jointnames[i],str(round(toDeg(joint_q[i]),2))+"°")

    for i in range(6):
        _jointdiff = joint_q_aim[i]-joint_q[i]
        drawBar(screen,(200,10+30*(3+i)),(100,20),mapFromTo(_jointdiff,-speed[i]/frequency,speed[i]/frequency,0.0,1.0),"diff_"+str(jointnames[i]),str(round(toDeg(_jointdiff),3))+"°")



    drawBar(screen,(10,10+30*10),(100,20),mapFromTo(basedeviation,-maxbasedeviation,maxbasedeviation,0.0,1.0),"basedeviation",str(round(basedeviation,2))+"°")

    drawBar(screen,(200,10+30*10),(100,20),mapFromTo(np.mean(receivetimearray),0,1.0,0.0,1.0),"Receive Interval",str(round(np.mean(receivetimearray),3))+"s")
    drawBar(screen,(200,10+30*11),(100,20),mapFromTo(np.max(receivetimearray),0,1.0,0.0,1.0),"Max Receive Int.",str(round(np.max(receivetimearray),3))+"s")


    text_to_screen(screen,"paused: "+str(paused),10,height-20-12*2-20*3,20,c_text)
    text_to_screen(screen,"speed: "+str(speed),10,height-20-12*4-20,20,c_text)
    text_to_screen(screen,"1: "+str(speed_1),10,height-20-12*3,12,c_text)
    text_to_screen(screen,"2: "+str(speed_2),10,height-20-12*2,12,c_text)
    text_to_screen(screen,"3: "+str(speed_3),10,height-20-12*1,12,c_text)
    text_to_screen(screen,"4: "+str(speed_4),10,height-20-12*0,12,c_text)

    text_to_screen(screen,"Press q to stop",250,height-20-20*3,20,c_text)
    text_to_screen(screen,"Press p to pause",250,height-20-20*2,20,c_text)
    text_to_screen(screen,"Press 1,2,3,4 to change speed",250,height-20-20*1,20,c_text)
    text_to_screen(screen,"Press s to resume",250,height-20-20*0,20,c_text)
    
    
    
        
    
    pygame.display.flip()
    
    end = time.time()
    stepduration = end - start
    
    if stepduration < dt:
        time.sleep(dt - stepduration)

print("Finished")

if use_rtde:
    rtde_c.servoStop()
    rtde_c.stopScript()





s.close()