"""
Dataset object for loading and unpacking an HDF5 dataset generated by
    sync.py

@author: derricw

Allen Institute for Brain Science
"""
from __future__ import annotations

import datetime
import functools
import io
import logging
import warnings
from collections.abc import Iterable, Sequence
from typing import TYPE_CHECKING, Any, Literal, Union

import h5py
import numpy as np
import numpy.typing as npt
from typing_extensions import Self, TypeAlias

if TYPE_CHECKING:
    import matplotlib.figure as fig
    import matplotlib.pyplot as plt

import npc_sessions.utils as utils

logger = logging.getLogger(__name__)

SyncPathOrDataset: TypeAlias = Union[utils.PathLike, h5py.File, "SyncDataset"]


def get_sync_data(sync_path_or_data: SyncPathOrDataset) -> SyncDataset:
    if isinstance(sync_path_or_data, SyncDataset):
        return sync_path_or_data
    return SyncDataset(sync_path_or_data)


def get_bit(uint_array: npt.NDArray, bit: int) -> npt.NDArray[np.uint8]:
    """
    Returns a bool array for a specific bit in a uint ndarray.

    Parameters
    ----------
    uint_array : (numpy.ndarray)
        The array to extract bits from.
    bit : (int)
        The bit to extract.

    """
    return np.bitwise_and(uint_array, 2**bit).astype(bool).astype(np.uint8)


