path: root/src/argaze/GazeFeatures.py

#!/usr/bin/env python

"""Generic gaze data and class definitions."""

__author__ = "Théo de la Hogue"
__credits__ = []
__copyright__ = "Copyright 2023, Ecole Nationale de l'Aviation Civile (ENAC)"
__license__ = "BSD"

""" """

__author__ = "Théo de la Hogue"
__credits__ = []
__copyright__ = "Copyright 2023, Ecole Nationale de l'Aviation Civile (ENAC)"
__license__ = "BSD"

from typing import TypeVar, Tuple, Any
from dataclasses import dataclass, field
import math
import ast
import json

from argaze import DataStructures

import numpy
import pandas
import cv2

@dataclass(frozen=True)
class GazePosition():
    """Define gaze position as a tuple of coordinates with precision."""

    value: tuple[int | float] = field(default=(0, 0))
    """Position's value."""

    precision: float = field(default=0., kw_only=True)
    """Position's precision represents the radius of a circle around \
    this gaze position value where other same gaze position measurements could be."""

    def __getitem__(self, axis: int) -> int | float:
        """Get position value along a particular axis."""

        return self.value[axis]

    def __iter__(self) -> iter:
        """Iterate over each position value axis."""

        return iter(self.value)
    
    def __len__(self) -> int:
        """Number of axis in position value."""

        return len(self.value)

    def __repr__(self):
        """String representation"""

        return json.dumps(self, ensure_ascii = False, default=vars)

    def __array__(self):
        """Cast as numpy array."""

        return numpy.array(self.value)

    @property
    def valid(self) -> bool:
        """Is the precision not None?"""

        return self.precision is not None

    def distance(self, gaze_position) -> float:
        """Distance to another gaze positions."""

        distance = (self.value[0] - gaze_position.value[0])**2 + (self.value[1] - gaze_position.value[1])**2
        distance = numpy.sqrt(distance)

        return distance

    def overlap(self, gaze_position, both=False) -> float:
        """Does this gaze position overlap another gaze position considering its precision?
        Set both to True to test if the other gaze position overlaps this one too."""

        distance = (self.value[0] - gaze_position.value[0])**2 + (self.value[1] - gaze_position.value[1])**2
        distance = numpy.sqrt(distance)

        if both:
            return distance < min(self.precision, gaze_position.precision)
        else:
            return distance < self.precision

    def draw(self, frame, color=(0, 255, 255), draw_precision=True):
        """Draw gaze position point and precision circle."""

        if self.valid:

            int_value = (int(self.value[0]), int(self.value[1]))

            # Draw point at position
            cv2.circle(frame, int_value, 2, color, -1)

            # Draw precision circle
            if self.precision > 0 and draw_precision:
                cv2.circle(frame, int_value, round(self.precision), color, 1)

class UnvalidGazePosition(GazePosition):
    """Unvalid gaze position."""

    def __init__(self, message=None):

        self.message = message

        super().__init__((None, None), precision=None)

TimeStampedGazePositionsType = TypeVar('TimeStampedGazePositions', bound="TimeStampedGazePositions")
# Type definition for type annotation convenience

class TimeStampedGazePositions(DataStructures.TimeStampedBuffer):
    """Define timestamped buffer to store gaze positions."""

    def __setitem__(self, key, value: GazePosition|dict):
        """Force GazePosition storage."""

        # Convert dict into GazePosition
        if type(value) == dict:

            assert(set(['value', 'precision']).issubset(value.keys()))

            if 'message' in value.keys():

                value = UnvalidGazePosition(value['message'])

            else:

                value = GazePosition(value['value'], precision=value['precision'])

        assert(type(value) == GazePosition or type(value) == UnvalidGazePosition)

        super().__setitem__(key, value)

    @classmethod
    def from_json(self, json_filepath: str) -> TimeStampedGazePositionsType:
        """Create a TimeStampedGazePositionsType from .json file."""

        with open(json_filepath, encoding='utf-8') as ts_buffer_file:

            json_buffer = json.load(ts_buffer_file)

            return TimeStampedGazePositions({ast.literal_eval(ts_str): json_buffer[ts_str] for ts_str in json_buffer})

