<commit_before>"""
This module implements file reader for AlphaOmega MPX file format version 4.

This module expect default channel names from the AlphaOmega record system (RAW
###, SPK ###, LFP ###, AI ###,…).

This module reads all \*.lsx and \*.mpx files in a directory (not recursively).
Listing files

The specifications are mostly extracted from the "AlphaRS User Manual V1.0.1.pdf"
manual provided with the AlphaRS hardware. The specifications are described in
the chapter 6: ALPHARS FILE FORMAT. See at the end of this file for file format
blocks description.
Some informations missing from the file specifications were kindly provided by
AlphaOmega engineers.

.. note::
    Not a lot of memory optimization effort was put into this module. You should
    expect a big memory usage when loading data with this module

.. todo::
    1. Make TTL events -> epochs? We could describe TTL 0/1 events as epochs:
       when the TTL goes to 1=beginning of epoch 0=end of epoch
    2. decode stram data (block type E) using StreamFormat.h file (not provided
       with this code)
    3. (2.bis) Once we know how to decode Stream AlphaOmega events add them to
       an event channel

Author: <NAME> <<EMAIL>>
"""

import io
import logging
import mmap
import os
import struct

from collections import defaultdict
from datetime import datetime
from itertools import chain
from pathlib import Path, PureWindowsPath

import numpy as np

from .baserawio import (
    BaseRawIO,
    _signal_stream_dtype,
    _signal_channel_dtype,
    _spike_channel_dtype,
    _event_channel_dtype,
)