class SyncDataset:
    """
    A sync dataset.  Contains methods for loading
        and parsing the binary data.

    Parameters
    ----------
    path : str
        Path to HDF5 file.

    Examples
    --------
    >>> dset = SyncDataset('my_h5_file.h5') # doctest: +SKIP
    >>> logger.info(dset.meta_data) # doctest: +SKIP
    >>> dset.stats() # doctest: +SKIP
    >>> dset.close() # doctest: +SKIP

    >>> with SyncDataset('my_h5_file.h5') as d: # doctest: +SKIP
    ...     logger.info(dset.meta_data)
    ...     dset.stats()

    The sync file documentation from MPE can be found at
    sharepoint > Instrumentation > Shared Documents > Sync_line_labels_discussion_2020-01-27-.xlsx  # NOQA E501
    Direct link:
    https://alleninstitute.sharepoint.com/:x:/s/Instrumentation/ES2bi1xJ3E9NupX-zQeXTlYBS2mVVySycfbCQhsD_jPMUw?e=Z9jCwH


    """

    FRAME_KEYS = ("frames", "stim_vsync", "vsync_stim")
    PHOTODIODE_KEYS = ("photodiode", "stim_photodiode")
    OPTOGENETIC_STIMULATION_KEYS = ("LED_sync", "opto_trial")
    EYE_TRACKING_KEYS = (
        "eye_frame_received",  # Expected eye tracking
        # line label after 3/27/2020
        # clocks eye tracking frame pulses (port 0, line 9)
        "cam2_exposure",
        # previous line label for eye tracking
        # (prior to ~ Oct. 2018)
        "eyetracking",
        "eye_cam_exposing",
        "eye_cam_exposure",
        "eye_cam_json",
        "eye_tracking",
    )  # An undocumented, but possible eye tracking line label  # E114
    BEHAVIOR_TRACKING_KEYS = (
        "beh_frame_received",  # Expected behavior line label after 3/27/2020  # NOQA E127
        # clocks behavior tracking frame # E127
        # pulses (port 0, line 8)
        "cam1_exposure",
        "behavior_monitoring",
    )

    DEPRECATED_KEYS: set[str] = set()

    def __init__(self, path) -> None:
        if isinstance(path, self.__class__):
            self = path
        else:
            self.dfile = self.load(path)
        self._check_line_labels()

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}({self.dfile.filename})"

    def _check_line_labels(self) -> None:
        if hasattr(self, "line_labels"):
            deprecated_keys = set(self.line_labels) & self.DEPRECATED_KEYS
            if deprecated_keys:
                warnings.warn(
                    (
                        f"The loaded sync file contains the "
                        f"following deprecated line label keys: "
                        f"{deprecated_keys}. Consider updating the "
                        f"sync file line labels."
                    ),
                    stacklevel=2,
                )
        else:
            warnings.warn(
                ("The loaded sync file has no line labels and may " "not be valid."),
                stacklevel=2,
            )

    def _process_times(self) -> npt.NDArray[np.int64]:
        """
        Preprocesses the time array to account for rollovers.
            This is only relevant for event-based sampling.

        """
        times = self.get_all_events()[:, 0:1].astype(np.int64)

        intervals = np.ediff1d(times, to_begin=0)
        rollovers = np.where(intervals < 0)[0]

        for i in rollovers:
            times[i:] += 4294967296

        return times

    def load(self, path) -> h5py.File:
        """
        Loads an hdf5 sync dataset.

        Parameters
        ----------
        path : str
            Path to hdf5 file.

        """
        if isinstance(path, h5py.File):
            self.dfile = path
        else:
            try:
                self.dfile = h5py.File(path, "r")
            except OSError:
                self.dfile = h5py.File(
                    io.BytesIO(utils.from_pathlike(path).read_bytes()), "r"
                )
        self.meta_data: dict[str, Any] = eval(self.dfile["meta"][()])
        self.line_labels: Sequence[str] = self.meta_data["line_labels"]
        self.times = self._process_times()
        return self.dfile

    @property
    def sample_freq(self) -> float:
        try:
            return float(self.meta_data["ni_daq"]["sample_freq"])
        except KeyError:
            return float(self.meta_data["ni_daq"]["counter_output_freq"])

    def get_bit(self, bit: int) -> npt.NDArray[np.uint8]:
        """
        Returns the values for a specific bit.

        Parameters
        ----------
        bit : int
            Bit to return.
        """
        return get_bit(self.get_all_bits(), bit)

    def get_line(self, line: str | int) -> npt.NDArray[np.uint8]:
        """
        Returns the values for a specific line.

        Parameters
        ----------
        line : str
            Line to return.

        """
        bit = self._line_to_bit(line)
        return self.get_bit(bit)

    def get_bit_changes(self, bit: int) -> npt.NDArray[np.uint8]:
        """
        Returns the first derivative of a specific bit.
            Data points are 1 on rising edges and 255 on falling edges.

        Parameters
        ----------
        bit : int
            Bit for which to return changes.

        """
        bit_array = self.get_bit(bit)
        return np.ediff1d(bit_array, to_begin=0)

    def get_line_changes(self, line: str | int) -> npt.NDArray[np.uint8]:
        """
        Returns the first derivative of a specific line.
            Data points are 1 on rising edges and 255 on falling edges.

        Parameters
        ----------
        line : (str)
            Line name for which to return changes.

        """
        bit = self._line_to_bit(line)
        return self.get_bit_changes(bit)

    def get_all_bits(self) -> npt.NDArray:
        """
        Returns the data for all bits.

        """
        return np.array(self.dfile["data"][()][:, -1])

    def get_all_times(self, units: Literal["samples", "seconds"]) -> npt.NDArray[Any]:
        """
        Returns all counter values.

        Parameters
        ----------
        units : str
            Return times in 'samples' or 'seconds'

        """
        if self.meta_data["ni_daq"]["counter_bits"] == 32:
            times = self.get_all_events()[:, 0]
        else:
            times = self.times
        if units.lower() == "samples":
            return times
        elif units.lower() in ["seconds", "sec", "secs"]:
            freq = self.sample_freq
            return times / freq
        else:
            raise ValueError("Only 'samples' or 'seconds' are valid units.")

    def get_all_events(self) -> npt.NDArray[Any]:
        """
        Returns all counter values and their cooresponding IO state.
        """
        return np.array(self.dfile["data"][()])

    def get_events_by_bit(self, bit: int, units=Literal["seconds", "samples"]):
        """
        Returns all counter values for transitions (both rising and falling)
            for a specific bit.

        Parameters
        ----------
        bit : int
            Bit for which to return events.

        """
        changes = self.get_bit_changes(bit)
        return self.get_all_times(units)[np.where(changes != 0)]

    def get_events_by_line(
        self, line: str | int, units: Literal["samples", "seconds"] = "samples"
    ):
        """
        Returns all counter values for transitions (both rising and falling)
            for a specific line.

        Parameters
        ----------
        line : str
            Line for which to return events.

        """
        line = self._line_to_bit(line)
        return self.get_events_by_bit(line, units)

    def _line_to_bit(self, line: str | int) -> int:
        """
        Returns the bit for a specified line.  Either line name and number is
            accepted.

        Parameters
        ----------
        line : str
            Line name for which to return corresponding bit.

        """
        if type(line) is int:
            return line
        elif type(line) is str:
            return self.line_labels.index(line)
        else:
            raise TypeError("Incorrect line type.  Try a str or int.")

    def _bit_to_line(self, bit: int) -> str:
        """
        Returns the line name for a specified bit.

        Parameters
        ----------
        bit : int
            Bit for which to return the corresponding line name.
        """
        return self.line_labels[bit]

    def get_rising_edges(
        self, line: str | int, units: Literal["samples", "seconds"] = "samples"
    ) -> npt.NDArray[Any]:
        """
        Returns the counter values for the rizing edges for a specific bit or
            line.

        Parameters
        ----------
        line : str
            Line for which to return edges.

        """
        bit = self._line_to_bit(line)
        changes = self.get_bit_changes(bit)
        return self.get_all_times(units)[np.where(changes == 1)]

    def get_edges(
        self,
        kind: Literal["rising", "falling", "all"],
        keys: str | Sequence[str],
        units: Literal["seconds", "samples"],
    ) -> npt.NDArray:
        """Utility function for extracting edge times from a line

        Parameters
        ----------
        kind : One of "rising", "falling", or "all". Should this method return
            timestamps for rising, falling or both edges on the appropriate
            line
        keys : These will be checked in sequence. Timestamps will be returned
            for the first which is present in the line labels
        units : one of "seconds", "samples", or "indices". The returned
            "time"stamps will be given in these units.
        raise_missing : If True and no matching line is found, a KeyError will
            be raised

        Returns
        -------
        An array of edge times. If raise_missing is False and none of the keys
            were found, returns None.

        Raises
        ------
        KeyError : none of the provided keys were found among this dataset's
            line labels

        """
        if kind == "falling":
            fn = self.get_falling_edges
        elif kind == "rising":
            fn = self.get_rising_edges
        elif kind == "all":
            rising = self.get_edges("rising", keys, units)
            falling = self.get_edges("falling", keys, units)
            return np.sort(np.concatenate([rising, falling]))

        if isinstance(keys, str):
            keys = [keys]

        for key in keys:
            try:
                result = fn(key, units)
            except ValueError:
                continue
            else:
                return result
        raise KeyError(f"none of {keys} were found in this dataset's line labels")

    def get_falling_edges(
        self, line: str | int, units: Literal["samples", "seconds"] = "samples"
    ):
        """
        Returns the counter values for the falling edges for a specific bit
            or line.

        Parameters
        ----------
        line : str
            Line for which to return edges.

        """
        bit = self._line_to_bit(line)
        changes = self.get_bit_changes(bit)
        return self.get_all_times(units)[np.where(changes == 255)]

    def get_nearest(
        self,
        source: str,
        target: str,
        source_edge: Literal["rising", "falling"] = "rising",
        target_edge: Literal["rising", "falling"] = "rising",
        direction: Literal["previous", "next"] = "previous",
        units: Literal["indices", "samples", "seconds"] = "indices",
    ) -> npt.NDArray:
        """
        For all values of the source line, finds the nearest edge from the
            target line.

        By default, returns the indices of the target edges.

        Args:
            source (str, int): desired source line
            target (str, int): desired target line
            source_edge [Optional(str)]: "rising" or "falling" source edges
            target_edge [Optional(str): "rising" or "falling" target edges
            direction (str): "previous" or "next". Whether to prefer the
                previous edge or the following edge.
            units (str): "indices"

        """
        source_edges = getattr(self, f"get_{source_edge.lower()}_edges")(
            source.lower(), units="samples"
        )  # E501
        target_edges = getattr(self, f"get_{target_edge.lower()}_edges")(
            target.lower(), units="samples"
        )  # E501
        indices = np.searchsorted(target_edges, source_edges, side="right")
        if direction.lower() == "previous":
            indices[np.where(indices != 0)] -= 1
        elif direction.lower() == "next":
            indices[np.where(indices == len(target_edges))] = -1
        if units in ["indices", "index"]:
            return indices
        elif units == "samples":
            return target_edges[indices]
        elif units in ["sec", "seconds", "second"]:
            return target_edges[indices] / self.sample_freq
        else:
            raise KeyError("Invalid units.  Try 'seconds', 'samples' or 'indices'")

    def get_analog_channel(
        self,
        channel: int,
        start_time: float = 0.0,
        stop_time: float | None = None,
        downsample: int = 1,
    ) -> npt.NDArray:
        """
        Returns the data from the specified analog channel between the
            timepoints.

        Args:
            channel (int, str): desired channel index or label
            start_time (Optional[float]): start time in seconds
            stop_time (Optional[float]): stop time in seconds
            downsample (Optional[int]): downsample factor

        Returns:
            ndarray: slice of data for specified channel

        Raises:
            KeyError: no analog data present

        """
        if isinstance(channel, str):
            channel_index = self.analog_meta_data["analog_labels"].index(channel)
            channel = self.analog_meta_data["analog_channels"].index(channel_index)

        if "analog_data" in self.dfile.keys():
            dset = np.array(self.dfile["analog_data"])
            analog_meta = self.get_analog_meta()
            sample_rate = analog_meta["analog_sample_rate"]
            start = int(start_time * sample_rate)
            if stop_time:
                stop = int(stop_time * sample_rate)
                return dset[start:stop:downsample, channel]
            else:
                return dset[start::downsample, channel]
        else:
            raise KeyError("No analog data was saved.")

    def get_analog_meta(self) -> Any:
        """
        Returns the metadata for the analog data.
        """
        if "analog_meta" in self.dfile.keys():
            return eval(self.dfile["analog_meta"].value)
        else:
            raise KeyError("No analog data was saved.")

    @property
    def analog_meta_data(self) -> Any:
        return self.get_analog_meta()

    def line_stats(self, line, print_results=True) -> dict[str, Any] | None:
        """
        Quick-and-dirty analysis of a bit.

        ##TODO: Split this up into smaller functions.

        """
        # convert to bit
        bit = self._line_to_bit(line)

        # get the bit's data
        bit_data = self.get_bit(bit)
        total_data_points = len(bit_data)

        # get the events
        events = self.get_events_by_bit(bit, units="samples")
        total_events = len(events)

        # get the rising edges
        rising = self.get_rising_edges(bit)
        total_rising = len(rising)

        # get falling edges
        falling = self.get_falling_edges(bit)
        total_falling = len(falling)

        # get labels
        label = self.line_labels[line]

        if total_events <= 0:
            if print_results:
                logger.info("*" * 70)
                logger.info("No events on line: %s" % line)
                logger.info("*" * 70)
            return None
        elif total_events <= 10:
            if print_results:
                logger.info("*" * 70)
                logger.info("Sparse events on line: %s" % line)
                logger.info("Rising: %s" % total_rising)
                logger.info("Falling: %s" % total_falling)
                logger.info("*" * 70)
            return {
                "line": line,
                "bit": bit,
                "total_rising": total_rising,
                "total_falling": total_falling,
                "avg_freq": None,
                "duty_cycle": None,
            }
        else:
            # period
            period = self.period(line)

            avg_period = period["avg"]
            max_period = period["max"]
            min_period = period["min"]
            period_sd = period["sd"]

            # freq
            avg_freq = self.frequency(line)

            # duty cycle
            duty_cycle = self.duty_cycle(line)

            if print_results:
                logger.info("*" * 70)

                logger.info("Quick stats for line: %s" % line)
                logger.info("Label: %s" % label)
                logger.info("Bit: %i" % bit)
                logger.info("Data points: %i" % total_data_points)
                logger.info("Total transitions: %i" % total_events)
                logger.info("Rising edges: %i" % total_rising)
                logger.info("Falling edges: %i" % total_falling)
                logger.info("Average period: %s" % avg_period)
                logger.info("Minimum period: %s" % min_period)
                logger.info("Max period: %s" % max_period)
                logger.info("Period SD: %s" % period_sd)
                logger.info("Average freq: %s" % avg_freq)
                logger.info("Duty cycle: %s" % duty_cycle)

                logger.info("*" * 70)

            return {
                "line": line,
                "label": label,
                "bit": bit,
                "total_data_points": total_data_points,
                "total_events": total_events,
                "total_rising": total_rising,
                "total_falling": total_falling,
                "avg_period": avg_period,
                "min_period": min_period,
                "max_period": max_period,
                "period_sd": period_sd,
                "avg_freq": avg_freq,
                "duty_cycle": duty_cycle,
            }

    def period(
        self, line: str | int, edge: Literal["rising", "falling"] = "rising"
    ) -> dict[str, Any]:
        """
        Returns a dictionary with avg, min, max, and st of period for a line.
        """
        bit = self._line_to_bit(line)

        if edge.lower() == "rising":
            edges = self.get_rising_edges(bit)
        elif edge.lower() == "falling":
            edges = self.get_falling_edges(bit)
        else:
            raise ValueError(
                f'edge should be one of ("rising", "falling") not {edge!r}'
            )

        if len(edges) > 2:
            timebase_freq = self.meta_data["ni_daq"]["counter_output_freq"]
            avg_period = np.mean(np.ediff1d(edges[1:])) / timebase_freq
            max_period = np.max(np.ediff1d(edges[1:])) / timebase_freq
            min_period = np.min(np.ediff1d(edges[1:])) / timebase_freq
            period_sd = np.std(avg_period)

        else:
            raise IndexError("Not enough edges for period: %i" % len(edges))

        return {
            "avg": avg_period,
            "max": max_period,
            "min": min_period,
            "sd": period_sd,
        }

    def frequency(
        self, line: str | int, edge: Literal["rising", "falling"] = "rising"
    ) -> float:
        """
        Returns the average frequency of a line.
        """

        period = self.period(line, edge)
        return 1.0 / period["avg"]

    def duty_cycle(self, line: str | int) -> Literal["fix me"]:
        """
        Doesn't work right now.  Freezes python for some reason.

        Returns the duty cycle of a line.

        """
        return "fix me"
        bit = self._line_to_bit(line)

        rising = self.get_rising_edges(bit)
        falling = self.get_falling_edges(bit)

        total_rising = len(rising)
        total_falling = len(falling)

        if total_rising > total_falling:
            rising = rising[:total_falling]
        elif total_rising < total_falling:
            falling = falling[:total_rising]
        else:
            pass

        if rising[0] < falling[0]:
            # line starts low
            high = falling - rising
        else:
            # line starts high
            high = np.concatenate(
                falling, self.get_all_events()[-1, 0]
            ) - np.concatenate(0, rising)

        total_high_time = np.sum(high)
        all_events = self.get_events_by_bit(bit)
        total_time = all_events[-1] - all_events[0]
        return 1.0 * total_high_time / total_time

    @property
    def stats(self) -> list[dict[str, Any]]:
        """
        Quick-and-dirty analysis of all bits.  Prints a few things about each
            bit where events are found.
        """
        bits = []
        for i in range(32):
            bits.append(self.line_stats(i, print_results=True))
        active_bits = [x for x in bits if x is not None]
        logger.info("Active bits: ", len(active_bits))
        for bit in active_bits:
            logger.info("*" * 70)
            logger.info("Bit: %i" % bit["bit"])
            logger.info("Label: %s" % self.line_labels[bit["bit"]])
            logger.info("Rising edges: %i" % bit["total_rising"])
            logger.info("Falling edges: %i" % bit["total_falling"])
            logger.info("Average freq: %s" % bit["avg_freq"])
            logger.info("Duty cycle: %s" % bit["duty_cycle"])
        logger.info("*" * 70)
        return active_bits

    @property
    def start_time(self) -> datetime.datetime:
        return datetime.datetime.fromisoformat(self.meta_data["start_time"])

    @functools.cached_property
    def stim_running_edges(
        self,
    ) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]:
        stim_running_rising_edges = self.get_rising_edges(
            "stim_running", units="seconds"
        )
        stim_running_falling_edges = self.get_falling_edges(
            "stim_running", units="seconds"
        )

        if any(stim_running_rising_edges) and any(stim_running_falling_edges):
            if stim_running_rising_edges[0] > stim_running_falling_edges[0]:
                stim_running_falling_edges[1:]
            if stim_running_falling_edges[-1] < stim_running_rising_edges[-1]:
                stim_running_falling_edges = np.concatenate(
                    [stim_running_falling_edges, [self.total_seconds]]
                )
        assert len(stim_running_rising_edges) == len(stim_running_falling_edges)
        return stim_running_rising_edges, stim_running_falling_edges

    def filter_on_stim_running(
        self, data: npt.NDArray[np.floating]
    ) -> npt.NDArray[np.floating]:
        """Filter data to only include times when stim_running is high.

        Data must be in seconds relative to first sample."""
        if self.stim_running_edges[0].size == 0:
            return data
        mask = [False] * len(data)
        for on, off in zip(*self.stim_running_edges):
            mask |= (data >= on) & (data <= off)

        return data[mask]

    @property
    def total_seconds(self) -> float:
        return self.meta_data["total_samples"] / self.sample_freq

    @functools.cached_property
    def vsync_times_in_blocks(self) -> tuple[npt.NDArray[np.floating], ...]:
        """Blocks of vsync falling edge times, in seconds relative to first
        sample: one block per stimulus.
        """
        vsync_rising_edges: npt.NDArray[np.floating] = self.get_rising_edges(
            "vsync_stim", units="seconds"
        )
        vsync_falling_edges: npt.NDArray[np.floating] = self.get_falling_edges(
            "vsync_stim", units="seconds"
        )
        # ensure first vsync is rising
        vsync_falling_edges = (
            vsync_falling_edges
            if vsync_rising_edges[0] < vsync_falling_edges[0]
            else vsync_falling_edges[1:]
        )

        vsync_times_in_blocks = list(
            reshape_into_blocks(vsync_falling_edges, min_gap=1.0)
        )

        block_lengths = np.array([len(block) for block in vsync_times_in_blocks])
        logger.info(
            f"Found {len(vsync_times_in_blocks)} blocks of vsync events with lengths {block_lengths}"
        )
        stim_running_rising_edges, stim_running_falling_edges = self.stim_running_edges
        if any(stim_running_rising_edges) and any(stim_running_falling_edges):
            assert len(stim_running_rising_edges) == len(vsync_times_in_blocks)
            for idx, block in enumerate(vsync_times_in_blocks):
                vsync_times_in_blocks[idx] = self.filter_on_stim_running(block)

        assert all(block.size > 0 for block in vsync_times_in_blocks)
        return tuple(vsync_times_in_blocks)

    @functools.cached_property
    def expected_diode_flip_rate(self) -> int:
        """Best-guess at what the diode flip period should be, e.g. 1 s for
        MPE/pipeline recordings, 1/60 s for Sam's TaskControl scripts."""
        med = np.median(
            np.diff(self.get_edges("all", "stim_photodiode", units="seconds"))
        )
        diode_assymmetry = 0.05  # s
        for period in (1 / 60, 1):
            if 0.9 * period - diode_assymmetry < med < 1.1 * period + diode_assymmetry:
                return int(1 / period)
        raise ValueError(f"Unexpected diode flip period: {med} sec")

    @functools.cached_property
    def expected_frame_display_rate(self) -> int:
        """Best-guess at what the screen display rate should be. Currently
        [2023] should only be 60 fps."""
        med = np.median(np.diff(self.get_falling_edges("vsync_stim", units="seconds")))
        for period in (1 / 60, 1 / 120, 1 / 144, 1 / 300):
            if 0.9 * period < med < 1.1 * period:
                return int(1 / period)
        raise ValueError(f"Unexpected vsync period: {med} sec")

    @functools.cached_property
    def frame_display_time_blocks(self) -> tuple[npt.NDArray[np.floating], ...]:
        """Blocks of adjusted diode times: one block per stimulus."""
        vsync_times_in_blocks = self.vsync_times_in_blocks

        diode_rising_edges = self.get_rising_edges("stim_photodiode", units="seconds")
        diode_falling_edges = self.get_falling_edges("stim_photodiode", units="seconds")
        assert abs(len(diode_rising_edges) - len(diode_falling_edges)) < 2

        diode_rising_edges_in_blocks = reshape_into_blocks(
            self.filter_on_stim_running(diode_rising_edges), min_gap=1.0
        )
        diode_falling_edges_in_blocks = reshape_into_blocks(
            self.filter_on_stim_running(diode_falling_edges), min_gap=1.0
        )

        if any(
            len(diode_edge_blocks) < len(vsync_times_in_blocks)
            for diode_edge_blocks in (
                diode_rising_edges_in_blocks,
                diode_falling_edges_in_blocks,
            )
        ):
            # too few blocks of diode flips - affected some early DRPilot
            # sessions with near-constant diode flips between stimulus
            # blocks, and no stim-running signal to filter them

            # - vsyncs necessarily precede diode flips
            # - split existing diode flip blocks at the first vsync time for
            #   each vsync block
            # - we'll clean up any excess diode flips at the end of the block later
            # to see the problem compare these two lists:
            # [(v[0], v[-1]) for v in vsync_times_in_blocks]
            # [(v[0], v[-1]) for v in diode_falling_edges_in_blocks]
            for v_start in (v[0] for v in vsync_times_in_blocks):
                assert len(diode_rising_edges_in_blocks) == len(
                    diode_falling_edges_in_blocks
                )
                new_rising_edges: list[npt.NDArray] = []
                new_falling_edges: list[npt.NDArray] = []
                for d_idx, (d_start, d_stop) in enumerate(
                    (d[0], d[-1]) for d in diode_falling_edges_in_blocks
                ):
                    if d_start < v_start < d_stop and (
                        split := np.searchsorted(
                            diode_falling_edges_in_blocks[d_idx], v_start
                        )
                    ) not in (0, len(diode_falling_edges_in_blocks[d_idx])):
                        for new, old in zip(
                            (new_rising_edges, new_falling_edges),
                            (
                                diode_rising_edges_in_blocks[d_idx],
                                diode_falling_edges_in_blocks[d_idx],
                            ),
                        ):
                            new.extend((old[:split], old[split:]))
                    else:
                        for new, old in zip(
                            (new_rising_edges, new_falling_edges),
                            (
                                diode_rising_edges_in_blocks[d_idx],
                                diode_falling_edges_in_blocks[d_idx],
                            ),
                        ):
                            new.append(old)
                diode_rising_edges_in_blocks, diode_falling_edges_in_blocks = (
                    tuple(new_rising_edges),
                    tuple(new_falling_edges),
                )

        if any(
            len(diode_edge_blocks) > len(vsync_times_in_blocks)
            for diode_edge_blocks in (
                diode_rising_edges_in_blocks,
                diode_falling_edges_in_blocks,
            )
        ):
            # more blocks of diode flips than blocks of vsyncs
            def is_mismatch(edges, block) -> bool:
                return edges[-1] < block[0]

            for idx, vsync_block in enumerate(vsync_times_in_blocks):
                # work through blocks in order,
                # discard blocks with diode flips that don't match any vsyncs
                while is_mismatch(diode_rising_edges_in_blocks[idx], vsync_block):
                    diode_rising_edges_in_blocks = tuple(
                        block
                        for i, block in enumerate(diode_rising_edges_in_blocks)
                        if i != idx
                    )
                while is_mismatch(diode_falling_edges_in_blocks[idx], vsync_block):
                    diode_falling_edges_in_blocks = tuple(
                        block
                        for i, block in enumerate(diode_falling_edges_in_blocks)
                        if i != idx
                    )

        frame_display_time_blocks: list[npt.NDArray[np.floating]] = []
        for block_idx, (vsyncs, rising, falling) in enumerate(
            zip(
                vsync_times_in_blocks,
                diode_rising_edges_in_blocks,
                diode_falling_edges_in_blocks,
            )
        ):
            # keep flips after first vsync + an empirically determined min latency
            # between vsync and screen update
            MIN_VSYNC_DIODE_FLIP_SEPARATION_SEC = 0.018  # 0.22 is typical
            falling = falling[
                falling > (vsyncs[0] + MIN_VSYNC_DIODE_FLIP_SEPARATION_SEC)
            ]
            rising = rising[rising > (vsyncs[0] + MIN_VSYNC_DIODE_FLIP_SEPARATION_SEC)]

            # keep flips only up to a certain time after the last vsync
            MAX_VSYNC_DIODE_FLIP_SEPARATION_SEC = 0.035
            falling = falling[
                falling < (vsyncs[-1] + MAX_VSYNC_DIODE_FLIP_SEPARATION_SEC)
            ]
            rising = rising[rising < (vsyncs[-1] + MAX_VSYNC_DIODE_FLIP_SEPARATION_SEC)]

            diode_flips = np.sort(np.concatenate((rising, falling)))

            short_interval_threshold = (
                0.1 * 1 / 60
            )  # at 60 fps we should never have diode-flip intervals this small: consider them anomalous

            def short_interval_indices(diode_flips):
                intervals = np.diff(diode_flips)
                return np.where(intervals < short_interval_threshold)[0]

            while any(short_interval_indices(diode_flips)):
                indices = short_interval_indices(diode_flips)
                diode_flips = np.delete(diode_flips, slice(indices[0], indices[0] + 2))

            one_diode_flip_per_vsync: bool = round(
                np.mean(np.diff(diode_flips)), 1
            ) == round(
                np.mean(np.diff(vsyncs)), 1
            )  # diode flip freq ~= vsync freq

            if one_diode_flip_per_vsync:
                # we have to assume that vsyncs are correct at this point (no
                # extras, none missing)

                def score(diode_flips, vsyncs) -> float:
                    """Similarity between diode and vsync intervals - higher is
                    more similar"""
                    common_len = min([len(vsyncs), len(diode_flips)])
                    return (
                        1
                        / np.mean(
                            np.abs(
                                np.diff(vsyncs[:common_len])
                                - np.diff(
                                    adjust_diode_flip_intervals(diode_flips)[
                                        :common_len
                                    ]
                                )
                            )
                        ).item()
                    )

                def corr_score(diode_flips, vsyncs) -> float:
                    """Similarity between diode and vsync intervals - higher is
                    more similar"""
                    common_len = min([len(vsyncs), len(diode_flips)])
                    return np.correlate(
                        np.diff(vsyncs[:common_len]),
                        np.diff(adjust_diode_flip_intervals(diode_flips)[:common_len]),
                    ).max()

                def median_diff(diode_flips, vsyncs) -> float:
                    common_len = min([len(vsyncs), len(diode_flips)])
                    return np.median(diode_flips[:common_len] - vsyncs[:common_len])

                counter = 0
                while (
                    median_diff(diode_flips, vsyncs)
                    > MAX_VSYNC_DIODE_FLIP_SEPARATION_SEC
                ) and (
                    MIN_VSYNC_DIODE_FLIP_SEPARATION_SEC
                    < median_diff(diode_flips, vsyncs[1:])
                    < median_diff(diode_flips, vsyncs)
                ):
                    if counter > 2:  #!  this may need to be dereased to allow only 1
                        raise ValueError(
                            f"Added two missed diode flips at start already, the max we can explain by pre-stim background & first diode square being similar luminance. This is a different problem: sync {self.start_time=}."
                        )
                    logger.debug("Missing first diode flip")
                    diode_flips = add_missing_diode_flip_at_stim_onset(
                        diode_flips, vsyncs
                    )
                    counter += 1

                while (
                    median_diff(diode_flips, vsyncs)
                    < MIN_VSYNC_DIODE_FLIP_SEPARATION_SEC
                ) and (
                    median_diff(diode_flips, vsyncs)
                    < median_diff(diode_flips[1:], vsyncs)
                    < MAX_VSYNC_DIODE_FLIP_SEPARATION_SEC
                ):
                    logger.debug("Removing extra first diode flip")
                    diode_flips = diode_flips[1:]

                if len(diode_flips) < len(vsyncs):
                    if (
                        is_last_vsync_close_to_end_of_sync := abs(
                            self.total_seconds - vsyncs[-1]
                        )
                        < MAX_VSYNC_DIODE_FLIP_SEPARATION_SEC
                    ) or (
                        any(self.stim_running_edges[1])
                        and (
                            is_last_vsync_close_to_last_stim_running_rising := abs(
                                self.stim_running_edges[1][block_idx] - vsyncs[-1]
                            )
                            < MAX_VSYNC_DIODE_FLIP_SEPARATION_SEC
                        )
                    ):
                        logger.debug(
                            "Missing last diode flips - truncated by end of sync of stim running"
                        )
                        diode_flips = add_missing_diode_flips_for_truncated_sync(
                            diode_flips, vsyncs
                        )

                if len(diode_flips) > len(vsyncs):
                    diode_flips = discard_erroneous_diode_flips_at_stim_offset(
                        diode_flips, vsyncs
                    )

                if len(diode_flips) != len(vsyncs):
                    import matplotlib.pyplot as plt

                    fig, _ = plt.subplots(1, 2)
                    vs = vsyncs
                    dvs = diode_flips
                    labels = []

                    # focus on the end of vsyncs if they occur well before the stim-TTL offset
                    # x0, x1 = min(soff - 2, vs[-1] - 2), min(soff + 1, vs[-1]
                    # + 1)
                    for idx, ax in enumerate(fig.axes):
                        padding = 5  # seconds
                        start_time = [min(vsyncs) - padding, max(vsyncs) - padding][idx]
                        end_time = [min(vsyncs) + padding, max(vsyncs) + padding][idx]

                        self.plot_bit(
                            4,
                            start_time=start_time,
                            end_time=end_time,
                            axes=ax,
                            auto_show=False,
                        )
                        labels.append("diode-measured sync square")
                        self.plot_bit(
                            5,
                            start_time=start_time,
                            end_time=end_time,
                            axes=ax,
                            auto_show=False,
                        )
                        labels.append("stim running")
                        ax.plot(vs, 0.5 * np.ones_like(vs), "|")
                        labels.append("stim vsyncs")
                        ax.plot(dvs, 0.5 * np.ones_like(dvs), "|", ms=20)
                        labels.append("diode_flips")
                        ax.set_xlim([start_time, end_time])
                        fig.axes[0].legend(
                            labels,
                            fontsize=8,
                            loc="upper center",
                            bbox_to_anchor=(0.5, 1.05),
                            ncol=len(labels),
                            fancybox=True,
                        )

                    # self.plot_lines(["stim_photodiode"])
                    plt.suptitle(
                        f"{len(diode_flips) = }, {len(vsyncs) = }, {block_idx = }"
                    )
                    plt.show()

                    raise IndexError(
                        f"Mismatch in stim {block_idx = }: {len(diode_flips) = }, {len(vsyncs) = }"
                    )

            else:
                pass
                # TODO adjust frametimes with diode data when flip is every 1 s

            diode_flips = adjust_diode_flip_intervals(diode_flips)

            AVERAGE_SCREEN_REFRESH_TIME = 0.008
            """Screen refreshes in stages top-to-bottom, total 16 ms measured by
            Corbett: use mid-point"""

            frametimes = diode_flips + AVERAGE_SCREEN_REFRESH_TIME
            frame_display_time_blocks.append(frametimes)

        assert len(frame_display_time_blocks) == len(self.vsync_times_in_blocks)
        return tuple(frame_display_time_blocks)

    def plot_all(
        self,
        start_time: float,
        stop_time: float | None = None,
        auto_show: bool = True,
    ) -> None:
        """
        Plot all active bits.

        Yikes.  Come up with a better way to show this.

        """
        import matplotlib.pyplot as plt

        for bit in range(32):
            if len(self.get_events_by_bit(bit)) > 0:
                self.plot_bit(
                    bit,
                    start_time,
                    stop_time,
                    auto_show=False,
                )
        if auto_show:
            plt.show()

    def plot_bits(
        self,
        bits: Sequence[int],
        start_time: float = 0.0,
        end_time: float | None = None,
        auto_show: bool = True,
    ) -> tuple[fig.Figure, plt.Axes | Iterable[plt.Axes]]:
        """
        Plots a list of bits.
        """
        import matplotlib.pyplot as plt

        subplots = len(bits)
        f, axes = plt.subplots(subplots, sharex=True, sharey=True)
        if not isinstance(axes, Iterable):
            axes = [axes]

        for bit, ax in zip(bits, axes):
            self.plot_bit(bit, start_time, end_time, auto_show=False, axes=ax)
        # f.set_size_inches(18, 10, forward=True)
        f.subplots_adjust(hspace=0)

        if auto_show:
            plt.show()

        return f, axes

    def plot_bit(
        self,
        bit,
        start_time: float = 0.0,
        end_time: float | None = None,
        auto_show: bool = True,
        axes=None,
        name="",
    ) -> tuple[fig.Figure, plt.Axes | Iterable[plt.Axes]]:
        """
        Plots a specific bit at a specific time period.
        """
        import matplotlib.pyplot as plt

        times = self.get_all_times(units="seconds")
        if not end_time:
            end_time = 2**32

        window = (times < end_time) & (times > start_time)

        if axes:
            ax = axes
        else:
            ax = plt

        if not name:
            name = self._bit_to_line(bit)
        if not name:
            name = str(bit)

        bit = self.get_bit(bit)
        ax.step(times[window], bit[window], where="post")
        if hasattr(ax, "set_ylim"):
            ax.set_ylim(-0.1, 1.1)
        else:
            axes_obj = plt.gca()
            axes_obj.set_ylim(-0.1, 1.1)
        # ax.set_ylabel('Logic State')
        # ax.yaxis.set_ticks_position('none')
        plt.setp(ax.get_yticklabels(), visible=False)
        ax.set_xlabel("time (seconds)")
        ax.legend([name])

        if auto_show:
            plt.show()

        return plt.gcf()

    def plot_line(
        self,
        line,
        start_time: float = 0.0,
        end_time: float | None = None,
        auto_show: bool = True,
    ) -> None:
        """
        Plots a specific line at a specific time period.
        """
        import matplotlib.pyplot as plt

        bit = self._line_to_bit(line)
        self.plot_bit(bit, start_time, end_time, auto_show=False)

        # plt.legend([line])
        if auto_show:
            plt.show()

    def plot_lines(
        self,
        lines,
        start_time: float = 0.0,
        end_time: float | None = None,
        auto_show: bool = True,
    ) -> tuple[fig.Figure, plt.Axes | Iterable[plt.Axes]]:
        """
        Plots specific lines at a specific time period.
        """
        import matplotlib.pyplot as plt

        bits = []
        for line in lines:
            bits.append(self._line_to_bit(line))
        f, axes = self.plot_bits(
            bits,
            start_time,
            end_time,
            auto_show=False,
        )

        plt.subplots_adjust(left=0.025, right=0.975, bottom=0.05, top=0.95)
        if auto_show:
            plt.show()

        return f, axes

    @property
    def stim_onsets(self) -> npt.NDArray[np.floating]:
        if any(stim_running := self.get_rising_edges("stim_running", units="seconds")):
            return stim_running
        return np.array([block[0] for block in self.vsync_times_in_blocks])

    @property
    def stim_offsets(self) -> npt.NDArray[np.floating]:
        if any(stim_running := self.get_falling_edges("stim_running", units="seconds")):
            return stim_running
        return np.array([block[-1] for block in self.frame_display_time_blocks])

    def plot_stim_onsets(self) -> tuple[fig.Figure, list[plt.Axes]]:
        import matplotlib.pyplot as plt

        # plot beginning of stims
        fig, _ = plt.subplots(len(self.stim_onsets))
        axes = fig.axes
        fig.suptitle("Stim onsets")
        for ind, (son, vs, dvs) in enumerate(
            zip(
                self.stim_onsets,
                self.vsync_times_in_blocks,
                self.frame_display_time_blocks,
            )
        ):
            labels = []
            axes[ind].plot(vs, 0.5 * np.ones(len(vs)), "|")
            labels.append("stim vsyncs")
            axes[ind].plot(dvs, 0.5 * np.ones(len(dvs)), "|", ms=20)
            labels.append("frame display (estimated)")

            # focus on the start of vsyncs if they occur well after the stim-TTL onset
            x0, x1 = max(son - 1, vs[0] - 1), max(son + 2, vs[0] + 2)

            self.plot_bit(4, x0, x1, axes=axes[ind], auto_show=False)
            labels.append("diode-measured sync square")
            self.plot_bit(5, x0, x1, axes=axes[ind], auto_show=False)
            labels.append("stim running or ends of block")
            axes[ind].set_xlim([x0, x1])
            axes[ind].legend(
                labels,
                fontsize=8,
                loc="upper center",
                bbox_to_anchor=(0.5, 1.05),
                ncol=len(labels),
                fancybox=True,
            )
            if len(fig.axes) > 1:
                axes[ind].set_title(f"visual stim {ind}", fontsize=8)
                legend = axes[ind].get_legend()
                if ind > 0 and legend is not None:
                    legend.remove()
        return fig, axes

    def plot_stim_offsets(self) -> tuple[fig.Figure, list[plt.Axes]]:
        import matplotlib.pyplot as plt

        # plot end of stims
        fig, _ = plt.subplots(len(self.stim_offsets))
        axes = fig.axes

        fig.suptitle("Stim offsets")
        for ind, (soff, vs, dvs) in enumerate(
            zip(
                self.stim_offsets,
                self.vsync_times_in_blocks,
                self.frame_display_time_blocks,
            )
        ):
            labels = []
            axes[ind].plot(vs, 0.5 * np.ones(len(vs)), "|")
            labels.append("stim vsyncs")
            axes[ind].plot(dvs, 0.5 * np.ones(len(dvs)), "|", ms=20)
            labels.append("frame display (estimated)")

            # focus on the end of vsyncs if they occur well before the stim-TTL offset
            x0, x1 = min(soff - 2, vs[-1] - 2), min(soff + 1, vs[-1] + 1)

            self.plot_bit(4, x0, x1, axes=axes[ind], auto_show=False)
            labels.append("diode-measured sync square")
            self.plot_bit(5, x0, x1, axes=axes[ind], auto_show=False)
            labels.append("stim running or ends of block")
            axes[ind].set_xlim([x0, x1])
            axes[ind].legend(
                labels,
                fontsize=8,
                loc="upper center",
                bbox_to_anchor=(0.5, 1.05),
                ncol=len(labels),
                fancybox=True,
            )
            if len(fig.axes) > 1:
                axes[ind].set_title(f"visual stim {ind}", fontsize=8)
                legend = axes[ind].get_legend()
                if ind > 0 and legend is not None:
                    legend.remove()
        return fig, axes

    def plot_diode_measured_sync_square_flips(
        self,
    ) -> tuple[fig.Figure, list[plt.Axes]]:
        """Plot the diode-measured sync-square changes that should occur every 1 s while stim is running."""
        import matplotlib.pyplot as plt

        stim_ons, stim_offs = self.stim_onsets, self.stim_offsets

        # we want the diode flips that occur after the stim-running TTL goes high
        # and after the vsyncs start
        all_diode_flips = np.concatenate(
            [
                self.get_rising_edges("stim_photodiode", units="seconds"),
                self.get_falling_edges("stim_photodiode", units="seconds"),
            ]
        )
        all_vsyncs = self.get_falling_edges("vsync_stim", units="seconds")

        frequency = self.expected_diode_flip_rate
        expected_period = 1 / frequency

        # get the intervals in parts (one for each stimulus)
        diode_flips_per_stim = []
        for son, soff in zip(stim_ons, stim_offs):
            # get the vsyncs that occur during this stimulus
            vsyncs = all_vsyncs[np.where((all_vsyncs > son) & (all_vsyncs < soff))]
            # get the diode flips that occur during this stimulus, while vsyncs are occurring
            diode_flips = all_diode_flips[
                np.where(
                    (all_diode_flips > son)
                    & (all_diode_flips < soff)
                    & (all_diode_flips > vsyncs[0])
                    & (all_diode_flips < vsyncs[-1])
                )
            ]
            diode_flips_per_stim.append(sorted(diode_flips))

        num_diode_flips_per_stim = np.array([len(_) for _ in diode_flips_per_stim])
        # add ` width_ratios=num_diode_flips/min(num_diode_flips)``
        fig, _ = plt.subplots(
            1,
            len(stim_ons),
            sharey=True,
            gridspec_kw={
                "width_ratios": num_diode_flips_per_stim / min(num_diode_flips_per_stim)
            },
        )
        fig.suptitle(
            f"diode-measured sync-square flip intervals, {expected_period = } s"
        )
        y_deviations_from_expected_period: list[float] = []
        for idx, (ax, d) in enumerate(zip(fig.axes, diode_flips_per_stim)):
            # add horizontal line at expected period
            ax.axhline(expected_period, linewidth=0.5, c="k", linestyle="--", alpha=0.3)
            plt.sca(ax)
            intervals = np.diff(d)
            times = np.diff(d) / 2 + d[:-1]  # plot at mid-point of interval
            markerline, stemline, baseline = plt.stem(
                times, intervals, bottom=expected_period
            )
            plt.setp(stemline, linewidth=0.5, alpha=0.3)
            plt.setp(markerline, markersize=0.5, alpha=0.8)
            plt.setp(baseline, visible=False)

            y_deviations_from_expected_period.append(max(intervals - expected_period))
            y_deviations_from_expected_period.append(max(expected_period - intervals))
            if len(fig.axes) > 1:
                ax.set_title(f"stim {idx}", fontsize=8)
            ax.set_xlabel("time (s)")
            if idx == 0:
                ax.set_ylabel("flip interval (s)")
            ax.set_xlim(min(d) - 20, max(d) + 20)

        for ax in fig.axes:
            # after all ylims are established
            ax.set_ylim(
                bottom=max(
                    0,
                    expected_period - np.max(np.abs(y_deviations_from_expected_period)),
                ),
            )
            ticks_with_period = sorted(set(ax.get_yticks()) | {expected_period})
            ax.set_yticks(ticks_with_period)
            if idx == 0:
                ax.set_yticklabels([f"{_:.3f}" for _ in ticks_with_period])
        fig.set_tight_layout(True)

        return fig, fig.axes

    def close(self) -> None:
        """
        Closes the dataset.
        """
        self.dfile.close()

    def __enter__(self) -> Self:
        """
        So we can use context manager (with...as) like any other open file.

        Examples
        --------
        >>> with Dataset('my_data.h5') as d: # doctest: +SKIP
        ...     d.stats()

        """
        return self

    def __exit__(self, type, value, traceback) -> None:
        """
        Exit statement for context manager.
        """
        self.close()