GazeMovementType = TypeVar('GazeMovement', bound="GazeMovement")
# Type definition for type annotation convenience

@dataclass(frozen=True)
class GazeMovement():
    """Define abstract gaze movement class as a buffer of timestamped positions."""

    positions: TimeStampedGazePositions
    """All timestamp gaze positions."""

    duration: float = field(init=False)
    """Inferred duration from first and last timestamps."""

    amplitude: float = field(init=False)
    """Inferred amplitude from first and last positions."""

    def __post_init__(self):

        start_position_ts, start_position = self.positions.first
        end_position_ts, end_position = self.positions.last

        # Update frozen duration attribute
        object.__setattr__(self, 'duration', end_position_ts - start_position_ts)

        _, start_position = self.positions.first
        _, end_position = self.positions.last

        amplitude = numpy.linalg.norm( numpy.array(start_position.value) - numpy.array(end_position.value))

        # Update frozen amplitude attribute
        object.__setattr__(self, 'amplitude', amplitude)

    def __str__(self) -> str:
        """String display"""

        output = f'{type(self)}:\n\tduration={self.duration}\n\tsize={len(self.positions)}'

        for ts, position in self.positions.items():

            output += f'\n\t{ts}:\n\t\tvalue={position.value},\n\t\taccurracy={position.precision}'

        return output

    def draw_positions(self, frame, color=(0, 55, 55)):
        """Draw gaze movement positions"""

        gaze_positions = self.positions.copy()

        while len(gaze_positions) >= 2:

            ts_start, start_gaze_position = gaze_positions.pop_first()
            ts_next, next_gaze_position = gaze_positions.first

            # Draw start gaze
            start_gaze_position.draw(frame, draw_precision=False)

            # Draw movement from start to next
            cv2.line(frame, start_gaze_position, next_gaze_position, color, 1)

FixationType = TypeVar('Fixation', bound="Fixation")
# Type definition for type annotation convenience

class Fixation(GazeMovement):
    """Define abstract fixation as gaze movement."""

    focus: tuple = field(init=False)
    """Representative position of the fixation."""

    def __post_init__(self):

        super().__post_init__()

    def merge(self, fixation) -> FixationType:
        """Merge another fixation into this fixation."""

        raise NotImplementedError('merge() method not implemented')

    def draw(self, frame, color):
        """Draw fixation into frame."""

        raise NotImplementedError('draw() method not implemented')

def is_fixation(gaze_movement):
    """Is a gaze movement a fixation?"""

    return type(gaze_movement).__bases__[0] == Fixation or type(gaze_movement) == Fixation

class Saccade(GazeMovement):
    """Define abstract saccade as gaze movement."""

    def __post_init__(self):

        super().__post_init__()

    def draw(self, frame, color):
        """Draw saccade into frame."""

        raise NotImplementedError('draw() method not implemented')

def is_saccade(gaze_movement):
    """Is a gaze movement a saccade?"""

    return type(gaze_movement).__bases__[0] == Saccade or type(gaze_movement) == Saccade

TimeStampedGazeMovementsType = TypeVar('TimeStampedGazeMovements', bound="TimeStampedGazeMovements")
# Type definition for type annotation convenience

class TimeStampedGazeMovements(DataStructures.TimeStampedBuffer):
    """Define timestamped buffer to store gaze movements."""

    def __setitem__(self, key, value: GazeMovement):
        """Force value to be or inherit from GazeMovement."""

        assert(isinstance(value, GazeMovement) or type(value).__bases__[0] == Fixation or type(value).__bases__[0] == Saccade)

        super().__setitem__(key, value)

    def __str__(self):

        output = ''
        for ts, item in self.items():

            output += f'\n{item}'

        return output

