from helpfunctions import *
import time

class UR:
        
    import rtde_receive
    import rtde_control
    
    def __init__(self,_joint_min,_joint_max,_frequency_rtde_send=500,_frequency_rtde_receive=500):
        self.frequency_rtde_send=_frequency_rtde_send
        self.frequency_rtde_receive=_frequency_rtde_receive

        self.controlEnabled=False
        
        self.last_receivertde=0
        self.last_sendrtde=0

        self.joint_min=_joint_min
        self.joint_max=_joint_max
                
        self.rtde_r=None
        self.rtde_c=None

            
        self.last_forwardkinematics=None
        self.last_forwardkinematics_joint_q_received=None
        

                
        # Parameters
        self.velocity = 0.1
        self.acceleration = 0.1
        self.lookahead_time = 0.1
        self.gain = 300
        self.joint_q_aim=[0,0,0,0,0,0] #joint rotations aiming for
        self.joint_q = [0,0,0,0,0,0] #joint rotation last sent
        self.velocitylimit=[1,1,1,3,3,3] #joint velocity limits
        #self.velocitylimit=[0.2,0.2,0.2,0.2,0.2,0.2] #joint velocity limits


        self.joint_q_received=[0,0,0,0,0,0]
        self.tcppose_received=[0,0,0,0,0,0]

    def connect(self,robotip):
        
        self.rtde_r = self.rtde_receive.RTDEReceiveInterface(robotip, self.frequency_rtde_receive)

        self.rtde_c = self.rtde_control.RTDEControlInterface(robotip)

    def setJointQAim(self,q):
        self.joint_q_aim=q

    def isJointsWithinSafetyLimits(self,q):
        return self.rtde_c.isJointsWithinSafetyLimits(q)

    def enableControl(self):
        if time.time()-self.last_receivertde<0.1: #if received position is recent
            self.joint_q=self.joint_q_received #update current position to received position
            self.joint_q_aim=self.joint_q_received
            self.controlEnabled=True
    def disableControl(self):
        self.controlEnabled=False
        
    def isControlEnabled(self):
        return self.controlEnabled
    
    def send(self,looptime):
    
        if looptime-self.last_sendrtde<1.0/self.frequency_rtde_send:
            return
        self.last_sendrtde=looptime

        if self.controlEnabled:
            for i in range(6): #update joint_q with max speeds
                self.joint_q[i] += constrain(self.joint_q_aim[i]-self.joint_q[i], -self.velocitylimit[i]/self.frequency_rtde_send,self.velocitylimit[i]/self.frequency_rtde_send) #constrain joint speeds

            if self.rtde_c is not None:
                self.rtde_c.servoJ(self.joint_q, self.velocity, self.acceleration, 1.0/self.frequency_rtde_send, self.lookahead_time, self.gain)

    def receiveJoints(self,looptime):
        
        #Get robot joint angles
        if looptime-self.last_receivertde<1.0/self.frequency_rtde_receive:
            return
        self.last_receivertde=looptime

        if self.rtde_r is not None:
            self.joint_q_received=self.rtde_r.getActualQ()
            self.tcppose_received=self.rtde_r.getActualTCPPose()
        else:
            _offsets=[12.2,5.3,0,6,12.4,612.91]
            _factors=[1.2,1.14,1.214,0.9921,1,1.092]
            
            for _i,_offset in enumerate(_offsets): #generate mockup values
                self.joint_q_received[_i]=(math.sin(time.time()*_factors[_i]+_offset)/2.0+0.5)*(self.joint_max[_i]-self.joint_min[_i])+self.joint_min[_i]

    def receiveStatus(self,looptime):
        
        if looptime-self.last_receivertde<1.0/self.frequency_rtde_receive:
            return
        self.last_receivertde=looptime
        
        self.status_connected=self.rtde_r.isConnected()
        self.status_timestamp=self.rtde_r.getTimestamp()
        self.status_protectivestopped=self.rtde_r.isProtectiveStopped()
        self.status_emergencystopped=self.rtde_r.isEmergencyStopped()
        self.status_robotstatus=self.rtde_c.getRobotStatus()
        '''
        self.status_=self.rtde_r.()
        self.status_=self.rtde_r.()
        self.status_=self.rtde_r.()
        self.status_=self.rtde_r.()
        self.status_=self.rtde_r.()
        self.status_=self.rtde_r.()
        self.status_=self.rtde_r.()
        '''

    def getReceivedQ(self):
        return self.joint_q_received
    
    '''def getForwardKinematics(self):
        if self.last_forwardkinematics_joint_q_received!=self.joint_q_received:
            self.last_forwardkinematics_joint_q_received=self.joint_q_received
            self.last_forwardkinematics=self.rtde_c.getForwardKinematics(self.joint_q_received)

        return self.last_forwardkinematics
        '''
    
    def getReceivedTCPPose(self):
        return self.tcppose_received
    
    def getInverseKinematics(self,pos_q):
        return self.rtde_c.getInverseKinematics(pos_q,self.joint_q)

    def disconnect(self):
        if self.rtde_c is not None:
            self.rtde_c.servoStop()
            self.rtde_c.stopScript()