def reshape_into_blocks(
    indices: Sequence[float] | npt.NDArray[np.floating],
    min_gap: int | float | None = None,
) -> tuple[npt.NDArray[np.floating], ...]:
    """
    Find the large gaps in indices and split at each gap.

    For example, if two blocks of stimuli were recorded in a single sync
    file, there will be one larger-than normal gap in frame timestamps.

    - default min gap threshold: median + 3 * std (won't work well for short seqs)

    >>> reshape_into_blocks([0, 1, 2, 103, 104, 105], min_gap=100)
    (array([0, 1, 2]), array([103, 104, 105]))

    >>> reshape_into_blocks([0, 1, 2, 3])
    (array([0, 1, 2, 3]),)
    """
    indices = np.array(indices)
    intervals = np.diff(indices)
    long_interval_threshold = (
        min_gap
        if min_gap is not None
        else (np.median(intervals) + 3 * np.std(intervals))
    )

    gaps_between_blocks = []
    for interval_index, interval in zip(
        intervals.argsort()[::-1], sorted(intervals)[::-1]
    ):
        if interval > long_interval_threshold:
            # large interval found
            gaps_between_blocks.append(interval_index + 1)
        else:
            break

    if not gaps_between_blocks:
        # a single block of timestamps
        return (np.array(indices),)

    # create blocks as intervals [start:end]
    gaps_between_blocks.sort()
    blocks = []
    start = 0
    for end in gaps_between_blocks:
        blocks.append(indices[start:end])
        start = end
    # add end of last block
    blocks.append(indices[start:])

    # filter out blocks with a single sample (not a block)
    blocks = [block for block in blocks if len(block) > 1]

    # filter out blocks with long avg timstamp interval (a few, widely-spaced timestamps)
    blocks = [
        block for block in blocks if np.median(np.diff(block)) < long_interval_threshold
    ]

    return tuple(blocks)