GazeStatusType = TypeVar('GazeStatus', bound="GazeStatus")
# Type definition for type annotation convenience

@dataclass(frozen=True)
class GazeStatus(GazePosition):
    """Define gaze status as a gaze position belonging to an identified and indexed gaze movement."""

    movement_type: str = field(kw_only=True)
    """GazeMovement type to which gaze position belongs."""

    movement_index: int = field(kw_only=True)
    """GazeMovement index to which gaze positon belongs."""

    @classmethod
    def from_position(cls, gaze_position: GazePosition, movement_type: str, movement_index: int) -> GazeStatusType:
        """Initialize from a gaze position instance."""

        return cls(gaze_position.value, precision=gaze_position.precision, movement_type=movement_type, movement_index=movement_index)

TimeStampedGazeStatusType = TypeVar('TimeStampedGazeStatus', bound="TimeStampedGazeStatus")
# Type definition for type annotation convenience

class TimeStampedGazeStatus(DataStructures.TimeStampedBuffer):
    """Define timestamped buffer to store gaze status."""

    def __setitem__(self, key, value: GazeStatus):
        super().__setitem__(key, value)

class GazeMovementIdentifier():
    """Abstract class to define what should provide a gaze movement identifier."""

    def identify(self, ts, gaze_position, terminate=False) -> GazeMovementType:
        """Identify gaze movement from successive timestamped gaze positions.

            The optional *terminate* argument allows to notify identification algorithm that given gaze position will be the last one.
        """

        raise NotImplementedError('identify() method not implemented')

    def browse(self, ts_gaze_positions: TimeStampedGazePositions) -> Tuple[TimeStampedGazeMovementsType, TimeStampedGazeMovementsType, TimeStampedGazeStatusType]:
        """Identify fixations and saccades browsing timestamped gaze positions."""

        assert(type(ts_gaze_positions) == TimeStampedGazePositions)

        ts_fixations = TimeStampedGazeMovements()
        ts_saccades = TimeStampedGazeMovements()
        ts_status = TimeStampedGazeStatus()

        # Get last ts to terminate identification on last gaze position
        last_ts, _ = ts_gaze_positions.last

        # Iterate on gaze positions
        for ts, gaze_position in ts_gaze_positions.items():

            gaze_movement = self.identify(ts, gaze_position, terminate=(ts == last_ts))

            if is_fixation(gaze_movement):

                start_ts, start_position = gaze_movement.positions.first

                ts_fixations[start_ts] = gaze_movement

                for ts, position in gaze_movement.positions.items():

                    ts_status[ts] = GazeStatus.from_position(position, 'Fixation', len(ts_fixations))

            elif is_saccade(gaze_movement):

                start_ts, start_position = gaze_movement.positions.first

                ts_saccades[start_ts] = gaze_movement

                for ts, position in gaze_movement.positions.items():

                    ts_status[ts] = GazeStatus.from_position(position, 'Saccade', len(ts_saccades))

            else:

                continue

        return ts_fixations, ts_saccades, ts_status

ScanStepType = TypeVar('ScanStep', bound="ScanStep")
# Type definition for type annotation convenience

class ScanStepError(Exception):
    """Exception raised at ScanStepError creation if a aoi scan step doesn't start by a fixation or doesn't end by a saccade."""

    def __init__(self, message):  

        super().__init__(message)

@dataclass(frozen=True)
class ScanStep():
    """Define a scan step as a fixation and a consecutive saccade.
    .. warning::
       Scan step have to start by a fixation and then end by a saccade."""

    first_fixation: Fixation
    """A fixation that comes before the next saccade."""

    last_saccade: Saccade
    """A saccade that comes after the previous fixation."""

    def __post_init__(self):

        # First movement have to be a fixation
        if not is_fixation(self.first_fixation):

            raise ScanStepError('First step movement is not a fixation')

        # Last movement have to be a saccade
        if not is_saccade(self.last_saccade):
            
            raise ScanStepError('Last step movement is not a saccade')

    @property
    def duration(self):
        """Time spent on AOI."""

        # Timestamp of first position of first fixation
        first_ts, _ = self.first_fixation.positions.first 

        # Timestamp of first position of last saccade
        last_ts, _ = self.last_saccade.positions.first 

        return last_ts - first_ts

