path: root/src/argaze/TobiiGlassesPro2/TobiiData.py

#!/usr/bin/env python

from dataclasses import dataclass
import threading
import uuid
import gzip
import json
import time
import queue

from argaze import DataStructures
from argaze.TobiiGlassesPro2 import TobiiNetworkInterface

from argaze.utils import MiscFeatures

import numpy

@dataclass
class DirSig():
    """Define dir sig data (dir sig)."""

    dir: int # meaning ?
    sig: int # meaning ?

@dataclass
class PresentationTimeStamp():
    """Define presentation time stamp (pts) data."""

    value: int
    """Pts value."""

@dataclass
class VideoTimeStamp():
    """Define video time stamp (vts) data."""

    value: int
    """Vts value."""

@dataclass
class EventSynch():
    """Define event synch (evts) data."""

    value: int # meaning ?
    """Evts value."""

@dataclass
class Event():
    """Define event data (ets type tag)."""

    ets: int # meaning ?
    type: str
    tag: str # dict ?

@dataclass
class Accelerometer():
    """Define accelerometer data (ac)."""

    value: numpy.array
    """Accelerometer value"""

@dataclass
class Gyroscope():
    """Define gyroscope data (gy)."""

    value: numpy.array
    """Gyroscope value"""

@dataclass
class PupilCenter():
    """Define pupil center data (gidx pc eye)."""

    validity: int
    index: int
    value: tuple((float, float, float))
    eye: str # 'right' or 'left'

@dataclass
class PupilDiameter():
    """Define pupil diameter data (gidx pd eye)."""

    validity: int
    index: int
    value: float
    eye: str # 'right' or 'left'

@dataclass
class GazeDirection():
    """Define gaze direction data (gidx gd eye)."""

    validity: int
    index: int
    value: tuple((float, float, float))
    eye: str # 'right' or 'left'

@dataclass
class GazePosition():
    """Define gaze position data (gidx l gp)."""

    validity: int
    index: int
    l: str # ?
    value: tuple((float, float))

@dataclass
class GazePosition3D():
    """Define gaze position 3D data (gidx gp3)."""

    validity: int
    index: int
    value: tuple((float, float))

@dataclass
class MarkerPosition():
    """Define marker data (marker3d marker2d)."""

    value_3d: tuple((float, float, float))
    value_2d: tuple((float, float))

class __TobiiJsonDataParser():

    def parse_dir_sig(self, status, json_data):

        return DirSig(json_data['dir'], json_data['sig'])

    def parse_pts(self, status, json_data):

        return PresentationTimeStamp(json_data['pts'])

    def parse_vts(self, status, json_data):

        return VideoTimeStamp(json_data['vts'])

    def parse_event_synch(self, status, json_data):

        return EventSynch(json_data['evts'])

    def parse_event(self, status, json_data):

        return Event(json_data['ets'], json_data['type'], json_data['tag'])

    def parse_accelerometer(self, status, json_data):

        return Accelerometer(json_data['ac'])

    def parse_gyroscope(self, status, json_data):

        return Gyroscope(json_data['gy'])

    def parse_pupil_center(self, status, gaze_index, json_data):

        return PupilCenter(status, gaze_index, json_data['pc'], json_data['eye'])

    def parse_pupil_diameter(self, status, gaze_index, json_data):

        return PupilDiameter(status, gaze_index, json_data['pd'], json_data['eye'])

    def parse_gaze_direction(self, status, gaze_index, json_data):

        return GazeDirection(status, gaze_index, json_data['gd'], json_data['eye'])

    def parse_gaze_position(self, status, gaze_index, json_data):

        return GazePosition(status, gaze_index, json_data['l'], json_data['gp'])

    def parse_gaze_position_3d(self, status, gaze_index, json_data):

        return GazePosition3D(status, gaze_index, json_data['gp3'])

    def parse_marker_position(self, status, json_data):

        return MarkerPosition(json_data['marker3d'], json_data['marker2d'])

    def parse_pupil_or_gaze(self, status, json_data):

        gaze_index = json_data.pop('gidx')

        # parse pupil or gaze data depending second json key
        second_key = next(iter(json_data))

        parse_map = {
            'pc': self.parse_pupil_center,
            'pd': self.parse_pupil_diameter,
            'gd': self.parse_gaze_direction,
            'l': self.parse_gaze_position,
            'gp3': self.parse_gaze_position_3d
        }

        return parse_map[second_key](status, gaze_index, json_data)

    def parse_data(self, status, json_data):

        # parse data depending first json key
        first_key = next(iter(json_data))

        parse_map = {
            'dir': self.parse_dir_sig,
            'pts': self.parse_pts,
            'vts': self.parse_vts,
            'evts': self.parse_event_synch,
            'ets': self.parse_event,
            'ac': self.parse_accelerometer,
            'gy': self.parse_gyroscope,
            'gidx': self.parse_pupil_or_gaze,
            'marker3d': self.parse_marker_position
        }

        return parse_map[first_key](status, json_data)

