<commit_before>from typing import Dict, Tuple, Union

from datetime import datetime
import io
import logging
import os

import numpy as np

from ..messages.core import *


class MessageData(object):
    def __init__(self, message_type, params):
        self.message_type = message_type
        self.message_class = message_type_to_class[self.message_type]
        self.params = params
        self.messages = []

    def to_numpy(self):
        """!
        @brief Convert the raw FusionEngine message data into numpy arrays that can be used for data analysis.

        On return, this function will add additional member variables to the class containing numpy arrays as returned
        by the `to_numpy()` function in the individual message data class (if supported).

        For example, if called on a @ref fusion_engine_client.messages.solution.PoseMessage "PoseMessage", this function
        will generate new members including `lla_deg`, a 3xN numpy array with the WGS-84 latitude, longitude, and
        altitude data for all available time epochs that can be used as follows:

        ```{.py}
        pose_data.to_numpy()
        mean_lla_deg = np.mean(pose_data.lla_deg, axis=1)
        ```
        """
        if hasattr(self.message_class, 'to_numpy'):
            have_numpy_data = 'p1_time' in self.__dict__
            if have_numpy_data:
                # If we don't have message data we can't do any conversion so the currently cached numpy data is as good
                # as it's gonna get. If it doesn't exist, so be it.
                if len(self.messages) == 0:
                    do_conversion = False
                else:
                    do_conversion = (len(self.messages) != len(self.p1_time) or
                                     float(self.messages[0].p1_time) != self.p1_time[0] or
                                     float(self.messages[-1].p1_time) != self.p1_time[-1])
            else:
                if len(self.messages) == 0:
                    raise ValueError('Raw %s data not available. Cannot convert to numpy.' %
                                     MessageType.get_type_string(self.message_type))
                else:
                    do_conversion = True

            if do_conversion:
                self.__dict__.update(self.message_class.to_numpy(self.messages))
        else:
            raise ValueError('Message type %s does not support numpy conversion.' %
                             MessageType.get_type_string(self.message_type))


class FileIndex(object):
    # Note: To reduce the index file size and load time, we've made the following limitations:
    # - Fractional timestamp is floored so time 123.4 becomes 123. The data read should not assume that an entry's
    #   timestamp is its exact time
    # - Offset is stored as a uint32, assuming the max binary file size is 4 GB. A file with high rate IMU enabled and
    #   recorded overnight was ~1.6 GB. It's reasonable to assume logs large than that will be segmented into multiple
    #   files in the future (i.e., a single 4+ GB binary file is a lot to process).
    RAW_DTYPE = np.dtype([('int', '<u4'), ('type', '<u2'), ('offset', '<u4')])

    DTYPE = np.dtype([('time', '<f8'), ('type', '<u2'), ('offset', '<u8')])

    @classmethod
    def load(cls, index_path):
        raw_data = np.fromfile(index_path, dtype=cls.RAW_DTYPE)
        return cls._from_raw(raw_data)

    @classmethod
    def save(cls, index_path, data: Union[np.ndarray, list]):
        if isinstance(data, np.ndarray) and data.dtype == cls.RAW_DTYPE:
            raw_data = data
            data = cls._from_raw(raw_data)
        else:
            if isinstance(data, list):
                data = np.array(data, dtype=cls.DTYPE)

            if data.dtype == cls.DTYPE:
                raw_data = cls._to_raw(data)
            else:
                raise ValueError('Unsupported array format.')

        raw_data.tofile(index_path)
        return data

    @classmethod
    def get_path(cls, data_path):
        return os.path.splitext(data_path)[0] + '.p1i'

    @classmethod
    def _from_raw(cls, raw_data):
        idx = raw_data['int'] == Timestamp._INVALID
        data = raw_data.astype(dtype=cls.DTYPE)
        data['time'][idx] = np.nan
        return data

    @classmethod
    def _to_raw(cls, data):
        time_sec = data['time']
        idx = np.isnan(time_sec)
        raw_data = data.astype(dtype=cls.RAW_DTYPE)
        raw_data['int'][idx] = Timestamp._INVALID
        return raw_data