ScanPathType = TypeVar('ScanPathType', bound="ScanPathType")
# Type definition for type annotation convenience

class ScanPath(list):
    """List of scan steps."""

    def __init__(self):

        super().__init__()
        
        self.__last_fixation = None

    def append_saccade(self, ts, saccade) -> ScanStepType:
        """Append new saccade to scan path and return last new scan step if one have been created."""

        # Ignore saccade if no fixation came before
        if self.__last_fixation != None:

            try: 

                # Edit new step
                new_step = ScanStep(self.__last_fixation, saccade)

                # Append new step
                super().append(new_step)

                # Return new step
                return new_step

            finally:

                # Clear last fixation
                self.__last_fixation = None

    def append_fixation(self, ts, fixation):
        """Append new fixation to scan path.
        !!! warning
                Consecutives fixations are ignored keeping the last fixation"""

        self.__last_fixation = fixation

    def draw(self, frame, fixation_color=(255, 255, 255), saccade_color=(255, 255, 255), deepness=0):
        """Draw scan path into frame."""

        last_step = None
        for step in self[-deepness:]:

            if last_step != None:

                cv2.line(frame, (int(last_step.first_fixation.focus[0]), int(last_step.first_fixation.focus[1])), (int(step.first_fixation.focus[0]), int(step.first_fixation.focus[1])), saccade_color, 2)

                last_step.first_fixation.draw(frame, fixation_color)

            last_step = step

class ScanPathAnalyzer():
    """Abstract class to define what should provide a scan path analyzer."""

    def analyze(self, scan_path: ScanPathType) -> Any:
        """Analyze scan path."""

        raise NotImplementedError('analyze() method not implemented')

AOIScanStepType = TypeVar('AOIScanStep', bound="AOIScanStep")
# Type definition for type annotation convenience

class AOIScanStepError(Exception):
    """Exception raised at AOIScanStepError creation if a aoi scan step doesn't start by a fixation or doesn't end by a saccade."""

    def __init__(self, message, aoi=''):  

        super().__init__(message)

        self.aoi = aoi

@dataclass(frozen=True)
class AOIScanStep():
    """Define a aoi scan step as a set of successive gaze movements onto a same AOI.
    .. warning::
       Aoi scan step have to start by a fixation and then end by a saccade."""

    movements: TimeStampedGazeMovements
    """All movements over an AOI and the last saccade that comes out."""

    aoi: str = field(default='')
    """AOI name."""

    letter: str = field(default='')
    """AOI unique letter to ease sequence analysis."""

    def __post_init__(self):

        # First movement have to be a fixation
        if not is_fixation(self.first_fixation):

            raise AOIScanStepError('First step movement is not a fixation', self.aoi)

        # Last movement have to be a saccade
        if not is_saccade(self.last_saccade):
            
            raise AOIScanStepError('Last step movement is not a saccade', self.aoi)

    @property
    def first_fixation(self):
        """First fixation on AOI."""

        _, first_movement = self.movements.first
        return first_movement

    @property
    def last_saccade(self):
        """Last saccade that comes out AOI."""

        _, last_movement = self.movements.last
        return last_movement

    @property
    def duration(self):
        """Time spent on AOI."""

        # Timestamp of first position of first fixation
        first_ts, _ = self.first_fixation.positions.first 

        # Timestamp of first position of last saccade
        last_ts, _ = self.last_saccade.positions.first 

        return last_ts - first_ts

AOIScanPathType = TypeVar('AOIScanPathType', bound="AOIScanPathType")
# Type definition for type annotation convenience