class TobiiDataSegment():
    """Handle Tobii Glasses Pro 2 segment data file."""

    def __init__(self, segment_data_path, start_timestamp = 0, end_timestamp = None):
        """Load segment data from segment directory then parse and register each recorded dataflow as a TimeStampedBuffer member of the TobiiSegmentData instance."""

        self.__segment_data_path = segment_data_path

        self.__vts_offset = 0
        self.__vts_ts = -1

        self.__json_data_parser = __TobiiJsonDataParser()

        self.__ts_data_buffer_dict = {
            'DirSig': DataStructures.TimeStampedBuffer(),
            'PresentationTimeStamp': DataStructures.TimeStampedBuffer(),
            'EventSynch': DataStructures.TimeStampedBuffer(),
            'Event': DataStructures.TimeStampedBuffer(),
            'Accelerometer': DataStructures.TimeStampedBuffer(),
            'Gyroscope': DataStructures.TimeStampedBuffer(),
            'PupilCenter': DataStructures.TimeStampedBuffer(),
            'PupilDiameter': DataStructures.TimeStampedBuffer(),
            'GazeDirection': DataStructures.TimeStampedBuffer(),
            'GazePosition': DataStructures.TimeStampedBuffer(),
            'GazePosition3D': DataStructures.TimeStampedBuffer(),
            'MarkerPosition': DataStructures.TimeStampedBuffer()
        }

        # define a decoder function
        def decode(json_data):

            # parse data status
            status = json_data.pop('s', -1)

            # convert timestamp
            ts = json_data.pop('ts')

            # watch for vts data to offset timestamps
            try:
                self.__vts_offset = json_data['vts']
                self.__vts_ts = ts

                return True # continue

            except KeyError:
                pass

            # ignore data before first vts entry
            if self.__vts_ts == -1:
                return True # continue

            ts -= self.__vts_ts
            ts += self.__vts_offset

            # ignore timestamps out of the given time range
            if ts < start_timestamp:
                return True # continue

            if ts >= end_timestamp:
                return False # stop

            # convert json data into data object
            data_object = self.__json_data_parser.parse_data(status, json_data)
            data_object_type = type(data_object).__name__

            # store data object into dedicated timestamped buffer
            self.__ts_data_buffer_dict[data_object_type][ts] = data_object

            return True # continue

        # start loading
        with gzip.open(self.__segment_data_path) as f:

            for item in f:
                if not json.loads(item.decode('utf-8'), object_hook=decode):
                    break

    def __getitem__(self, key):
        return self.__ts_data_buffer_dict[key]

    def keys(self):
        """Get all registered data keys.
        **Returns:** dict_keys"""

        return list(self.__ts_data_buffer_dict.keys())

    def get_path(self):
        """Get segment data path.
        **Returns:** str"""

        return self.__segment_data_path