def adjust_diode_flip_intervals(
    diode_flips: npt.NDArray[np.floating],
) -> npt.NDArray[np.floating]:
    """
    diode flip intervals have a bimodal distribution due to asymmetry of
    photodiode thresholding: adjust every other interval to get a closer
    estimate of actual transition time for each diode flip
    """
    diode_flips = np.array(diode_flips)  # make a copy
    original_intervals = np.diff(diode_flips)
    outliers = np.logical_or(
        np.percentile(original_intervals, 5) > original_intervals,
        original_intervals > np.percentile(original_intervals, 95),
    )
    intervals = original_intervals[~outliers]

    # the two distributions are symmetric about the mean
    deviation = np.mean(np.abs(intervals - np.mean(intervals)))
    if np.mean(original_intervals[0::2]) > np.mean(original_intervals[1::2]):
        sign = 1
    else:
        sign = -1

    for idx in range(1, len(diode_flips) - 1, 2):
        # alternate on every short/long interval and expand/contract
        # interval
        diode_flips[idx] -= sign * 0.5 * deviation
        diode_flips[idx + 1] += sign * 0.5 * deviation
    return diode_flips


def add_missing_diode_flips_for_truncated_sync(
    diode_flips: npt.NDArray,
    vsyncs: npt.NDArray,
) -> npt.NDArray[np.float64]:
    while len(diode_flips) < len(vsyncs):
        logger.info(
            "Creating extra diode flip at end of stim block to account for truncated sync recording"
        )
        common_len = min([len(vsyncs) - 1, len(diode_flips)])
        avg_vsync_to_flip_interval = np.median(
            adjust_diode_flip_intervals(diode_flips[:common_len])
            - vsyncs[1 : common_len + 1]
        )
        diode_flips = np.array([*diode_flips, vsyncs[-1] + avg_vsync_to_flip_interval])
    return diode_flips