class AOIScanPath(list):
    """List of aoi scan steps over successive aoi."""

    def __init__(self, expected_aois: list[str] = []):

        super().__init__()
        
        self.__expected_aois = expected_aois
        self.__movements = TimeStampedGazeMovements()
        self.__current_aoi = ''
        self.__index = ord('A')
        self.__aoi_letter = {}
        self.__letter_aoi = {}

        size = len(self.__expected_aois)
        self.__transition_matrix = pandas.DataFrame(numpy.zeros((size, size)), index=self.__expected_aois, columns=self.__expected_aois)

    def __repr__(self):
        """String representation."""

        return str(super())

    def __get_aoi_letter(self, aoi):

        try :

            return self.__aoi_letter[aoi]

        except KeyError:

            letter = chr(self.__index)
            self.__aoi_letter[aoi] = letter
            self.__index += 1
            return letter

    def get_letter_aoi(self, letter):
        """Get which aoi is related to an unique letter."""

        return self.__letter_aoi[letter]

    def __str__(self) -> str:
        """Convert aoi scan path into a string with unique letter per aoi step."""

        sequence = ''
        for step in self:
            sequence += step.letter

        return sequence

    @property
    def expected_aois(self):
        """List of all expected aoi."""

        return self.__expected_aois
    
    @property
    def current_aoi(self):
        """AOI name of aoi scan step under construction"""

        return self.__current_aoi

    @property
    def transition_matrix(self) -> pandas.DataFrame:
        """Pandas DataFrame where indexes are transition departures and columns are transition destinations."""

        return self.__transition_matrix

    def append_saccade(self, ts, saccade):
        """Append new saccade to aoi scan path."""

        # Ignore saccade if no fixation have been stored before
        if len(self.__movements) > 0:

            self.__movements[ts] = saccade

    def append_fixation(self, ts, fixation, looked_aoi: str) -> bool:
        """Append new fixation to aoi scan path and return last new aoi scan step if one have been created.

        !!! warning
                It could raise AOIScanStepError"""

        if looked_aoi not in self.__expected_aois:

            raise AOIScanStepError('AOI not expected', looked_aoi)

        # Is it fixation onto a new aoi?
        if looked_aoi != self.__current_aoi and len(self.__movements) > 0:

            try: 

                # Edit unique letter per aoi
                letter = self.__get_aoi_letter(self.__current_aoi)

                # Remember which letter identify which aoi
                self.__letter_aoi[letter] = self.__current_aoi

                # Edit new step
                new_step = AOIScanStep(self.__movements, self.__current_aoi, letter)

                # Edit transition matrix
                if len(self) > 0:

                    # Increment [index: source, columns: destination] value
                    self.__transition_matrix.loc[self[-1].aoi, self.__current_aoi,] += 1

                # Append new step
                super().append(new_step)

                # Return new step
                return new_step

            finally:

                # Clear movements
                self.__movements = TimeStampedGazeMovements()

                # Append new fixation
                self.__movements[ts] = fixation

                # Remember new aoi
                self.__current_aoi = looked_aoi
        else:

            # Append new fixation
            self.__movements[ts] = fixation

            # Remember aoi
            self.__current_aoi = looked_aoi

            return None

    def fixations_count(self):
        """Get how many fixations are there in the scan path and how many fixation are there in each aoi."""

        scan_fixations_count = 0
        aoi_fixations_count = {aoi: 0 for aoi in self.__expected_aois}

        for aoi_scan_step in self:

            step_fixations_count = len(aoi_scan_step.movements) - 1 # -1: to ignore last saccade

            scan_fixations_count += step_fixations_count
            aoi_fixations_count[aoi_scan_step.aoi] += step_fixations_count
            
        return scan_fixations_count, aoi_fixations_count

class AOIScanPathAnalyzer():
    """Abstract class to define what should provide a aoi scan path analyzer."""

    def analyze(self, aoi_scan_path: AOIScanPathType) -> Any:
        """Analyze aoi scan path."""

        raise NotImplementedError('analyze() method not implemented')