class TobiiDataStream():
    """Capture Tobii Glasses Pro 2 data stream in separate thread."""

    reading_callback = None
    """Reading callback function to get incoming (data_ts, data_object, data_object_type) back."""

    def __init__(self, network_interface: TobiiNetworkInterface.TobiiNetworkInterface):
        """Initialise network connection and prepare data reception."""

        # Enable network connection
        self.__network = network_interface
        self.__data_socket = self.__network.make_socket()

        # Data reception
        self.__data_thread = None
        self.__data_queue = queue.Queue(50) # TODO : set queue size according technical reason
        self.__json_data_parser = __TobiiJsonDataParser()
        self.__first_ts = 0

        # Data capture
        self.__data_stream_selector = ''
        self.__data_ts_buffer = None
        self.__data_ts_buffer_size = 0

        # Keep connection alive
        self.__keep_alive_msg = "{\"type\": \"live.data.unicast\", \"key\": \""+ str(uuid.uuid4()) +"\", \"op\": \"start\"}"
        self.__keep_alive_thread = threading.Thread(target = self.__keep_alive)
        self.__keep_alive_thread.daemon = True
        self.__keep_alive_thread.start()

    def __del__(self):
        """Stop data reception and network connection before destruction."""

        if self.__data_thread != None:

            threading.Thread.join(self.__data_thread)
            self.__data_thread = None

        threading.Thread.join(self.__keep_alive_thread)

        self.__data_socket.close()

    def __keep_alive(self):
        """Maintain network connection and clear socket when is not read."""

        while True:

            # if no thread is reading data socket
            if self.__data_thread == None:

                self.__network.send_keep_alive_msg(self.__data_socket, self.__keep_alive_msg)

                clear_count = 0
                while clear_count < 1000 and self.__data_thread == None:

                    # Clear socket each milli second
                    time.sleep(0.001)
                    self.__network.grab_data(self.__data_socket)
                    clear_count += 1

            else:

                self.__network.send_keep_alive_msg(self.__data_socket, self.__keep_alive_msg)
                time.sleep(1)

    def open(self):
        """Start data reception."""

        if self.__data_thread == None:

            self.__first_ts = 0

            self.__data_thread = threading.Thread(target = self.__run)
            self.__data_thread.daemon = True

            self.__stop_event = threading.Event()
            self.__read_lock = threading.Lock() 

            self.__data_thread.start()

    def close(self):
        """Stop data reception definitively."""

        if self.__data_thread != None:

            self.__stop_event.set()

            threading.Thread.join(self.__data_thread)
            self.__data_thread = None

    def __run(self):
        """Managed received data for sync and async reading case.
            - Sync: send data to callback function.
            - Async: store data into a locked queue for further reading."""

        while not self.__stop_event.isSet():

            # Sync reading case
            if self.reading_callback != None:

                # grab data
                data = self.__network.grab_data(self.__data_socket)

                # decode data
                json_data = json.loads(data.decode('utf-8'))

                # parse json into timestamped data object
                data_ts, data_object, data_object_type = self.__parse_json_data(json_data)

                self.reading_callback(data_ts, data_object, data_object_type)

            # Async reading case
            else:

                # wait for queue reading
                if self.__data_queue.full():

                    # sleep 1 micro second
                    time.sleep(0.0001)
                    continue

                # lock data queue access
                self.__read_lock.acquire()

                # grab data
                data = self.__network.grab_data(self.__data_socket)

                # decode data
                json_data = json.loads(data.decode('utf-8'))

                # write data in queue
                self.__data_queue.put(json_data)

                # unlock data queue access
                self.__read_lock.release()

    def read(self):
        """Read incoming timestamped data.
        **Returns:** int, dict"""

        # no data to read
        if self.__data_queue.empty():
            raise ValueError

        # create a dictionary of timestamped data buffers
        ts_data_buffer_dict = {
            'DirSig': DataStructures.TimeStampedBuffer(),
            'PresentationTimeStamp': DataStructures.TimeStampedBuffer(),
            'VideoTimeStamp': DataStructures.TimeStampedBuffer(),
            'EventSynch': DataStructures.TimeStampedBuffer(),
            'Event': DataStructures.TimeStampedBuffer(),
            'Accelerometer': DataStructures.TimeStampedBuffer(),
            'Gyroscope': DataStructures.TimeStampedBuffer(),
            'PupilCenter': DataStructures.TimeStampedBuffer(),
            'PupilDiameter': DataStructures.TimeStampedBuffer(),
            'GazeDirection': DataStructures.TimeStampedBuffer(),
            'GazePosition': DataStructures.TimeStampedBuffer(),
            'GazePosition3D': DataStructures.TimeStampedBuffer(),
            'MarkerPosition': DataStructures.TimeStampedBuffer()
        }

        # if the data acquisition thread is not running
        if self.__stop_event.isSet():
            return ts_data_buffer_dict
        
        # lock data queue access
        self.__read_lock.acquire()

        # search for earliest timestamp in the queue
        earliest_ts = 0
   
        # read data queue
        while not self.__data_queue.empty():

            # parse json into timestamped data object
            data_ts, data_object, data_object_type = self.__parse_json_data(self.__data_queue.get())

            # ignore negative timestamp
            if data_ts < 0:
                break

            # keep earliest timestamp
            if data_ts > earliest_ts:
                earliest_ts = data_ts

            # store data object into dedicated timestamped buffer
            ts_data_buffer_dict[data_object_type][data_ts] = data_object

        # unlock data queue access
        self.__read_lock.release()

        return earliest_ts, ts_data_buffer_dict

    def __parse_json_data(self, json_data):

        # parse data status
        status = json_data.pop('s', -1)

        # convert timestamp
        data_ts = json_data.pop('ts')

        # convert json data into data object
        data_object = self.__json_data_parser.parse_data(status, json_data)
        data_object_type = type(data_object).__name__

        # keep first timestamp to offset all timestamps
        if self.__first_ts == 0:
            self.__first_ts = data_ts
 
        data_ts -= self.__first_ts

        return data_ts, data_object, data_object_type

    def __capture_callback(self, data_ts, data_object, data_object_type):
        """Store incoming data into timestamped buffer"""

        if data_object_type == self.__data_stream_selector:
        
            if len(self.__data_ts_buffer.keys()) < self.__data_ts_buffer_size:

                # update first timestamp if next timestamps are negative
                if data_ts < 0:
                    self.__first_ts += data_ts
                    data_ts = 0

                self.__data_ts_buffer[data_ts] = data_object

    def capture(self, data_ts_buffer, data_object_type = '', sample_number = 500):
        """Capture a data stream into buffer."""

        # Prepare for data acquisition
        self.__data_stream_selector = data_object_type
        self.__data_ts_buffer = data_ts_buffer
        self.__data_ts_buffer_size = sample_number

        # Subscribe to data stream
        memo_reading_callback = self.reading_callback 
        self.reading_callback = self.__capture_callback

        # Start data reception
        self.open()

        # Share data acquisition progress
        buffer_size = 0
        while buffer_size < sample_number:

            time.sleep(0.1)

            buffer_size = len(self.__data_ts_buffer.keys())

            yield buffer_size

        # Stop data reception
        self.close()

        # Unsubscribe to tobii data stream
        self.reading_callback = memo_reading_callback