def add_missing_diode_flip_at_stim_onset(
    diode_flips: npt.NDArray,
    vsyncs: npt.NDArray,
) -> npt.NDArray[np.float64]:
    """
    Transition from white screen to white sync-square on first frame results
    in a missing diode flip (same for black-to-black).
    Create a flip after the first vsync with an interval
    based on the statistics of other intervals in the
    recording.
    Shouldn't have any side-effects, and the first frame likely contains nothing important.
    """
    logger.info(
        "Creating extra diode flip at start of stim block to account for white-to-white or black-to-black transition"
    )
    common_len = min([len(vsyncs) - 1, len(diode_flips)])
    avg_vsync_to_flip_interval = np.median(
        adjust_diode_flip_intervals(diode_flips[:common_len])
        - vsyncs[1 : common_len + 1]
    )
    # TODO could be more accurate by adding short/long interval as appropriate
    # - diode_flips are pre-adjustment, but we're currently adding the expected
    #   post-adjustment interval
    return np.array([vsyncs[0] + avg_vsync_to_flip_interval, *diode_flips])


def discard_erroneous_diode_flips_at_stim_offset(
    diode_flips: npt.NDArray, vsyncs: npt.NDArray
) -> npt.NDArray[np.float64]:
    """
    - after a stimulus the screen usually turns grey, causing
    the diode to flip at least one more time after all stim
    vsyncs are finished.
    - the grey screen itself may be on the threshold of diode activation, causing
    multiple additional flips

    remove the last flip if doesn't look time-locked with a
    vsync (ie. vsync-to-flip interval is an outlier)
    """
    while len(diode_flips[diode_flips > vsyncs[-1]]) > 1 and len(diode_flips) > len(
        vsyncs
    ):
        diode_flips = diode_flips[:-1]
    return diode_flips


if __name__ == "__main__":
    import doctest

    import dotenv

    dotenv.load_dotenv(dotenv.find_dotenv(usecwd=True))
    doctest.testmod(
        optionflags=(doctest.IGNORE_EXCEPTION_DETAIL | doctest.NORMALIZE_WHITESPACE)
    )
