'''
Module for implementing a UR controller real-time monitor over socket port 30003.
Confer http://support.universal-robots.com/Technical/RealTimeClientInterface
Note: The packet lenght given in the web-page is 740. What is actually received from the controller is 692. 
It is assumed that the motor currents, the last group of 48 bytes, are not send.
Originally Written by Morten Lind
Parsing for Firmware 5.9 is added by Byeongdu Lee
'''
import logging
import socket
import struct
import time
import threading
from copy import deepcopy

import numpy as np

import math3d as m3d

__author__ = "Morten Lind, Olivier Roulet-Dubonnet"
__copyright__ = "Copyright 2011, NTNU/SINTEF Raufoss Manufacturing AS"
__credits__ = ["Morten Lind, Olivier Roulet-Dubonnet"]
__license__ = "LGPLv3"


class URRTMonitor(threading.Thread):

    # Struct for revision of the UR controller giving 692 bytes
    rtstruct692 = struct.Struct('>d6d6d6d6d6d6d6d6d18d6d6d6dQ')

    # for revision of the UR controller giving 540 byte. Here TCP
    # pose is not included!
    rtstruct540 = struct.Struct('>d6d6d6d6d6d6d6d6d18d')

    rtstruct5_1 = struct.Struct('>d1d6d6d6d6d6d6d6d6d6d6d6d6d6d6d1d6d1d1d1d6d1d6d3d6d1d1d1d1d1d1d1d6d1d1d3d3d')
    rtstruct5_9 = struct.Struct('>d6d6d6d6d6d6d6d6d6d6d6d6d6d6d1d6d1d1d1d6d1d6d3d6d1d1d1d1d1d1d1d6d1d1d3d3d1d')

    def __init__(self, urHost, urFirm=None):
        threading.Thread.__init__(self)
        self.logger = logging.getLogger(self.__class__.__name__)
        self.daemon = True
        self._stop_event = True
        self._dataEvent = threading.Condition()
        self._dataAccess = threading.Lock()
        self._rtSock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self._rtSock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        self._urHost = urHost
        self.urFirm = urFirm
        # Package data variables
        self._timestamp = None
        self._ctrlTimestamp = None
        self._qActual = None
        self._qTarget = None
        self._tcp = None
        self._tcp_force = None
        self._joint_temperature = None
        self._joint_voltage = None
        self._joint_current = None
        self._main_voltage = None
        self._robot_voltage = None
        self._robot_current = None
        self._qdTarget = None
        self._qddTarget = None
        self._iTarget = None
        self._mTarget = None
        self._qdActual = None
        self._tcp_speed = None
        self._tcp = None
        self._robot_mode = None
        self._joint_modes = None
        self._digital_outputs = None
        self._program_state = None
        self._safety_status = None

        self.__recvTime = 0
        self._last_ctrl_ts = 0
        # self._last_ts = 0
        self._buffering = False
        self._buffer_lock = threading.Lock()
        self._buffer = []
        self._csys = None
        self._csys_lock = threading.Lock()

    def set_csys(self, csys):
        with self._csys_lock:
            self._csys = csys

    def __recv_bytes(self, nBytes):
        ''' Facility method for receiving exactly "nBytes" bytes from
        the robot connector socket.'''
        # Record the time of arrival of the first of the stream block
        recvTime = 0
        pkg = b''
        while len(pkg) < nBytes:
            pkg += self._rtSock.recv(nBytes - len(pkg))
            if recvTime == 0:
                recvTime = time.time()
        self.__recvTime = recvTime
        return pkg

    def wait(self):
        with self._dataEvent:
            self._dataEvent.wait()

    def q_actual(self, wait=False, timestamp=False):
        """ Get the actual joint position vector."""
        if wait:
            self.wait()
        with self._dataAccess:
            if timestamp:
                return self._timestamp, self._qActual
            else:
                return self._qActual
    getActual = q_actual

    def qd_actual(self, wait=False, timestamp=False):
        """ Get the actual joint velocity vector."""
        if wait:
            self.wait()
        with self._dataAccess:
            if timestamp:
                return self._timestamp, self._qdActual
            else:
                return self._qdActual

    def q_target(self, wait=False, timestamp=False):
        """ Get the target joint position vector."""
        if wait:
            self.wait()
        with self._dataAccess:
            if timestamp:
                return self._timestamp, self._qTarget
            else:
                return self._qTarget
    getTarget = q_target

    def tcf_pose(self, wait=False, timestamp=False, ctrlTimestamp=False):
        """ Return the tool pose values."""
        if wait:
            self.wait()
        with self._dataAccess:
            tcf = self._tcp
            if ctrlTimestamp or timestamp:
                ret = [tcf]
                if timestamp:
                    ret.insert(-1, self._timestamp)
                if ctrlTimestamp:
                    ret.insert(-1, self._ctrlTimestamp)
                return ret
            else:
                return tcf
    getTCF = tcf_pose

    def tcf_force(self, wait=False, timestamp=False):
        """ Get the tool force. The returned tool force is a
        six-vector of three forces and three moments."""
        if wait:
            self.wait()
        with self._dataAccess:
            # tcf = self._fwkin(self._qActual)
            tcf_force = self._tcp_force
            if timestamp:
                return self._timestamp, tcf_force
            else:
                return tcf_force
    getTCFForce = tcf_force

    def joint_temperature(self, wait=False, timestamp=False):
        """ Get the joint temperature."""
        if wait:
            self.wait()
        with self._dataAccess:
            joint_temperature = self._joint_temperature
            if timestamp:
                return self._timestamp, joint_temperature
            else:
                return joint_temperature
    getJOINTTemperature = joint_temperature

    def joint_voltage(self, wait=False, timestamp=False):
        """ Get the joint voltage."""
        if wait:
            self.wait()
        with self._dataAccess:
            joint_voltage = self._joint_voltage
            if timestamp:
                return self._timestamp, joint_voltage
            else:
                return joint_voltage
    getJOINTVoltage = joint_voltage

    def joint_current(self, wait=False, timestamp=False):
        """ Get the joint current."""
        if wait:
            self.wait()
        with self._dataAccess:
            joint_current = self._joint_current
            if timestamp:
                return self._timestamp, joint_current
            else:
                return joint_current
    getJOINTCurrent = joint_current

    def main_voltage(self, wait=False, timestamp=False):
        """ Get the Safety Control Board: Main voltage."""
        if wait:
            self.wait()
        with self._dataAccess:
            main_voltage = self._main_voltage
            if timestamp:
                return self._timestamp, main_voltage
            else:
                return main_voltage
    getMAINVoltage = main_voltage

    def robot_voltage(self, wait=False, timestamp=False):
        """ Get the Safety Control Board: Robot voltage (48V)."""
        if wait:
            self.wait()
        with self._dataAccess:
            robot_voltage = self._robot_voltage
            if timestamp:
                return self._timestamp, robot_voltage
            else:
                return robot_voltage
    getROBOTVoltage = robot_voltage

    def robot_current(self, wait=False, timestamp=False):
        """ Get the Safety Control Board: Robot current."""
        if wait:
            self.wait()
        with self._dataAccess:
            robot_current = self._robot_current
            if timestamp:
                return self._timestamp, robot_current
            else:
                return robot_current
    getROBOTCurrent = robot_current

    def __recv_rt_data(self):
        head = self.__recv_bytes(4)
        # Record the timestamp for this logical package
        timestamp = self.__recvTime
        pkgsize = struct.unpack('>i', head)[0]
        self.logger.debug(
            'Received header telling that package is %s bytes long',
            pkgsize)
        payload = self.__recv_bytes(pkgsize - 4)
        if self.urFirm is not None:
            if self.urFirm == 5.1:
                unp = self.rtstruct5_1.unpack(payload[:self.rtstruct5_1.size])
            if self.urFirm == 5.9:
                unp = self.rtstruct5_9.unpack(payload[:self.rtstruct5_9.size])
        else:
            if pkgsize >= 692:
                unp = self.rtstruct692.unpack(payload[:self.rtstruct692.size])
            elif pkgsize >= 540:
                unp = self.rtstruct540.unpack(payload[:self.rtstruct540.size])
            else:
                self.logger.warning(
                    'Error, Received packet of length smaller than 540: %s ',
                    pkgsize)
                return
        

        with self._dataAccess:
            self._timestamp = timestamp
            # it seems that packet often arrives packed as two... maybe TCP_NODELAY is not set on UR controller??
            # if (self._timestamp - self._last_ts) > 0.010:
            # self.logger.warning("Error the we did not receive a packet for {}s ".format( self._timestamp - self._last_ts))
            # self._last_ts = self._timestamp
            self._ctrlTimestamp = np.array(unp[0])
            if self._last_ctrl_ts != 0 and (
                    self._ctrlTimestamp -
                    self._last_ctrl_ts) > 0.010:
                self.logger.warning(
                    "Error the controller failed to send us a packet: time since last packet %s s ",
                    self._ctrlTimestamp - self._last_ctrl_ts)
            self._last_ctrl_ts = self._ctrlTimestamp
            self._qActual = np.array(unp[31:37])
            self._qdActual = np.array(unp[37:43])
            self._qTarget = np.array(unp[1:7])
            self._tcp_force = np.array(unp[67:73])
            self._tcp = np.array(unp[73:79])            
            self._joint_current = np.array(unp[43:49])
            if self.urFirm >= 3.1:
                self._joint_temperature = np.array(unp[86:92])
                self._joint_voltage = np.array(unp[124:130])
                self._main_voltage = unp[121]
                self._robot_voltage = unp[122]
                self._robot_current = unp[123]

            if self.urFirm>= 5.9:
                self._qdTarget = np.array(unp[7:13])
                self._qddTarget = np.array(unp[13:19])
                self._iTarget = np.array(unp[19:25])
                self._mTarget = np.array(unp[25:31])
                self._tcp_speed = np.array(unp[61:67])
                self._joint_current = np.array(unp[49:55])
                self._joint_voltage = np.array(unp[124:130])
                self._robot_mode = unp[94]
                self._joint_modes = np.array(unp[95:101])
                self._digital_outputs = unp[130]
                self._program_state = unp[131]
                self._safety_status = unp[138]

            if self._csys:
                with self._csys_lock:
                    # might be a godd idea to remove dependancy on m3d
                    tcp = self._csys.inverse * m3d.Transform(self._tcp)
                self._tcp = tcp.pose_vector
        if self._buffering:
            with self._buffer_lock:
                self._buffer.append(
                    (self._timestamp,
                     self._ctrlTimestamp,
                     self._tcp,
                     self._qActual))  # FIXME use named arrays of allow to configure what data to buffer

        with self._dataEvent:
            self._dataEvent.notifyAll()

    def start_buffering(self):
        """
        start buffering all data from controller
        """
        self._buffer = []
        self._buffering = True

    def stop_buffering(self):
        self._buffering = False

    def try_pop_buffer(self):
        """
        return oldest value in buffer
        """
        with self._buffer_lock:
            if len(self._buffer) > 0:
                return self._buffer.pop(0)
            else:
                return None

    def pop_buffer(self):
        """
        return oldest value in buffer
        """
        while True:
            with self._buffer_lock:
                if len(self._buffer) > 0:
                    return self._buffer.pop(0)
            time.sleep(0.001)

    def get_buffer(self):
        """
        return a copy of the entire buffer
        """
        with self._buffer_lock:
            return deepcopy(self._buffer)

    def get_all_data(self, wait=True):
        """
        return all data parsed from robot as a dict
        """
        if wait:
            self.wait()
        with self._dataAccess:
            return dict(
                timestamp=self._timestamp,
                ctrltimestamp=self._ctrlTimestamp,
                qActual=self._qActual,
                qTarget=self._qTarget,
                qdActual=self._qdActual,
                qdTarget=self._qdTarget,
                tcp=self._tcp,
                tcp_force=self._tcp_force,
                tcp_speed=self._tcp_speed,
                joint_temperature=self._joint_temperature,
                joint_voltage=self._joint_voltage,
                joint_current=self._joint_current,
                joint_modes=self._joint_modes,
                robot_modes=self._robot_mode,
                main_voltage=self._main_voltage,
                robot_voltage=self._robot_voltage,
                robot_current=self._robot_current,
                digital_outputs=self._digital_outputs,
                program_state=self._program_state,
                safety_status=self._safety_status)
    getALLData = get_all_data

    def stop(self):
        # print(self.__class__.__name__+': Stopping')
        self._stop_event = True

    def close(self):
        self.stop()
        self.join()

    def run(self):
        self._stop_event = False
        self._rtSock.connect((self._urHost, 30003))
        while not self._stop_event:
            self.__recv_rt_data()
        self._rtSock.close()