class FileReader(object):
    logger = logging.getLogger('point_one.fusion_engine.analysis.file_reader')

    def __init__(self, path=None):
        """!
        @brief Create a new reader instance.

        @param path The path to a FusionEngine binary file to open.
        """
        self.file = None
        self.file_size = 0
        self.data: Dict[MessageType, MessageData] = {}
        self.t0 = None

        self.index = None

        if path is not None:
            self.open(path)

    def open(self, path):
        """!
        @brief Open a FusionEngine binary file.

        @param path The path to the file, or an existing Python file object.
        """
        self.close()

        if isinstance(path, io.IOBase):
            self.file = path
        else:
            self.file = open(path, 'rb')

        self.file.seek(0, io.SEEK_END)
        self.file_size = self.file.tell()
        self.file.seek(0, 0)

        # Load the data index file if present.
        index_path = FileIndex.get_path(self.file.name)
        if os.path.exists(index_path):
            self.logger.debug("Reading index file '%s'." % index_path)
            self.index = FileIndex.load(index_path)
        else:
            self.index = None

        # Read the first message (with P1 time) in the file to set self.t0.
        self.read(time_range=(0.0, None), max_messages=1, generate_index=False)

    def close(self):
        """!
        @brief Close the file.
        """
        if self.file is not None:
            self.file = None

    def read(self, message_types: Union[list, tuple] = None,
             time_range: Tuple[Union[float, Timestamp], Union[float, Timestamp]] = None, absolute_time: bool = False,
             max_messages: int = None,
             return_numpy: bool = False, keep_messages: bool = False,
             generate_index: bool = True, show_progress: bool = False) \
            -> Dict[MessageType, MessageData]:
        """!
        @brief Read data for one or more desired message types.
<commit_msg>Added a function to explicitly generate an index.<commit_after>from typing import Dict, Tuple, Union

from datetime import datetime
import io
import logging
import os

import numpy as np

from ..messages.core import *


class MessageData(object):
    def __init__(self, message_type, params):
        self.message_type = message_type
        self.message_class = message_type_to_class[self.message_type]
        self.params = params
        self.messages = []

    def to_numpy(self):
        """!
        @brief Convert the raw FusionEngine message data into numpy arrays that can be used for data analysis.

        On return, this function will add additional member variables to the class containing numpy arrays as returned
        by the `to_numpy()` function in the individual message data class (if supported).

        For example, if called on a @ref fusion_engine_client.messages.solution.PoseMessage "PoseMessage", this function
        will generate new members including `lla_deg`, a 3xN numpy array with the WGS-84 latitude, longitude, and
        altitude data for all available time epochs that can be used as follows:

        ```{.py}
        pose_data.to_numpy()
        mean_lla_deg = np.mean(pose_data.lla_deg, axis=1)
        ```
        """
        if hasattr(self.message_class, 'to_numpy'):
            have_numpy_data = 'p1_time' in self.__dict__
            if have_numpy_data:
                # If we don't have message data we can't do any conversion so the currently cached numpy data is as good
                # as it's gonna get. If it doesn't exist, so be it.
                if len(self.messages) == 0:
                    do_conversion = False
                else:
                    do_conversion = (len(self.messages) != len(self.p1_time) or
                                     float(self.messages[0].p1_time) != self.p1_time[0] or
                                     float(self.messages[-1].p1_time) != self.p1_time[-1])
            else:
                if len(self.messages) == 0:
                    raise ValueError('Raw %s data not available. Cannot convert to numpy.' %
                                     MessageType.get_type_string(self.message_type))
                else:
                    do_conversion = True

            if do_conversion:
                self.__dict__.update(self.message_class.to_numpy(self.messages))
        else:
            raise ValueError('Message type %s does not support numpy conversion.' %
                             MessageType.get_type_string(self.message_type))


class FileIndex(object):
    # Note: To reduce the index file size and load time, we've made the following limitations:
    # - Fractional timestamp is floored so time 123.4 becomes 123. The data read should not assume that an entry's
    #   timestamp is its exact time
    # - Offset is stored as a uint32, assuming the max binary file size is 4 GB. A file with high rate IMU enabled and
    #   recorded overnight was ~1.6 GB. It's reasonable to assume logs large than that will be segmented into multiple
    #   files in the future (i.e., a single 4+ GB binary file is a lot to process).
    RAW_DTYPE = np.dtype([('int', '<u4'), ('type', '<u2'), ('offset', '<u4')])

    DTYPE = np.dtype([('time', '<f8'), ('type', '<u2'), ('offset', '<u8')])

    @classmethod
    def load(cls, index_path):
        raw_data = np.fromfile(index_path, dtype=cls.RAW_DTYPE)
        return cls._from_raw(raw_data)

    @classmethod
    def save(cls, index_path, data: Union[np.ndarray, list]):
        if isinstance(data, np.ndarray) and data.dtype == cls.RAW_DTYPE:
            raw_data = data
            data = cls._from_raw(raw_data)
        else:
            if isinstance(data, list):
                data = np.array(data, dtype=cls.DTYPE)

            if data.dtype == cls.DTYPE:
                raw_data = cls._to_raw(data)
            else:
                raise ValueError('Unsupported array format.')

        raw_data.tofile(index_path)
        return data

    @classmethod
    def get_path(cls, data_path):
        return os.path.splitext(data_path)[0] + '.p1i'

    @classmethod
    def _from_raw(cls, raw_data):
        idx = raw_data['int'] == Timestamp._INVALID
        data = raw_data.astype(dtype=cls.DTYPE)
        data['time'][idx] = np.nan
        return data

    @classmethod
    def _to_raw(cls, data):
        time_sec = data['time']
        idx = np.isnan(time_sec)
        raw_data = data.astype(dtype=cls.RAW_DTYPE)
        raw_data['int'][idx] = Timestamp._INVALID
        return raw_data


class FileReader(object):
    logger = logging.getLogger('point_one.fusion_engine.analysis.file_reader')

    def __init__(self, path=None):
        """!
        @brief Create a new reader instance.

        @param path The path to a FusionEngine binary file to open.
        """
        self.file = None
        self.file_size = 0
        self.data: Dict[MessageType, MessageData] = {}
        self.t0 = None

        self.index = None

        if path is not None:
            self.open(path)

    def open(self, path):
        """!
        @brief Open a FusionEngine binary file.

        @param path The path to the file, or an existing Python file object.
        """
        self.close()

        if isinstance(path, io.IOBase):
            self.file = path
        else:
            self.file = open(path, 'rb')

        self.file.seek(0, io.SEEK_END)
        self.file_size = self.file.tell()
        self.file.seek(0, 0)

        # Load the data index file if present.
        index_path = FileIndex.get_path(self.file.name)
        if os.path.exists(index_path):
            self.logger.debug("Reading index file '%s'." % index_path)
            self.index = FileIndex.load(index_path)
        else:
            self.index = None

        # Read the first message (with P1 time) in the file to set self.t0.
        #
        # Note that we explicitly set a start time since, if the time range is not specified, read() will include
        # messages that do not have P1 time. We want to make sure the 1 message is one with time.
        self.read(time_range=(0.0, None), max_messages=1, generate_index=False)

    def close(self):
        """!
        @brief Close the file.
        """
        if self.file is not None:
            self.file = None

    def generate_index(self):
        """!
        @brief Generate an index file for the current binary file if one does not already exist.
        """
        if self.index is None:
            # We'll read pose data (doesn't actually matter which). Store the currently cached data and restore it when
            # we're done. That way if the user already did a read (with generate_index == False), they don't have to
            # re-read the data if they try to use it again.
            prev_data = self.data.get(MessageType.POSE, None)

            self.read(message_types=[MessageType.POSE], max_messages=1, generate_index=True)

            if prev_data is not None:
                self.data[MessageType.POSE] = prev_data

    def read(self, message_types: Union[list, tuple] = None,
             time_range: Tuple[Union[float, Timestamp], Union[float, Timestamp]] = None, absolute_time: bool = False,
             max_messages: int = None,
             return_numpy: bool = False, keep_messages: bool = False,
             generate_index: bool = True, show_progress: bool = False) \
            -> Dict[MessageType, MessageData]:
        """!
        @brief Read data for one or more desired message types.

{'fraction_non_alphanumeric': 0.09561639948928927, 'fraction_numerical': 0.00688040856859129, 'mean_word_length': 2.1164898320070735, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3102417', 'n_tokens_mistral': 4122, 'n_tokens_neox': 3961, 'n_words': 1287}