class AlphaOmegaRawIO(BaseRawIO):
    """
    AlphaOmega MPX file format 4 reader. Handles several blocks and segments.

    A block is a recording defined in a \*.lsx file. AlphaOmega record system
    creates such files every time the software is opened.

    A segment is a continuous record (when record starts/stops).
    If file are not referenced in a \*.lsx file, they are put in the same block.

    :param dirname: folder from where to load the data
    :type dirname: str or Path-like
    :param prune_channels: if True removes the empty channels, defaults to True
    :type prune_channels: bool

    .. warning::
        Because channels must be gathered into coherent streams, channels names
        MUST be the default channel names in AlphaRS software (or Alpha LAB SNR).
    """

    extensions = ["lsx", "mpx"]
    rawmode = "one-dir"

    STREAM_CHANNELS = (
        ("Accelerometer", "ACC", "ACC"),
        ("Spike Train", "SPK", "SPK"),
        ("Local Field Potential", "LFP", "LFP"),
        ("Raw Analog", "RAW", "RAW"),
        ("Analog Input", "AI", "AI"),
        ("User Defined", "UD", "UD"),
    )

    EVENT_CHANNELS = (
        ("TTL", "TTL", "TTL"),
        ("Digital Output", "DOUT", "DOUT"),
        ("User Event", "UD", "UD"),
        ("Stim Marker", "StimMarker", "StimMarker"),
        ("Digital Port", "InPort", "InPort"),
        ("Internal Detection", "Internal Detection", "Internal Detection"),
    )

    def __init__(self, dirname="", prune_channels=True):
        super().__init__(dirname=dirname)
        self.dirname = Path(dirname)
        self._filenames = {}
        if self.dirname.is_dir():
            self._explore_folder()
        else:
            self.logger.error(f"{self.dirname} is not a folder")
        self._prune_channels = prune_channels
        self._opened_files = {}
        self._ignore_unknown_blocks = True  # internal debug property

    def _explore_folder(self):
        """
        Explores an AlphaOmega folder and tries to load the index files (*.lsx)
        to get the data files (*.mpx). If the folder contains only *.mpx files,
        it will load them without splitting them into blocks.
        It does not explores the folder recursively.
        """
        self._filenames = defaultdict(list)

        # first check if there is any *.mpx files
        filenames = list(filter(lambda x: x.is_file(), self.dirname.glob("*.mpx")))
        if not filenames:
            self.logger.error(f"Found no AlphaOmega *.mpx files in {self.dirname}")
        else:
            index_files = list(
                filter(lambda x: x.is_file(), self.dirname.glob("*.lsx"))
            )
            # the following is not very usefull but we should expect the original
            # names to be in increasing lexical order
            index_files.sort()
            if not index_files:
                self.logger.warning(
                    "No *.lsx files found. Will try to load all *.mpx files"
                )
                self._filenames[""].extend(filenames)
            else:
                for index_file in index_files:
                    with open(index_file, "r") as f:
                        for line in f:
                            # a line is a Microsoft Windows path. As we cannot
                            # instanciate a WindowsPath on other OS than MS
                            # Windows, we use the PureWindowsPath class
                            filename = PureWindowsPath(line.strip())
                            filename = self.dirname / filename.name
                            if not filename.is_file():
                                self.logger.warning(f"File {filename} does not exist")
                            else:
                                self._filenames[index_file.name].append(filename)
                                filenames.remove(filename)
                if filenames:
                    self.logger.info(
                        "Some *.mpx files not referenced. Will try to load them."
                    )
                    self._filenames[""].extend(filenames)

    def _source_name(self):
        return str(self.dirname)

    def _read_file_blocks(self, filename, prune_channels=True):
        """Read blocks from AlphaOmega MPX file version 4.

        :param filename: the MPX filename to read blocks from
        :type filename: Path-like object or str
        :param prune_channels: Remove references to channels and ports which
            doesn't contain any data recorded. Be careful when using this option
            with multiple-file data since it could theoretically leads to
            expection raised when data recorded in further files are merged into
            the first file pruned from these channels.
        :type prune_channels: bool
        """
        continuous_analog_channels = {}
        segmented_analog_channels = {}
        digital_channels = {}
        channel_type = {}
        stream_data_channels = {}
        ports = {}
        events = []
        unknown_blocks = []
        with open(filename, "rb") as f:
            length, block_type = HeaderType.unpack(f.read(HeaderType.size))
            # First block is always type h and size 60
            if block_type != b"h":
                try:
                    bt = block_type.decode()
                    self.logger.error("First block must be of type h not type {bt}")
                except UnicodeDecodeError:
                    self.logger.error(
                        (
                            f"First block must be of type h not type "
                            f"{int.from_bytes(bt, 'little')} (int format)"
                        )
                    )
                raise Exception(
                    "First block of AlphaOmega MPX file format must be of type h"
                )
            if not length == 60:
                self.logger.error("First block must be of size 60 (got size {length})")
                raise Exception(
                    "First block of AlphaOmega MPX file format must be of size 60"
                )
            (
                next_block,
                version,
                hour,
                minute,
                second,
                hsecond,
                day,
                month,
                year,
                dayofweek,
                minimum_time,
                maximum_time,
                erase_count,
                map_version,
                application_name,
                resource_version,
                reserved,
            ) = SDataHeader.unpack(f.read(SDataHeader.size))
            if map_version != 4:
                self.logger.error("Only map version 4 is supported")
                raise Exception("Only AlphaOmega MPX file format 4 is supported")
            resource_version = "".join([str(b) for b in resource_version])
            try:
                resource_version = int(resource_version)
            except ValueError:
                self.logger.error(
                    f"m_ResourceVersion should be an integer (got: {resource_version}"
                )
            metadata = {
                "application_version": version,
                "application_name": decode_string(application_name),
                "record_date": datetime(
                    year, month, day, hour, minute, second, 10000 * hsecond
                ),
                "start_time": minimum_time,
                "stop_time": maximum_time,
                "erase_count": erase_count,
                "map_version": map_version,
                "resource_version": resource_version,
                "max_sample_rate": 0,
            }

            pos = 0
            while True:
                pos += length
                f.seek(pos)
                header_bytes = f.read(HeaderType.size)

                length, block_type = HeaderType.unpack(header_bytes)
                if length == 65535:
                    # The stop condition for reading MPX blocks is base on the
                    # length of the block: if the block has length 65535 (or -1
                    # in signed integer value) we know we have reached the end
                    # of the file
                    # We could also check that we are at the end of the file
                    # after this block
                    break

                if block_type == b"h":
                    self.logger.error(
                        "Type h block must exist only at the beginning of file"
                    )
                    raise Exception(
                        "AlphaOmega MPX file format must not have type h block after first block"
                    )

                if block_type == b"2":
                    (
                        next_block,
                        is_analog,
                        is_input,
                        channel_number,
                        *spike_color,
                    ) = Type2Block.unpack(f.read(Type2Block.size))
                    if is_analog and is_input:
                        (
                            mode,
                            amplitude,
                            sample_rate,
                            spike_count,
                            mode_spike,
                        ) = SDefAnalog.unpack(f.read(SDefAnalog.size))
                        mode_spike_mode = (mode_spike & 0b1111000000000000) >> 12
                        mode_spike_master = mode_spike_mode == 1
                        mode_spike_slave = mode_spike_mode == 2
                        mode_spike_linked_channel = mode_spike & 0b0000111111111111
                        mode_spike = {
                            "mode_spike_orig": mode_spike,
                            "mode_spike_master": mode_spike_master,
                            "mode_spike_slave": mode_spike_slave,
                            "mode_spike_linked_channel": mode_spike_linked_channel,
                        }
                        metadata["max_sample_rate"] = max(
                            metadata["max_sample_rate"], sample_rate * 1000
                        )
                        if amplitude <= 5:
                            # This is true for any logging software for map
                            # version >4 (specs say only for ALab SNR but AO
                            # engineer says it's true for any software)
                            amplitude = 1250000 / 2 ** 15
                        if mode == 0:
                            # continuous analog channel definition block
                            assert channel_number not in channel_type
                            channel_type[channel_number] = "continuous_analog"
                            duration, total_gain_100 = SDefContinAnalog.unpack(
                                f.read(SDefContinAnalog.size)
                            )
                            name_length = length - 38
                            name = get_name(f, name_length)
                            assert channel_number not in continuous_analog_channels
                            continuous_analog_channels[channel_number] = {
                                "spike_color": spike_color,
                                "bit_resolution": amplitude,
                                "sample_rate": sample_rate * 1000,
                                "spike_count": spike_count,
                                "mode_spike": mode_spike,
                                "duration": duration,
                                "gain": total_gain_100 / 100,
                                "name": name,
                                "positions": defaultdict(list),
                            }
                        elif mode == 1:
                            # segmented analog channel definition block
                            assert channel_number not in channel_type
                            channel_type[channel_number] = "segmented_analog"
                            (
                                pre_trig_ms,
                                post_trig_ms,
                                level_value,
                                trg_mode,
                                yes_rms,
                                total_gain_100,
                            ) = SDefLevelAnalog.unpack(f.read(SDefLevelAnalog.size))
                            name_length = length - 48
                            name = get_name(f, name_length)
                            assert channel_number not in segmented_analog_channels
                            segmented_analog_channels[channel_number] = {
                                "spike_color": spike_color,
                                "bit_resolution": amplitude,
                                "sample_rate": sample_rate * 1000,
                                "spike_count": spike_count,
                                "mode_spike": mode_spike,
                                "pre_trig_duration": pre_trig_ms / 1000,
                                "post_trig_duration": post_trig_ms / 1000,
                                "level_value": level_value,
                                "trg_mode": trg_mode,
                                "automatic_level_base_rms": yes_rms,
                                "gain": total_gain_100 / 100,
                                "name": name,
                                "positions": defaultdict(list),
                            }
                        else:
                            self.logger.error(
                                f"Unknown type 2 analog block mode: {mode}"
                            )
                            continue
                    elif is_analog == 0 and is_input == 1:
                        # digital input channel definition
                        assert channel_number not in channel_type
                        channel_type[channel_number] = "digital"
                        (
                            sample_rate,
                            save_trigger,
                            duration,
                            prev_status,
                        ) = SDefDigitalInput.unpack(f.read(SDefDigitalInput.size))
                        metadata["max_sample_rate"] = max(
                            metadata["max_sample_rate"], sample_rate * 1000
                        )
                        assert channel_number not in digital_channels
                        name_length = length - 30
                        name = get_name(f, name_length)
                        digital_channels[channel_number] = {
                            "spike_color": spike_color,
                            "sample_rate": sample_rate * 1000,
                            "save_trigger": save_trigger,
                            "duration": duration,
                            "prev_status": prev_status,
                            "name": name,
                            "samples": [],
                        }
                    else:
                        self.logger.error(
                            f"Unknown type 2 block: analog={is_analog}, input={is_input}"
                        )
                        continue
                elif block_type == b"S":
                    # stream data definition block
                    next_block, channel_number, sample_rate = SDefStream.unpack(
                        f.read(SDefStream.size)
                    )
                    metadata["max_sample_rate"] = max(
                        metadata["max_sample_rate"], sample_rate * 1000
                    )
                    assert channel_number not in channel_type
                    channel_type[channel_number] = "stream_data"
                    name_length = length - 18
                    name = get_name(f, name_length)
                    stream_data_channels[channel_number] = {
                        "sample_rate": sample_rate * 1000,
                        "name": name,
                    }
                elif block_type == b"b":
                    # digital input/output port definition block
                    board_number, port, sample_rate, prev_value = SDefPortX.unpack(
                        f.read(SDefPortX.size)
                    )
                    metadata["max_sample_rate"] = max(
                        metadata["max_sample_rate"], sample_rate * 1000
                    )
                    assert port not in channel_type
                    channel_type[port] = "port"
                    name_length = length - 18
                    name = get_name(f, name_length)
                    ports[port] = {
                        "board_number": board_number,
                        "sample_rate": sample_rate * 1000,
                        "prev_value": prev_value,
                        "name": name,
                        "samples": [],
                    }
                elif block_type == b"5":
                    # channel data block
                    unit_number, channel_number = SDataChannel.unpack(
                        f.read(SDataChannel.size)
                    )
                    assert channel_number in channel_type
                    unit_number = int.from_bytes(unit_number, "little")
                    if "analog" in channel_type[channel_number]:
                        data_length = (length - 10) / 2
                        assert int(data_length) == data_length
                        data_length = int(data_length)
                        data_start = f.tell()
                        f.seek(2 * data_length, io.SEEK_CUR)
                        if channel_type[channel_number].startswith("continuous"):
                            assert channel_number in continuous_analog_channels
                            continuous_analog_channels[channel_number]["positions"][
                                filename
                            ].append(
                                (
                                    SDataChannel_sample_id.unpack(
                                        f.read(SDataChannel_sample_id.size)
                                    )[0],
                                    data_start,
                                    data_length,
                                )
                            )
                        elif channel_type[channel_number].startswith("segmented"):
                            assert channel_number in segmented_analog_channels
                            if unit_number > 0 and unit_number <= 4:
                                segmented_analog_channels[channel_number]["positions"][
                                    filename
                                ].append(
                                    (
                                        SDataChannel_sample_id.unpack(
                                            f.read(SDataChannel_sample_id.size)
                                        )[0],
                                        data_start,
                                        data_length,
                                    )
                                )
                            elif unit_number == 0:
                                segmented_analog_channels[channel_number]["positions"][
                                    filename
                                ].append(
                                    (
                                        SDataChannel_sample_id.unpack(
                                            f.read(SDataChannel_sample_id.size)
                                        )[0],
                                        data_start,
                                        data_length,
                                    )
                                )
                            else:
                                self.logger.error(
                                    f"Unknown unit_number={unit_number} in channel data block"
                                )
                                continue
                    elif channel_type[channel_number] == "digital":
                        assert channel_number in digital_channels
                        sample_number, value = SDataChannelDigital.unpack(
                            f.read(SDataChannelDigital.size)
                        )
                        digital_channels[channel_number]["samples"].append(
                            (
                                sample_number,
                                value,
                            )
                        )
                    elif channel_type[channel_number] == "port":
                        assert channel_number in ports
                        # specifications says that for ports it should be "<Lh"
                        # but the data shows clearly "<HL"
                        value, sample_number = SDataChannelPort.unpack(
                            f.read(SDataChannelPort.size)
                        )
                        ports[channel_number]["samples"].append(
                            (
                                sample_number,
                                value,
                            )
                        )
                    else:
                        self.logger.error(
                            f"Unknown channel_type={channel_type[channel_number]} for block type 5"
                        )
                elif block_type == b"E":
                    type_event, timestamp = SAOEvent.unpack(f.read(SAOEvent.size))
                    stream_data_length = length - 8
                    events.append(
                        {
                            "timestamp": timestamp,
                            "stream_data": struct.unpack(
                                f"<{stream_data_length}s", f.read(stream_data_length)
                            )[0],
                        }
                    )
                else:
                    if not self._ignore_unknown_blocks:
                        try:
                            bt = block_type.decode()
                            self.logger.debug(
                                f"Unknown block type: block length: {length}, block_type: {bt}"
                            )
                        except UnicodeDecodeError:
                            self.logger.debug(
                                (
                                    f"Unknown block type: block length: {length}, "
                                    f"block_type: {int.from_bytes(block_type, 'little')} "
                                    "(int format)"
                                )
                            )
                        unknown_blocks.append(
                            {
                                "length": length,
                                "block_type": block_type,
                                "data": f.read(length),
                            }
                        )
        if prune_channels:
            to_remove = []
            for channel_id, channel in continuous_analog_channels.items():
                if not channel["positions"]:
                    to_remove.append(channel_id)
            for channel_id in to_remove:
                del continuous_analog_channels[channel_id]
            to_remove = []
            for channel_id, channel in segmented_analog_channels.items():
                if not channel["positions"]:
                    to_remove.append(channel_id)
            for channel_id in to_remove:
                del segmented_analog_channels[channel_id]
            to_remove = []
            for channel_id, channel in digital_channels.items():
                if not channel["samples"]:
                    to_remove.append(channel_id)
            for channel_id in to_remove:
                del digital_channels[channel_id]
            to_remove = []
            for port_id, port in ports.items():
                if not port["samples"]:
                    to_remove.append(port_id)
            for port_id in to_remove:
                del ports[port_id]
            del to_remove
        return (
            metadata,
            continuous_analog_channels,
            segmented_analog_channels,
            digital_channels,
            channel_type,
            stream_data_channels,
            ports,
            events,
            unknown_blocks,
        )

    def _merge_segments(self, factor_period=1.5):
        """This method merge segments that are the same segment but comes from
        different files. The AlphaOmega data recording system split files with
        a configurable maximum time and size but with a global non-configurable
        maximum size of 1GB.
        Two segment are merged if they validate the following conditions:
            - start of next segment is less than `factor_period` (default 1.5)
            the system's sampling period + end of previous segment. We can't
            check with exactly 1 sampling period because timings are floats
            which are not exact period values
            - The two segment must also have the same record date
            (YEAR-MONTH-DAY). This could potentially lead to errors if
            recordings are longer than a day or run over the night
            - The next segment must have a datetime greater or equal than the
            previous segment (datetime : YEAR-MONTH-DAY-HOUR-MINUTE-SECOND)

        :param factor_period: how many sample period to consider the next segment
            is part of the previous one. This should be 1 <= `factor_period` < 2,
            defaults to 1.5
        :type factor_period: float
        """
        for block in self._blocks:
            segments_to_merge = []
            for segment in block:
                possible_same_segments = [
                    s
                    for s in block
                    if s["metadata"]["record_date"].date()
                    == segment["metadata"]["record_date"].date()
                    and s["metadata"]["record_date"]
                    >= segment["metadata"]["record_date"]
                    and 0
                    <= (s["metadata"]["start_time"] - segment["metadata"]["stop_time"])
                    <= factor_period / segment["metadata"]["max_sample_rate"]
                    and s is not segment
                ]
                assert len(possible_same_segments) in (0, 1)
                if possible_same_segments:
                    existing_merges = [s for segs in segments_to_merge for s in segs]
                    if existing_merges:
                        existing_merges.sort(key=lambda x: x["metadata"]["start_time"])
                        segment = existing_merges[0]
                    segments_to_merge.append((segment, possible_same_segments[0]))
            for segment, segment_to_merge in segments_to_merge:
                assert (
                    segment["metadata"]["max_sample_rate"]
                    == segment_to_merge["metadata"]["max_sample_rate"]
                )
                segment["metadata"]["stop_time"] = segment_to_merge["metadata"][
                    "stop_time"
                ]
                segment["metadata"]["filenames"].extend(
                    segment_to_merge["metadata"]["filenames"]
                )
                for stream in segment_to_merge["streams"]:
                    for channel_id in segment_to_merge["streams"][stream]:
                        assert channel_id in segment["streams"][stream]
                        for f in segment_to_merge["streams"][stream][channel_id][
                            "positions"
                        ]:
                            segment["streams"][stream][channel_id]["positions"][
                                f
                            ].extend(
                                segment_to_merge["streams"][stream][channel_id][
                                    "positions"
                                ][f]
                            )
                for channel_id in segment_to_merge["events"]:
                    # it's possible  that the first segment doesn't record any
                    # port channel data and the port could have been pruned
                    if channel_id in segment["events"]:
                        segment["events"][channel_id]["samples"].extend(
                            segment_to_merge["events"][channel_id]["samples"]
                        )
                    else:
                        segment["events"][channel_id] = segment_to_merge["events"][
                            channel_id
                        ]
                for channel_id in segment_to_merge["spikes"]:
                    assert channel_id in segment["spikes"]
                    for f in segment_to_merge["spikes"][channel_id]["positions"]:
                        segment["spikes"][channel_id]["positions"][f].extend(
                            segment_to_merge["spikes"][channel_id]["positions"][f]
                        )
                segment["ao_events"].extend(segment_to_merge["ao_events"])
                for channel_id in segment_to_merge["stream_data"]:
                    # To be honest, I have no idea what is a
                    # stream_data_channels so let's just overwrite it here
                    segment["stream_data"][channel_id] = segment_to_merge[
                        "stream_data"
                    ][channel_id]
                block.remove(segment_to_merge)

    def _parse_header(self):
        blocks = []
        for index_file, filenames in self._filenames.items():
            segments = []
            continuous_analog_channels = {}
            for i, filename in enumerate(filenames):
                metadata, cac, sac, dc, ct, sd, p, e, ub = self._read_file_blocks(
                    filename, self._prune_channels
                )
                metadata["filenames"] = [filename]
                streams = {}
                for stream_name, channel_name_start, stream_id in self.STREAM_CHANNELS:
                    channels = {
                        i: c
                        for i, c in cac.items()
                        if c["name"].startswith(channel_name_start)
                    }
                    streams[stream_id] = channels
                events = dc.copy()
                events.update(p)
                segment = {
                    "metadata": metadata,
                    "streams": streams,
                    "events": events,
                    "spikes": sac,
                    "ao_events": e,
                    "stream_data": sd,
                }

                segments.append(segment)
            blocks.append(segments)
        self._blocks = blocks
        # We merge segments after having loading all the files because they
        # could be loaded in any order
        self._merge_segments()

        signal_streams = set(
            (stream_name, stream_id)
            for block in self._blocks
            for segment in block
            for stream in segment["streams"]
            for stream_name, _, stream_id in self.STREAM_CHANNELS
            if stream_id == stream and segment["streams"][stream]
        )
        signal_streams = list(signal_streams)
        signal_streams.sort(key=lambda x: x[1])
        signal_streams = np.array(signal_streams, dtype=_signal_stream_dtype)

        signal_channels = set(
            (
                channel["name"],
                channel_id,
                channel["sample_rate"],
                np.dtype(np.short).name,
                "uV",
                channel["gain"] / channel["bit_resolution"],
                0,
                stream,
            )
            for block in self._blocks
            for segment in block
            for stream in segment["streams"]
            for channel_id, channel in segment["streams"][stream].items()
        )
        signal_channels = list(signal_channels)
        signal_channels.sort(key=lambda x: (x[7], x[0]))
        signal_channels = np.array(signal_channels, dtype=_signal_channel_dtype)

        spike_channels = set(
            (
                c["name"],
                i,
                "uV",
                c["gain"] / c["bit_resolution"],
                0,
                round(c["pre_trig_duration"] * c["sample_rate"]),
                c["sample_rate"],
            )
            for block in self._blocks
            for segment in block
            for i, c in segment["spikes"].items()
        )
        spike_channels = list(spike_channels)
        spike_channels.sort(key=lambda x: x[0])
        spike_channels = np.array(spike_channels, dtype=_spike_channel_dtype)

        event_channels = set(
            (event["name"], i, "event")
            for block in self._blocks
            for segment in block
            for i, event in segment["events"].items()
        )
        event_channels = list(event_channels)
        event_channels.sort(key=lambda x: x[1])
        event_channels = np.array(event_channels, dtype=_event_channel_dtype)

        self.header = {}
        self.header["nb_block"] = len(self._filenames)
        self.header["nb_segment"] = [len(segment) for segment in self._blocks]
        self.header["signal_streams"] = signal_streams
        self.header["signal_channels"] = signal_channels
        self.header["spike_channels"] = spike_channels
        self.header["event_channels"] = event_channels

        self._generate_minimal_annotations()

        for block_index, filename in enumerate(self._filenames):
            bl_ann = self.raw_annotations["blocks"][block_index]
            bl_ann["name"] = "Block #{}{}".format(
                block_index, " from lsx file" if filename else ""
            )
            bl_ann["file_origin"] = (
                str(self.dirname / filename) if filename else bl_ann["file_origin"]
            )
            bl_ann["rec_datetime"] = self._blocks[block_index][0]["metadata"][
                "record_date"
            ]
            for seg_index, segment in enumerate(self._blocks[block_index]):
                seg_ann = bl_ann["segments"][seg_index]
                seg_ann["name"] = "Seg #{} Block #{}".format(seg_index, block_index)
                seg_ann["file_origin"] = "\n".join(
                    str(f)
                    for f in self._blocks[block_index][seg_index]["metadata"][
                        "filenames"
                    ]
                )
                seg_ann["rec_datetime"] = self._blocks[block_index][seg_index][
                    "metadata"
                ]["record_date"]
                for c_index, c in enumerate(seg_ann["signals"]):
                    c = c.copy()
                    c["file_origin"] = "\n".join(
                        set(
                            str(f)
                            for channels in self._blocks[block_index][seg_index][
                                "streams"
                            ][c["stream_id"]].values()
                            for f in channels["positions"]
                        )
                    )
                    seg_ann["signals"][c_index] = c
                for e_index, e in enumerate(seg_ann["events"]):
                    e = e.copy()
                    e["file_origin"] = seg_ann["file_origin"]
                    seg_ann["events"][e_index] = e

        # We open files and create mmap objects
        for block in self._filenames:
            for filename in self._filenames[block]:
                if filename not in self._opened_files:
                    self._opened_files[filename] = {}
                    self._opened_files[filename]["file"] = filename.open(mode="rb")
                    self._opened_files[filename]["mmap"] = mmap.mmap(
                        self._opened_files[filename]["file"].fileno(),
                        0,
                        access=mmap.ACCESS_READ,
                    )

    def __del__(self):
        # To be sure we close the file when object is deleted. Be aware that the
        # __del__ method is not necessarily called when interpreter exits so we
        # could still leave file opened. This is probably bad…
        for filename in self._opened_files:
            self._opened_files[filename]["mmap"].close()
            self._opened_files[filename]["file"].close()
        if hasattr(super(), "__del__"):
            super().__del__()

    def _segment_t_start(self, block_index, seg_index):
        return self._blocks[block_index][seg_index]["metadata"]["start_time"]

    def _segment_t_stop(self, block_index, seg_index):
        return self._blocks[block_index][seg_index]["metadata"]["stop_time"]

    def _get_signal_size(self, block_index, seg_index, stream_index):
        stream_id = self.header["signal_streams"][stream_index]["id"]
        sizes = [
            sum(
                sample[2]
                for sample_by_file in channel["positions"].values()
                for sample in sample_by_file
            )
            for channel in self._blocks[block_index][seg_index]["streams"][
                stream_id
            ].values()
        ]
        assert all(s == sizes[0] for s in sizes)
        return sizes[0]

    def _get_signal_t_start(self, block_index, seg_index, stream_index):
        return self._segment_t_start(block_index, seg_index)

    def _get_analogsignal_chunk(
        self, block_index, seg_index, i_start, i_stop, stream_index, channel_indexes
    ):
        if i_start is None:
            i_start = 0
        signal_size = self._get_signal_size(block_index, seg_index, stream_index)
        if i_stop is None or i_stop > signal_size:
            i_stop = signal_size
        stream_id = self.header["signal_streams"][stream_index]["id"]
        mask = self.header["signal_channels"]["stream_id"] == stream_id
        channel_ids = self.header["signal_channels"][mask]["id"][
            channel_indexes
        ].flatten()

        # the data refers to timestamp (see docstrings) which does not start at
        # 0 but at time_start (see type H docstring)
        first_pos = []
        for channel_id in channel_ids:
            channel_id = int(channel_id)
            first_pos.append(
                min(
                    p[0]
                    for f in self._blocks[block_index][seg_index]["streams"][stream_id][
                        channel_id
                    ]["positions"].values()
                    for p in f
                )
            )

        # now we need to get all file data block indexes for the signal and
        # timestamp we want
        file_chunks = defaultdict(list)
        for i, channel_id in enumerate(channel_ids):
            channel_id = int(channel_id)
            effective_start = i_start + first_pos[i]
            effective_stop = i_stop + first_pos[i]
            for filename, positions in self._blocks[block_index][seg_index]["streams"][
                stream_id
            ][channel_id]["positions"].items():
                file_chunks[filename].extend(
                    [
                        (i, *p)
                        for p in positions
                        if p[0] + p[2] > effective_start and p[0] < effective_stop
                    ]
                )

        # we almost surely loaded more than asked (because the blocks are not
        # contiguous) so we need to know where to cut the results
        slices_channels = []
        for channel_index, channel_id in enumerate(channel_ids):
            channel_id = int(channel_id)
            min_pos = (
                min(
                    p[0]
                    for f in file_chunks.values()
                    for i, *p in f
                    if i == channel_index
                )
                - first_pos[channel_index]
            )
            max_pos = (
                max(
                    p[0] + p[2]
                    for f in file_chunks.values()
                    for i, *p in f
                    if i == channel_index
                )
                - first_pos[channel_index]
            )
            slices_channels.append((min_pos, max_pos))
        min_size = min(s[0] for s in slices_channels)
        max_size = max(s[1] for s in slices_channels)
        sigs = np.ndarray((max_size - min_size, len(channel_ids)), dtype=np.short)

        for filename in file_chunks:
            # we sort by chunk position in the file because we want to optimize
            # IO access and possibly read in sequential access. This is mainly
            # true for hard drives but shouldn't hurt flash memory
            file_chunks[filename].sort(key=lambda x: x[0])
        for filename in file_chunks:
            for channel_index, chunk_index, file_position, chunk_size in file_chunks[
                filename
            ]:
                sig_offset = chunk_index - first_pos[channel_index] - min_size
                sigs[
                    sig_offset : sig_offset + chunk_size, channel_index
                ] = np.frombuffer(
                    self._opened_files[filename]["mmap"],
                    dtype=np.short,<commit_msg>Fix sorting by file chunk position<commit_after>"""
This module implements file reader for AlphaOmega MPX file format version 4.

This module expect default channel names from the AlphaOmega record system (RAW
###, SPK ###, LFP ###, AI ###,…).

This module reads all \*.lsx and \*.mpx files in a directory (not recursively).
Listing files

The specifications are mostly extracted from the "AlphaRS User Manual V1.0.1.pdf"
manual provided with the AlphaRS hardware. The specifications are described in
the chapter 6: ALPHARS FILE FORMAT. See at the end of this file for file format
blocks description.
Some informations missing from the file specifications were kindly provided by
AlphaOmega engineers.

.. note::
    Not a lot of memory optimization effort was put into this module. You should
    expect a big memory usage when loading data with this module

.. todo::
    1. Make TTL events -> epochs? We could describe TTL 0/1 events as epochs:
       when the TTL goes to 1=beginning of epoch 0=end of epoch
    2. decode stram data (block type E) using StreamFormat.h file (not provided
       with this code)
    3. (2.bis) Once we know how to decode Stream AlphaOmega events add them to
       an event channel

Author: <NAME> <<EMAIL>>
"""

import io
import logging
import mmap
import os
import struct

from collections import defaultdict
from datetime import datetime
from itertools import chain
from pathlib import Path, PureWindowsPath

import numpy as np

from .baserawio import (
    BaseRawIO,
    _signal_stream_dtype,
    _signal_channel_dtype,
    _spike_channel_dtype,
    _event_channel_dtype,
)


class AlphaOmegaRawIO(BaseRawIO):
    """
    AlphaOmega MPX file format 4 reader. Handles several blocks and segments.

    A block is a recording defined in a \*.lsx file. AlphaOmega record system
    creates such files every time the software is opened.

    A segment is a continuous record (when record starts/stops).
    If file are not referenced in a \*.lsx file, they are put in the same block.

    :param dirname: folder from where to load the data
    :type dirname: str or Path-like
    :param prune_channels: if True removes the empty channels, defaults to True
    :type prune_channels: bool

    .. warning::
        Because channels must be gathered into coherent streams, channels names
        MUST be the default channel names in AlphaRS software (or Alpha LAB SNR).
    """

    extensions = ["lsx", "mpx"]
    rawmode = "one-dir"

    STREAM_CHANNELS = (
        ("Accelerometer", "ACC", "ACC"),
        ("Spike Train", "SPK", "SPK"),
        ("Local Field Potential", "LFP", "LFP"),
        ("Raw Analog", "RAW", "RAW"),
        ("Analog Input", "AI", "AI"),
        ("User Defined", "UD", "UD"),
    )

    EVENT_CHANNELS = (
        ("TTL", "TTL", "TTL"),
        ("Digital Output", "DOUT", "DOUT"),
        ("User Event", "UD", "UD"),
        ("Stim Marker", "StimMarker", "StimMarker"),
        ("Digital Port", "InPort", "InPort"),
        ("Internal Detection", "Internal Detection", "Internal Detection"),
    )

    def __init__(self, dirname="", prune_channels=True):
        super().__init__(dirname=dirname)
        self.dirname = Path(dirname)
        self._filenames = {}
        if self.dirname.is_dir():
            self._explore_folder()
        else:
            self.logger.error(f"{self.dirname} is not a folder")
        self._prune_channels = prune_channels
        self._opened_files = {}
        self._ignore_unknown_blocks = True  # internal debug property

    def _explore_folder(self):
        """
        Explores an AlphaOmega folder and tries to load the index files (*.lsx)
        to get the data files (*.mpx). If the folder contains only *.mpx files,
        it will load them without splitting them into blocks.
        It does not explores the folder recursively.
        """
        self._filenames = defaultdict(list)

        # first check if there is any *.mpx files
        filenames = list(filter(lambda x: x.is_file(), self.dirname.glob("*.mpx")))
        if not filenames:
            self.logger.error(f"Found no AlphaOmega *.mpx files in {self.dirname}")
        else:
            index_files = list(
                filter(lambda x: x.is_file(), self.dirname.glob("*.lsx"))
            )
            # the following is not very usefull but we should expect the original
            # names to be in increasing lexical order
            index_files.sort()
            if not index_files:
                self.logger.warning(
                    "No *.lsx files found. Will try to load all *.mpx files"
                )
                self._filenames[""].extend(filenames)
            else:
                for index_file in index_files:
                    with open(index_file, "r") as f:
                        for line in f:
                            # a line is a Microsoft Windows path. As we cannot
                            # instanciate a WindowsPath on other OS than MS
                            # Windows, we use the PureWindowsPath class
                            filename = PureWindowsPath(line.strip())
                            filename = self.dirname / filename.name
                            if not filename.is_file():
                                self.logger.warning(f"File {filename} does not exist")
                            else:
                                self._filenames[index_file.name].append(filename)
                                filenames.remove(filename)
                if filenames:
                    self.logger.info(
                        "Some *.mpx files not referenced. Will try to load them."
                    )
                    self._filenames[""].extend(filenames)

    def _source_name(self):
        return str(self.dirname)

    def _read_file_blocks(self, filename, prune_channels=True):
        """Read blocks from AlphaOmega MPX file version 4.

        :param filename: the MPX filename to read blocks from
        :type filename: Path-like object or str
        :param prune_channels: Remove references to channels and ports which
            doesn't contain any data recorded. Be careful when using this option
            with multiple-file data since it could theoretically leads to
            expection raised when data recorded in further files are merged into
            the first file pruned from these channels.
        :type prune_channels: bool
        """
        continuous_analog_channels = {}
        segmented_analog_channels = {}
        digital_channels = {}
        channel_type = {}
        stream_data_channels = {}
        ports = {}
        events = []
        unknown_blocks = []
        with open(filename, "rb") as f:
            length, block_type = HeaderType.unpack(f.read(HeaderType.size))
            # First block is always type h and size 60
            if block_type != b"h":
                try:
                    bt = block_type.decode()
                    self.logger.error("First block must be of type h not type {bt}")
                except UnicodeDecodeError:
                    self.logger.error(
                        (
                            f"First block must be of type h not type "
                            f"{int.from_bytes(bt, 'little')} (int format)"
                        )
                    )
                raise Exception(
                    "First block of AlphaOmega MPX file format must be of type h"
                )
            if not length == 60:
                self.logger.error("First block must be of size 60 (got size {length})")
                raise Exception(
                    "First block of AlphaOmega MPX file format must be of size 60"
                )
            (
                next_block,
                version,
                hour,
                minute,
                second,
                hsecond,
                day,
                month,
                year,
                dayofweek,
                minimum_time,
                maximum_time,
                erase_count,
                map_version,
                application_name,
                resource_version,
                reserved,
            ) = SDataHeader.unpack(f.read(SDataHeader.size))
            if map_version != 4:
                self.logger.error("Only map version 4 is supported")
                raise Exception("Only AlphaOmega MPX file format 4 is supported")
            resource_version = "".join([str(b) for b in resource_version])
            try:
                resource_version = int(resource_version)
            except ValueError:
                self.logger.error(
                    f"m_ResourceVersion should be an integer (got: {resource_version}"
                )
            metadata = {
                "application_version": version,
                "application_name": decode_string(application_name),
                "record_date": datetime(
                    year, month, day, hour, minute, second, 10000 * hsecond
                ),
                "start_time": minimum_time,
                "stop_time": maximum_time,
                "erase_count": erase_count,
                "map_version": map_version,
                "resource_version": resource_version,
                "max_sample_rate": 0,
            }

            pos = 0
            while True:
                pos += length
                f.seek(pos)
                header_bytes = f.read(HeaderType.size)

                length, block_type = HeaderType.unpack(header_bytes)
                if length == 65535:
                    # The stop condition for reading MPX blocks is base on the
                    # length of the block: if the block has length 65535 (or -1
                    # in signed integer value) we know we have reached the end
                    # of the file
                    # We could also check that we are at the end of the file
                    # after this block
                    break

                if block_type == b"h":
                    self.logger.error(
                        "Type h block must exist only at the beginning of file"
                    )
                    raise Exception(
                        "AlphaOmega MPX file format must not have type h block after first block"
                    )

                if block_type == b"2":
                    (
                        next_block,
                        is_analog,
                        is_input,
                        channel_number,
                        *spike_color,
                    ) = Type2Block.unpack(f.read(Type2Block.size))
                    if is_analog and is_input:
                        (
                            mode,
                            amplitude,
                            sample_rate,
                            spike_count,
                            mode_spike,
                        ) = SDefAnalog.unpack(f.read(SDefAnalog.size))
                        mode_spike_mode = (mode_spike & 0b1111000000000000) >> 12
                        mode_spike_master = mode_spike_mode == 1
                        mode_spike_slave = mode_spike_mode == 2
                        mode_spike_linked_channel = mode_spike & 0b0000111111111111
                        mode_spike = {
                            "mode_spike_orig": mode_spike,
                            "mode_spike_master": mode_spike_master,
                            "mode_spike_slave": mode_spike_slave,
                            "mode_spike_linked_channel": mode_spike_linked_channel,
                        }
                        metadata["max_sample_rate"] = max(
                            metadata["max_sample_rate"], sample_rate * 1000
                        )
                        if amplitude <= 5:
                            # This is true for any logging software for map
                            # version >4 (specs say only for ALab SNR but AO
                            # engineer says it's true for any software)
                            amplitude = 1250000 / 2 ** 15
                        if mode == 0:
                            # continuous analog channel definition block
                            assert channel_number not in channel_type
                            channel_type[channel_number] = "continuous_analog"
                            duration, total_gain_100 = SDefContinAnalog.unpack(
                                f.read(SDefContinAnalog.size)
                            )
                            name_length = length - 38
                            name = get_name(f, name_length)
                            assert channel_number not in continuous_analog_channels
                            continuous_analog_channels[channel_number] = {
                                "spike_color": spike_color,
                                "bit_resolution": amplitude,
                                "sample_rate": sample_rate * 1000,
                                "spike_count": spike_count,
                                "mode_spike": mode_spike,
                                "duration": duration,
                                "gain": total_gain_100 / 100,
                                "name": name,
                                "positions": defaultdict(list),
                            }
                        elif mode == 1:
                            # segmented analog channel definition block
                            assert channel_number not in channel_type
                            channel_type[channel_number] = "segmented_analog"
                            (
                                pre_trig_ms,
                                post_trig_ms,
                                level_value,
                                trg_mode,
                                yes_rms,
                                total_gain_100,
                            ) = SDefLevelAnalog.unpack(f.read(SDefLevelAnalog.size))
                            name_length = length - 48
                            name = get_name(f, name_length)
                            assert channel_number not in segmented_analog_channels
                            segmented_analog_channels[channel_number] = {
                                "spike_color": spike_color,
                                "bit_resolution": amplitude,
                                "sample_rate": sample_rate * 1000,
                                "spike_count": spike_count,
                                "mode_spike": mode_spike,
                                "pre_trig_duration": pre_trig_ms / 1000,
                                "post_trig_duration": post_trig_ms / 1000,
                                "level_value": level_value,
                                "trg_mode": trg_mode,
                                "automatic_level_base_rms": yes_rms,
                                "gain": total_gain_100 / 100,
                                "name": name,
                                "positions": defaultdict(list),
                            }
                        else:
                            self.logger.error(
                                f"Unknown type 2 analog block mode: {mode}"
                            )
                            continue
                    elif is_analog == 0 and is_input == 1:
                        # digital input channel definition
                        assert channel_number not in channel_type
                        channel_type[channel_number] = "digital"
                        (
                            sample_rate,
                            save_trigger,
                            duration,
                            prev_status,
                        ) = SDefDigitalInput.unpack(f.read(SDefDigitalInput.size))
                        metadata["max_sample_rate"] = max(
                            metadata["max_sample_rate"], sample_rate * 1000
                        )
                        assert channel_number not in digital_channels
                        name_length = length - 30
                        name = get_name(f, name_length)
                        digital_channels[channel_number] = {
                            "spike_color": spike_color,
                            "sample_rate": sample_rate * 1000,
                            "save_trigger": save_trigger,
                            "duration": duration,
                            "prev_status": prev_status,
                            "name": name,
                            "samples": [],
                        }
                    else:
                        self.logger.error(
                            f"Unknown type 2 block: analog={is_analog}, input={is_input}"
                        )
                        continue
                elif block_type == b"S":
                    # stream data definition block
                    next_block, channel_number, sample_rate = SDefStream.unpack(
                        f.read(SDefStream.size)
                    )
                    metadata["max_sample_rate"] = max(
                        metadata["max_sample_rate"], sample_rate * 1000
                    )
                    assert channel_number not in channel_type
                    channel_type[channel_number] = "stream_data"
                    name_length = length - 18
                    name = get_name(f, name_length)
                    stream_data_channels[channel_number] = {
                        "sample_rate": sample_rate * 1000,
                        "name": name,
                    }
                elif block_type == b"b":
                    # digital input/output port definition block
                    board_number, port, sample_rate, prev_value = SDefPortX.unpack(
                        f.read(SDefPortX.size)
                    )
                    metadata["max_sample_rate"] = max(
                        metadata["max_sample_rate"], sample_rate * 1000
                    )
                    assert port not in channel_type
                    channel_type[port] = "port"
                    name_length = length - 18
                    name = get_name(f, name_length)
                    ports[port] = {
                        "board_number": board_number,
                        "sample_rate": sample_rate * 1000,
                        "prev_value": prev_value,
                        "name": name,
                        "samples": [],
                    }
                elif block_type == b"5":
                    # channel data block
                    unit_number, channel_number = SDataChannel.unpack(
                        f.read(SDataChannel.size)
                    )
                    assert channel_number in channel_type
                    unit_number = int.from_bytes(unit_number, "little")
                    if "analog" in channel_type[channel_number]:
                        data_length = (length - 10) / 2
                        assert int(data_length) == data_length
                        data_length = int(data_length)
                        data_start = f.tell()
                        f.seek(2 * data_length, io.SEEK_CUR)
                        if channel_type[channel_number].startswith("continuous"):
                            assert channel_number in continuous_analog_channels
                            continuous_analog_channels[channel_number]["positions"][
                                filename
                            ].append(
                                (
                                    SDataChannel_sample_id.unpack(
                                        f.read(SDataChannel_sample_id.size)
                                    )[0],
                                    data_start,
                                    data_length,
                                )
                            )
                        elif channel_type[channel_number].startswith("segmented"):
                            assert channel_number in segmented_analog_channels
                            if unit_number > 0 and unit_number <= 4:
                                segmented_analog_channels[channel_number]["positions"][
                                    filename
                                ].append(
                                    (
                                        SDataChannel_sample_id.unpack(
                                            f.read(SDataChannel_sample_id.size)
                                        )[0],
                                        data_start,
                                        data_length,
                                    )
                                )
                            elif unit_number == 0:
                                segmented_analog_channels[channel_number]["positions"][
                                    filename
                                ].append(
                                    (
                                        SDataChannel_sample_id.unpack(
                                            f.read(SDataChannel_sample_id.size)
                                        )[0],
                                        data_start,
                                        data_length,
                                    )
                                )
                            else:
                                self.logger.error(
                                    f"Unknown unit_number={unit_number} in channel data block"
                                )
                                continue
                    elif channel_type[channel_number] == "digital":
                        assert channel_number in digital_channels
                        sample_number, value = SDataChannelDigital.unpack(
                            f.read(SDataChannelDigital.size)
                        )
                        digital_channels[channel_number]["samples"].append(
                            (
                                sample_number,
                                value,
                            )
                        )
                    elif channel_type[channel_number] == "port":
                        assert channel_number in ports
                        # specifications says that for ports it should be "<Lh"
                        # but the data shows clearly "<HL"
                        value, sample_number = SDataChannelPort.unpack(
                            f.read(SDataChannelPort.size)
                        )
                        ports[channel_number]["samples"].append(
                            (
                                sample_number,
                                value,
                            )
                        )
                    else:
                        self.logger.error(
                            f"Unknown channel_type={channel_type[channel_number]} for block type 5"
                        )
                elif block_type == b"E":
                    type_event, timestamp = SAOEvent.unpack(f.read(SAOEvent.size))
                    stream_data_length = length - 8
                    events.append(
                        {
                            "timestamp": timestamp,
                            "stream_data": struct.unpack(
                                f"<{stream_data_length}s", f.read(stream_data_length)
                            )[0],
                        }
                    )
                else:
                    if not self._ignore_unknown_blocks:
                        try:
                            bt = block_type.decode()
                            self.logger.debug(
                                f"Unknown block type: block length: {length}, block_type: {bt}"
                            )
                        except UnicodeDecodeError:
                            self.logger.debug(
                                (
                                    f"Unknown block type: block length: {length}, "
                                    f"block_type: {int.from_bytes(block_type, 'little')} "
                                    "(int format)"
                                )
                            )
                        unknown_blocks.append(
                            {
                                "length": length,
                                "block_type": block_type,
                                "data": f.read(length),
                            }
                        )
        if prune_channels:
            to_remove = []
            for channel_id, channel in continuous_analog_channels.items():
                if not channel["positions"]:
                    to_remove.append(channel_id)
            for channel_id in to_remove:
                del continuous_analog_channels[channel_id]
            to_remove = []
            for channel_id, channel in segmented_analog_channels.items():
                if not channel["positions"]:
                    to_remove.append(channel_id)
            for channel_id in to_remove:
                del segmented_analog_channels[channel_id]
            to_remove = []
            for channel_id, channel in digital_channels.items():
                if not channel["samples"]:
                    to_remove.append(channel_id)
            for channel_id in to_remove:
                del digital_channels[channel_id]
            to_remove = []
            for port_id, port in ports.items():
                if not port["samples"]:
                    to_remove.append(port_id)
            for port_id in to_remove:
                del ports[port_id]
            del to_remove
        return (
            metadata,
            continuous_analog_channels,
            segmented_analog_channels,
            digital_channels,
            channel_type,
            stream_data_channels,
            ports,
            events,
            unknown_blocks,
        )

    def _merge_segments(self, factor_period=1.5):
        """This method merge segments that are the same segment but comes from
        different files. The AlphaOmega data recording system split files with
        a configurable maximum time and size but with a global non-configurable
        maximum size of 1GB.
        Two segment are merged if they validate the following conditions:
            - start of next segment is less than `factor_period` (default 1.5)
            the system's sampling period + end of previous segment. We can't
            check with exactly 1 sampling period because timings are floats
            which are not exact period values
            - The two segment must also have the same record date
            (YEAR-MONTH-DAY). This could potentially lead to errors if
            recordings are longer than a day or run over the night
            - The next segment must have a datetime greater or equal than the
            previous segment (datetime : YEAR-MONTH-DAY-HOUR-MINUTE-SECOND)

        :param factor_period: how many sample period to consider the next segment
            is part of the previous one. This should be 1 <= `factor_period` < 2,
            defaults to 1.5
        :type factor_period: float
        """
        for block in self._blocks:
            segments_to_merge = []
            for segment in block:
                possible_same_segments = [
                    s
                    for s in block
                    if s["metadata"]["record_date"].date()
                    == segment["metadata"]["record_date"].date()
                    and s["metadata"]["record_date"]
                    >= segment["metadata"]["record_date"]
                    and 0
                    <= (s["metadata"]["start_time"] - segment["metadata"]["stop_time"])
                    <= factor_period / segment["metadata"]["max_sample_rate"]
                    and s is not segment
                ]
                assert len(possible_same_segments) in (0, 1)
                if possible_same_segments:
                    existing_merges = [s for segs in segments_to_merge for s in segs]
                    if existing_merges:
                        existing_merges.sort(key=lambda x: x["metadata"]["start_time"])
                        segment = existing_merges[0]
                    segments_to_merge.append((segment, possible_same_segments[0]))
            for segment, segment_to_merge in segments_to_merge:
                assert (
                    segment["metadata"]["max_sample_rate"]
                    == segment_to_merge["metadata"]["max_sample_rate"]
                )
                segment["metadata"]["stop_time"] = segment_to_merge["metadata"][
                    "stop_time"
                ]
                segment["metadata"]["filenames"].extend(
                    segment_to_merge["metadata"]["filenames"]
                )
                for stream in segment_to_merge["streams"]:
                    for channel_id in segment_to_merge["streams"][stream]:
                        assert channel_id in segment["streams"][stream]
                        for f in segment_to_merge["streams"][stream][channel_id][
                            "positions"
                        ]:
                            segment["streams"][stream][channel_id]["positions"][
                                f
                            ].extend(
                                segment_to_merge["streams"][stream][channel_id][
                                    "positions"
                                ][f]
                            )
                for channel_id in segment_to_merge["events"]:
                    # it's possible  that the first segment doesn't record any
                    # port channel data and the port could have been pruned
                    if channel_id in segment["events"]:
                        segment["events"][channel_id]["samples"].extend(
                            segment_to_merge["events"][channel_id]["samples"]
                        )
                    else:
                        segment["events"][channel_id] = segment_to_merge["events"][
                            channel_id
                        ]
                for channel_id in segment_to_merge["spikes"]:
                    assert channel_id in segment["spikes"]
                    for f in segment_to_merge["spikes"][channel_id]["positions"]:
                        segment["spikes"][channel_id]["positions"][f].extend(
                            segment_to_merge["spikes"][channel_id]["positions"][f]
                        )
                segment["ao_events"].extend(segment_to_merge["ao_events"])
                for channel_id in segment_to_merge["stream_data"]:
                    # To be honest, I have no idea what is a
                    # stream_data_channels so let's just overwrite it here
                    segment["stream_data"][channel_id] = segment_to_merge[
                        "stream_data"
                    ][channel_id]
                block.remove(segment_to_merge)

    def _parse_header(self):
        blocks = []
        for index_file, filenames in self._filenames.items():
            segments = []
            continuous_analog_channels = {}
            for i, filename in enumerate(filenames):
                metadata, cac, sac, dc, ct, sd, p, e, ub = self._read_file_blocks(
                    filename, self._prune_channels
                )
                metadata["filenames"] = [filename]
                streams = {}
                for stream_name, channel_name_start, stream_id in self.STREAM_CHANNELS:
                    channels = {
                        i: c
                        for i, c in cac.items()
                        if c["name"].startswith(channel_name_start)
                    }
                    streams[stream_id] = channels
                events = dc.copy()
                events.update(p)
                segment = {
                    "metadata": metadata,
                    "streams": streams,
                    "events": events,
                    "spikes": sac,
                    "ao_events": e,
                    "stream_data": sd,
                }

                segments.append(segment)
            blocks.append(segments)
        self._blocks = blocks
        # We merge segments after having loading all the files because they
        # could be loaded in any order
        self._merge_segments()

        signal_streams = set(
            (stream_name, stream_id)
            for block in self._blocks
            for segment in block
            for stream in segment["streams"]
            for stream_name, _, stream_id in self.STREAM_CHANNELS
            if stream_id == stream and segment["streams"][stream]
        )
        signal_streams = list(signal_streams)
        signal_streams.sort(key=lambda x: x[1])
        signal_streams = np.array(signal_streams, dtype=_signal_stream_dtype)

        signal_channels = set(
            (
                channel["name"],
                channel_id,
                channel["sample_rate"],
                np.dtype(np.short).name,
                "uV",
                channel["gain"] / channel["bit_resolution"],
                0,
                stream,
            )
            for block in self._blocks
            for segment in block
            for stream in segment["streams"]
            for channel_id, channel in segment["streams"][stream].items()
        )
        signal_channels = list(signal_channels)
        signal_channels.sort(key=lambda x: (x[7], x[0]))
        signal_channels = np.array(signal_channels, dtype=_signal_channel_dtype)

        spike_channels = set(
            (
                c["name"],
                i,
                "uV",
                c["gain"] / c["bit_resolution"],
                0,
                round(c["pre_trig_duration"] * c["sample_rate"]),
                c["sample_rate"],
            )
            for block in self._blocks
            for segment in block
            for i, c in segment["spikes"].items()
        )
        spike_channels = list(spike_channels)
        spike_channels.sort(key=lambda x: x[0])
        spike_channels = np.array(spike_channels, dtype=_spike_channel_dtype)

        event_channels = set(
            (event["name"], i, "event")
            for block in self._blocks
            for segment in block
            for i, event in segment["events"].items()
        )
        event_channels = list(event_channels)
        event_channels.sort(key=lambda x: x[1])
        event_channels = np.array(event_channels, dtype=_event_channel_dtype)

        self.header = {}
        self.header["nb_block"] = len(self._filenames)
        self.header["nb_segment"] = [len(segment) for segment in self._blocks]
        self.header["signal_streams"] = signal_streams
        self.header["signal_channels"] = signal_channels
        self.header["spike_channels"] = spike_channels
        self.header["event_channels"] = event_channels

        self._generate_minimal_annotations()

        for block_index, filename in enumerate(self._filenames):
            bl_ann = self.raw_annotations["blocks"][block_index]
            bl_ann["name"] = "Block #{}{}".format(
                block_index, " from lsx file" if filename else ""
            )
            bl_ann["file_origin"] = (
                str(self.dirname / filename) if filename else bl_ann["file_origin"]
            )
            bl_ann["rec_datetime"] = self._blocks[block_index][0]["metadata"][
                "record_date"
            ]
            for seg_index, segment in enumerate(self._blocks[block_index]):
                seg_ann = bl_ann["segments"][seg_index]
                seg_ann["name"] = "Seg #{} Block #{}".format(seg_index, block_index)
                seg_ann["file_origin"] = "\n".join(
                    str(f)
                    for f in self._blocks[block_index][seg_index]["metadata"][
                        "filenames"
                    ]
                )
                seg_ann["rec_datetime"] = self._blocks[block_index][seg_index][
                    "metadata"
                ]["record_date"]
                for c_index, c in enumerate(seg_ann["signals"]):
                    c = c.copy()
                    c["file_origin"] = "\n".join(
                        set(
                            str(f)
                            for channels in self._blocks[block_index][seg_index][
                                "streams"
                            ][c["stream_id"]].values()
                            for f in channels["positions"]
                        )
                    )
                    seg_ann["signals"][c_index] = c
                for e_index, e in enumerate(seg_ann["events"]):
                    e = e.copy()
                    e["file_origin"] = seg_ann["file_origin"]
                    seg_ann["events"][e_index] = e

        # We open files and create mmap objects
        for block in self._filenames:
            for filename in self._filenames[block]:
                if filename not in self._opened_files:
                    self._opened_files[filename] = {}
                    self._opened_files[filename]["file"] = filename.open(mode="rb")
                    self._opened_files[filename]["mmap"] = mmap.mmap(
                        self._opened_files[filename]["file"].fileno(),
                        0,
                        access=mmap.ACCESS_READ,
                    )

    def __del__(self):
        # To be sure we close the file when object is deleted. Be aware that the
        # __del__ method is not necessarily called when interpreter exits so we
        # could still leave file opened. This is probably bad…
        for filename in self._opened_files:
            self._opened_files[filename]["mmap"].close()
            self._opened_files[filename]["file"].close()
        if hasattr(super(), "__del__"):
            super().__del__()

    def _segment_t_start(self, block_index, seg_index):
        return self._blocks[block_index][seg_index]["metadata"]["start_time"]

    def _segment_t_stop(self, block_index, seg_index):
        return self._blocks[block_index][seg_index]["metadata"]["stop_time"]

    def _get_signal_size(self, block_index, seg_index, stream_index):
        stream_id = self.header["signal_streams"][stream_index]["id"]
        sizes = [
            sum(
                sample[2]
                for sample_by_file in channel["positions"].values()
                for sample in sample_by_file
            )
            for channel in self._blocks[block_index][seg_index]["streams"][
                stream_id
            ].values()
        ]
        assert all(s == sizes[0] for s in sizes)
        return sizes[0]

    def _get_signal_t_start(self, block_index, seg_index, stream_index):
        return self._segment_t_start(block_index, seg_index)

    def _get_analogsignal_chunk(
        self, block_index, seg_index, i_start, i_stop, stream_index, channel_indexes
    ):
        if i_start is None:
            i_start = 0
        signal_size = self._get_signal_size(block_index, seg_index, stream_index)
        if i_stop is None or i_stop > signal_size:
            i_stop = signal_size
        stream_id = self.header["signal_streams"][stream_index]["id"]
        mask = self.header["signal_channels"]["stream_id"] == stream_id
        channel_ids = self.header["signal_channels"][mask]["id"][
            channel_indexes
        ].flatten()

        # the data refers to timestamp (see docstrings) which does not start at
        # 0 but at time_start (see type H docstring)
        first_pos = []
        for channel_id in channel_ids:
            channel_id = int(channel_id)
            first_pos.append(
                min(
                    p[0]
                    for f in self._blocks[block_index][seg_index]["streams"][stream_id][
                        channel_id
                    ]["positions"].values()
                    for p in f
                )
            )

        # now we need to get all file data block indexes for the signal and
        # timestamp we want
        file_chunks = defaultdict(list)
        for i, channel_id in enumerate(channel_ids):
            channel_id = int(channel_id)
            effective_start = i_start + first_pos[i]
            effective_stop = i_stop + first_pos[i]
            for filename, positions in self._blocks[block_index][seg_index]["streams"][
                stream_id
            ][channel_id]["positions"].items():
                file_chunks[filename].extend(
                    [
                        (i, *p)
                        for p in positions
                        if p[0] + p[2] > effective_start and p[0] < effective_stop
                    ]
                )

        # we almost surely loaded more than asked (because the blocks are not
        # contiguous) so we need to know where to cut the results
        slices_channels = []
        for channel_index, channel_id in enumerate(channel_ids):
            channel_id = int(channel_id)
            min_pos = (
                min(
                    p[0]
                    for f in file_chunks.values()
                    for i, *p in f
                    if i == channel_index
                )
                - first_pos[channel_index]
            )
            max_pos = (
                max(
                    p[0] + p[2]
                    for f in file_chunks.values()
                    for i, *p in f
                    if i == channel_index
                )
                - first_pos[channel_index]
            )
            slices_channels.append((min_pos, max_pos))
        min_size = min(s[0] for s in slices_channels)
        max_size = max(s[1] for s in slices_channels)
        sigs = np.ndarray((max_size - min_size, len(channel_ids)), dtype=np.short)

        for filename in file_chunks:
            # we sort by chunk position in the file because we want to optimize
            # IO access and possibly read in sequential access. This is mainly
            # true for hard drives but shouldn't hurt flash memory
            file_chunks[filename].sort(key=lambda x: x[2])
        for filename in file_chunks:
            for channel_index, chunk_index, file_position, chunk_size in file_chunks[
                filename
            ]:
                sig_offset = chunk_index - first_pos[channel_index] - min_size
                sigs[
                    sig_offset : sig_offset + chunk_size, channel_index
                ] = np.frombuffer(
                    self._opened_files[filename]["mmap"],
                    dtype=np.short,
{'fraction_non_alphanumeric': 0.07566722292889569, 'fraction_numerical': 0.005321445659404138, 'mean_word_length': 1.211977284460506, 'pattern_counts': {'":': 138, '<': 29, '<?xml version=': 0, '>': 25, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 4, 'word_list_counts': {'cursed_substrings.json': 62, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1978428', 'n_tokens_mistral': 21241, 'n_tokens_neox': 19787, 'n_words': 5252}