import sys
from functools import lru_cache
import traceback
import warnings
from matplotlib import pyplot as plt
import os
import datetime
import numpy as np
import h5py
import tables
import pandas as pd
import json
from tqdm.auto import tqdm
from scipy import interpolate
if sys.version_info >= (3,9):
    from importlib.resources import files, as_file
else:
    from importlib_resources import files, as_file

from .. import precondition
from .. import preproc
from ..preproc.base import get_data_segment, get_data_segments, get_trial_segments, get_trial_segments_and_times, interp_timestamps2timeseries, sample_timestamped_data, trial_align_data
from ..preproc.bmi3d import get_target_events, get_ref_dis_frequencies
from ..whitematter import ChunkedStream, Dataset
from ..utils import derivative, get_pulse_edge_times, compute_pulse_duty_cycles, convert_digital_to_channels, detect_edges
from .base import load_preproc_exp_data, load_preproc_eye_data, load_preproc_lfp_data, yaml_read, get_preprocessed_filename, load_hdf_data, load_hdf_ts_trial, load_hdf_ts_segment, load_hdf_group

############
# Raw data #
############
def get_ecube_data_sources(data_dir):
    '''
    Lists the available data sources in a given data directory

    Args: 
        data_dir (str): eCube data directory

    Returns:
        str array: available sources (AnalogPanel, Headstages, etc.)
    '''
    dat = Dataset(data_dir)
    return dat.listsources()
    
def load_ecube_metadata(data_dir, data_source):
    '''
    Sums the number of channels and samples across all files in the data_dir

    Args: 
        data_dir (str): eCube data directory
        source (str): selects the source (AnalogPanel, Headstages, etc.)

    Returns:
        dict: Dictionary of metadata with fields:
            | **samplerate (float):** sampling rate of data for this source
            | **data_source (str):** copied from the function argument
            | **n_channels (int):** number of channels
            | **n_samples (int):** number of samples for one channel
    '''

    # For now just load the metadata provieded by pyECubeSig
    # TODO: Really need the channel names and voltage per bit from the xml file
    # For now I am hard coding these. Please change!
    HEADSTAGE_VOLTSPERBIT = 1.907348633e-7
    ANALOG_VOLTSPERBIT = 3.0517578125e-4
    if data_source == 'Headstages':
        voltsperbit = HEADSTAGE_VOLTSPERBIT
    elif data_source == 'AnalogPanel':
        voltsperbit = ANALOG_VOLTSPERBIT
    else:
        voltsperbit = None

    n_channels = 0
    n_samples = 0
    dat = Dataset(data_dir)
    recordings = dat.listrecordings()
    for r in recordings: # r: (data_source, n_channels, n_samples)
        if data_source in r[0]:
            n_samples += r[2]  
            n_channels = r[1]
    if n_channels == 0:
        raise Exception('No data found for data source: {}'.format(data_source))
    samplerate = dat.samplerate
    metadata = dict(
        samplerate = samplerate,
        data_source = data_source,
        n_channels = n_channels,
        n_samples = n_samples,
        voltsperbit = voltsperbit,
    )
    return metadata

def load_ecube_data(data_dir, data_source, channels=None):
    '''
    Loads data from eCube for a given directory and datasource

    Requires load_ecube_metadata(), process_channels()

    Args:
        data_dir (str): folder containing the data you want to load
        data_source (str): type of data ("Headstages", "AnalogPanel", "DigitalPanel")
        channels (int array or None): list of channel numbers (0-indexed) to load. If None, will load all channels by default

    Returns:
        (nt, nch): all the data for the given source
    '''

    # Read metadata, check inputs
    metadata = load_ecube_metadata(data_dir, data_source)
    if channels is None:
        channels = range(metadata['n_channels'])
    elif len(channels) > metadata['n_channels']:
        raise ValueError("Supplied channel numbers are invalid")

    # Datatype is currently fixed for each data_source
    if data_source == 'DigitalPanel':
        dtype = np.uint64
    else:
        dtype = np.int16

    # Fetch all the data for all the channels
    n_samples = metadata['n_samples']
    timeseries_data = np.zeros((n_samples, len(channels)), dtype=dtype)
    n_read = 0
    for chunk in _process_channels(data_dir, data_source, channels, metadata['n_samples'], dtype=dtype):
        chunk_len = min(chunk.shape[0], n_samples-n_read)
        timeseries_data[n_read:n_read+chunk_len,:] = chunk[:chunk_len,:] # Deal with potentially corrupted data
        n_read += chunk_len
    return timeseries_data

def load_ecube_data_chunked(data_dir, data_source, channels=None, chunksize=728):
    '''
    Loads a data file one "chunk" at a time. Useful for replaying files as if they were online data.

    Args:
        data_dir (str): folder containing the data you want to load
        data_source (str): type of data ("Headstages", "AnalogPanel", "DigitalPanel")
        channels (int array or None): list of channel numbers (0-indexed) to load. If None, will load all channels by default
        chunksize (int): how many samples to include in each chunk

    Yields:
        (chunksize, nch): one chunk of data for the given source
    '''
    # Read metadata, check inputs
    metadata = load_ecube_metadata(data_dir, data_source)
    if channels is None:
        channels = range(metadata['n_channels'])
    elif len(channels) > metadata['n_channels']:
        raise ValueError("Supplied channel numbers are invalid")

    # Datatype is currently fixed for each data_source
    if data_source == 'DigitalPanel':
        dtype = np.uint64
    else:
        dtype = np.int16
    
    # Fetch all the channels but just return the generator
    n_samples = metadata['n_samples']
    n_channels = metadata['n_channels']
    kwargs = dict(maxchunksize=chunksize*n_channels*np.dtype(dtype).itemsize)
    return _process_channels(data_dir, data_source, channels, n_samples, **kwargs)

def proc_ecube_data(data_path, data_source, result_filepath, result_name='broadband_data', max_memory_gb=1.):
    '''
    Loads and saves eCube data into an HDF file

    Requires load_ecube_metadata()

    Args:
        data_path (str): path to folder containing the ecube data you want to load
        data_source (str): type of data ("Headstages", "AnalogPanel", "DigitalPanel")
        result_filepath (str): path to hdf file to be written (or appended)
        max_memory_gb (float, optional): max memory used to load binary data at one time

    Returns:
        tuple: tuple containing:
            | **dset (h5py.Dataset):** the new hdf dataset
            | **metadata (dict):** the ecube metadata
    '''

    # Load the metadata to figure out the datatypes
    metadata = load_ecube_metadata(data_path, data_source)
    if data_source == 'DigitalPanel':
        dtype = np.uint64
    else:
        dtype = np.int16
    chunksize = int(max_memory_gb * 1e9 / np.dtype(dtype).itemsize / metadata['n_channels'])

    # Create an hdf dataset
    hdf = h5py.File(result_filepath, 'a')
    dset = hdf.create_dataset(result_name, (metadata['n_samples'], metadata['n_channels']), dtype=dtype)

    # Write broadband data directly into the hdf file
    n_samples = 0
    for broadband_chunk in load_ecube_data_chunked(data_path, 'Headstages', chunksize=chunksize):
        chunk_len = broadband_chunk.shape[0]
        dset[n_samples:n_samples+chunk_len,:] = broadband_chunk
        n_samples += chunk_len

    return dset, metadata

def filter_lfp_from_broadband(broadband_filepath, result_filepath, mean_subtract=True, dtype='int16', max_memory_gb=1., **filter_kwargs):
    '''
    Filters local field potential (LFP) data from a given broadband signal file into an hdf file.

    Args:
        broadband_filepath (str): Path to the input broadband signal file.
        result_filepath (str): Path to save the filtered LFP data.
        mean_subtract (bool, optional): Whether to subtract the mean from the filtered LFP signal.
                                        Default is True.
        dtype (str, optional): Data type for the filtered LFP signal. Default is 'int16'.
        max_memory_gb (float, optional): Maximum memory (in gigabytes) to use for filtering. Default is 1.0 GB.
        **filter_kwargs: Additional keyword arguments to customize the filtering process.
                        These arguments will be passed to the filtering function.

    Raises:
        IOError: If the input broadband file is not found.
        MemoryError: If the specified max_memory_gb is insufficient for the filtering process.

    Note:
        This function is used in the :func:`~aopy.preproc.warppers.proc_lfp` wrapper.
    '''
    lfp_samplerate = filter_kwargs.pop('lfp_samplerate', 1000)

    metadata = load_hdf_group('', broadband_filepath, 'broadband_metadata')
    samplerate = metadata['samplerate']
    n_channels = int(metadata['n_channels'])
    n_samples = int(metadata['n_samples'])
    downsample_factor = int(samplerate/lfp_samplerate)
    lfp_samples = int(np.ceil(n_samples/downsample_factor))
    if 'low_cut' not in filter_kwargs:
        filter_kwargs['low_cut'] = 500
    if 'buttord' not in filter_kwargs:
        filter_kwargs['buttord'] = 4

    # Create an hdf dataset
    lfp_hdf = h5py.File(result_filepath, 'a') # should append existing or write new?
    dset = lfp_hdf.create_dataset('lfp_data', (lfp_samples, n_channels), dtype=dtype)

    # Figure out how much data we can load at once
    max_samples = int(max_memory_gb * 1e9 / np.dtype(dtype).itemsize)
    channel_chunksize = max(min(n_channels, max_samples // n_samples), 1)
    time_chunksize = min(n_samples, max_samples // channel_chunksize)
    print(f"{channel_chunksize} channels and {time_chunksize} samples in each chunk")
    
    # Load the broadband dataset
    bb_hdf = h5py.File(broadband_filepath, 'r')
    if 'broadband_data' not in bb_hdf:
        raise ValueError(f'broadband_data not found in file {broadband_filepath}')
    bb_data = bb_hdf['broadband_data']
    
    # Filter broadband data into LFP directly into the hdf file
    n_bb_samples = 0
    n_lfp_samples = 0
    while n_bb_samples < n_samples:
        n_ch = 0
        while n_ch < n_channels:
            broadband_chunk = bb_data[n_bb_samples:n_bb_samples+time_chunksize, n_ch:n_ch+channel_chunksize]
            lfp_chunk, _ = precondition.filter_lfp(broadband_chunk, samplerate, **filter_kwargs)
            chunk_len = lfp_chunk.shape[0]
            dset[n_lfp_samples:n_lfp_samples+chunk_len,n_ch:n_ch+channel_chunksize] = lfp_chunk
            n_ch += channel_chunksize
        n_bb_samples += time_chunksize
        n_lfp_samples += chunk_len
        
    if mean_subtract:
        dset -= np.mean(dset, axis=0, dtype=dtype) # hopefully this isn't constrained by memory
    lfp_hdf.close()
    bb_hdf.close()

    # Append the lfp metadata to the file
    lfp_metadata = metadata
    lfp_metadata['lfp_samplerate'] = lfp_samplerate # for backwards compatibility
    lfp_metadata['samplerate'] = lfp_samplerate
    lfp_metadata['n_samples'] = lfp_samples
    lfp_metadata.update(filter_kwargs)
    
    return dset, lfp_metadata

def filter_lfp_from_ecube(ecube_filepath, result_filepath, mean_subtract=True, dtype='int16', max_memory_gb=1., **filter_kwargs):
    '''
    Filters local field potential (LFP) data from an eCube recording file.

    Args:
        ecube_filepath (str): Path to the input eCube recording file.
        result_filepath (str): Path to save the filtered LFP data.
        mean_subtract (bool, optional): Whether to subtract the mean from the filtered LFP signal.
                                        Default is True.
        dtype (str, optional): Data type for the filtered LFP signal. Default is 'int16'.
        max_memory_gb (float, optional): Maximum memory (in gigabytes) to use for filtering. Default is 1.0 GB.
        **filter_kwargs: Additional keyword arguments to customize the filtering process.
                        These arguments will be passed to the filtering function.

    Raises:
        IOError: If the input eCube recording file is not found.
        MemoryError: If the specified max_memory_gb is insufficient for the filtering process.

    Note:
        This function is used in the :func:`~aopy.preproc.warppers.proc_lfp` wrapper.
    '''
    lfp_samplerate = filter_kwargs.pop('lfp_samplerate', 1000)

    metadata = load_ecube_metadata(ecube_filepath, 'Headstages')
    samplerate = metadata['samplerate']
    downsample_factor = int(samplerate/lfp_samplerate)
    n_channels = int(metadata['n_channels'])
    n_samples = int(metadata['n_samples'])
    if 'low_cut' not in filter_kwargs:
        filter_kwargs['low_cut'] = 500
    if 'buttord' not in filter_kwargs:
        filter_kwargs['buttord'] = 4

    # Figure out how much data we can load at once
    max_samples = int(max_memory_gb * 1e9 / np.dtype(dtype).itemsize)
    chunksize = int(max_samples / metadata['n_channels'])
    n_whole_chunks = int(n_samples / chunksize)
    n_remaining = n_samples - (chunksize * n_whole_chunks)
    lfp_samples = np.ceil(chunksize/downsample_factor) * n_whole_chunks + np.ceil(n_remaining/downsample_factor)
    print("lfp chunks should be ", np.ceil(chunksize/downsample_factor))
    print("last chunk should be ", np.ceil(n_remaining/downsample_factor))
    print("total samples: ", lfp_samples)

    # Create an hdf dataset
    hdf = h5py.File(result_filepath, 'a') # should append existing or write new?
    dset = hdf.create_dataset('lfp_data', (lfp_samples, n_channels), dtype=dtype)

    # Filter broadband data into LFP directly into the hdf file
    n_lfp_samples = 0
    for broadband_chunk in load_ecube_data_chunked(ecube_filepath, 'Headstages', chunksize=chunksize):
        lfp_chunk, _ = precondition.filter_lfp(broadband_chunk, samplerate, **filter_kwargs)
        chunk_len = lfp_chunk.shape[0]
        dset[n_lfp_samples:n_lfp_samples+chunk_len,:] = lfp_chunk
        n_lfp_samples += chunk_len
        
    if mean_subtract:
        dset -= np.mean(dset, axis=0, dtype=dtype) # hopefully this isn't constrained by memory
    hdf.close()

    # Append the lfp metadata to the file
    lfp_metadata = metadata
    lfp_metadata['lfp_samplerate'] = lfp_samplerate # for backwards compatibility
    lfp_metadata['samplerate'] = lfp_samplerate
    lfp_metadata['n_samples'] = lfp_samples
    lfp_metadata.update(filter_kwargs)
    
    return dset, lfp_metadata

def _process_channels(data_dir, data_source, channels, n_samples, dtype=None, debug=False, **dataset_kwargs):
    '''
    Reads data from an ecube data source by channel until the number of samples requested 
    has been loaded. If a processing function is supplied, it will be applied to 
    each batch of data. If not, the data will be appended 

    Args:
        data_dir (str): folder containing the data you want to load
        data_source (str): type of data ("Headstages", "AnalogPanel", "DigitalPanel")
        channels (int array): list of channels to process
        n_samples (int): number of samples to read. Must be geq than a single chunk
        dtype (np.dtype): format for data_out if none supplied
        data_out (nt, nch): array of data to be written to. If None, it will be created
        debug (bool): whether the data is read in debug mode
        dataset_kwargs (kwargs): list of key value pairs to pass to the ecube dataset
        
    Yields:
        (nt, nch): Chunks of the requested samples for requested channels
    '''

    dat = Dataset(data_dir, **dataset_kwargs)
    dat.selectsource(data_source)
    chunk = dat.emitchunk(startat=0, debug=debug)
    datastream = ChunkedStream(chunkemitter=chunk)

    idx_samples = 0 # keeps track of the number of samples already read/written
    while idx_samples < n_samples:
        try:
            data_chunk = next(datastream)
            data_len = np.shape(data_chunk)[1]
            if len(channels) == 1:
                yield data_chunk[channels,:].T 
            else:
                yield np.squeeze(data_chunk[channels,:]).T # this might be where you filter data
            idx_samples += data_len
        except StopIteration:
            break

def load_ecube_digital(path, data_dir):
    '''
    Just a wrapper around load_ecube_data() and load_ecube_metadata()

    Args:
        path (str): base directory where ecube data is stored
        data_dir (str): folder you want to load

    Returns:
        tuple: Tuple containing:
            | **data (nt):** digital data, arranged as 64-bit numbers representing the 64 channels
            | **metadata (dict):** metadata (see load_ecube_metadata() for details)
    '''
    data = load_ecube_data(os.path.join(path, data_dir), 'DigitalPanel')
    metadata = load_ecube_metadata(os.path.join(path, data_dir), 'DigitalPanel')
    return data, metadata

def load_ecube_analog(path, data_dir, channels=None):
    '''
    Just a wrapper around load_ecube_data() and load_ecube_metadata()

    Args:
        path (str): base directory where ecube data is stored
        data_dir (str): folder you want to load
        channels (int array, optional): which channels to load

    Returns:
        tuple: Tuple containing:
            | **data (nt, nch):** analog data for the requested channels
            | **metadata (dict):** metadata (see load_ecube_metadata() for details)
    '''
    data = load_ecube_data(os.path.join(path, data_dir), 'AnalogPanel', channels)
    metadata = load_ecube_metadata(os.path.join(path, data_dir), 'AnalogPanel')
    return data, metadata

def load_ecube_headstages(path, data_dir, channels=None):
    '''
    Just a wrapper around load_ecube_data() and load_ecube_metadata()

    Args:
        path (str): base directory where ecube data is stored
        data_dir (str): folder you want to load
        channels (int array, optional): which channels to load

    Returns:
        tuple: Tuple containing:
            | **data (nt, nch):** analog data for the requested channels
            | **metadata (dict):** metadata (see load_ecube_metadata() for details)
    '''
    data = load_ecube_data(os.path.join(path, data_dir), 'Headstages', channels)
    metadata = load_ecube_metadata(os.path.join(path, data_dir), 'Headstages')
    return data, metadata


def get_e3v_video_frame_data( digital_data, sync_channel_idx, trigger_channel_idx, samplerate ):
    
    """get_e3v_video_frame_data

    Compute pulse times and duty cycles from e3vision video data frames collected on an ecube digital panel.

    Args:
        digital_data (nt, nch): array of data read from ecube digital panel
        sync_channel_idx (int): sync channel to read from digital_data. Indicates each video frame.
        trigger_channel_idx (int): trigger channel to read from digital_data. Indicates start/end video triggers.
        sample_rate (numeric): data sampling rate (Hz)

    Returns:
        pulse_times (np.array): array of floats indicating pulse start times
        duty_cycle (np.array): array of floats indicating pulse duty cycle (quotient of pulse width and pulse period)
    """

    trig_pulse_edges = get_pulse_edge_times(digital_data[:,trigger_channel_idx],samplerate)
    # watchtower triggers (start, end) are a triplet of pulses within a ~33ms window.
    trig_pulse_times = trig_pulse_edges[:,0]
    start_trig_time = trig_pulse_times[0]
    end_trig_time = trig_pulse_times[3]
    sync_pulse_edges = get_pulse_edge_times(digital_data[:,sync_channel_idx],samplerate)
    sync_pulse_times = sync_pulse_edges[:,0]
    start_pulse_idx = np.where(np.abs(sync_pulse_times-start_trig_time) == np.abs(sync_pulse_times-start_trig_time).min())[0][0]
    end_pulse_idx = np.where(np.abs(sync_pulse_times-end_trig_time) == np.abs(sync_pulse_times-end_trig_time).min())[0][0] + 1
    sync_pulse_times = sync_pulse_times[start_pulse_idx:end_pulse_idx]
    sync_duty_cycles = compute_pulse_duty_cycles(sync_pulse_edges[start_pulse_idx:end_pulse_idx,:])

    return sync_pulse_times, sync_duty_cycles

def load_bmi3d_hdf_table(data_dir, filename, table_name):
    '''
    Loads data and metadata from a table in an hdf file generated by BMI3D

    Args:
        data_dir (str): path to the data
        filename (str): name of the file to load from
        table_name (str): name of the table you want to load

    Returns:
        tuple: Tuple containing:
            | **data (ndarray):** data from bmi3d
            | **metadata (dict):** attributes associated with the table
    '''
    filepath = os.path.join(data_dir, filename)
    with tables.open_file(filepath, 'r') as f:
        if table_name not in f.root:
            raise ValueError(f"{table_name} not found in {filename}")
        table = getattr(f.root, table_name)
        param_keys = table.attrs._f_list("user")
        metadata = {k : getattr(table.attrs, k) for k in param_keys}
        return table.read(), metadata

def load_bmi3d_root_metadata(data_dir, filename):
    '''
    Root metadata not accessible using pytables, instead use h5py

    Args:
        data_dir (str): path to the data
        filename (str): name of the file to load from

    Returns:
        dict: key-value attributes
    '''
    with h5py.File(os.path.join(data_dir, filename), 'r') as f:
        return dict(f['/'].attrs.items())

def get_ecube_digital_input_times(path, data_dir, ch):
    '''
    Computes the times when digital input turns on or off in ecube
    For synchronizing openephys with ecube, use ch=-1.
    
    Args:
    path (str): base directory where ecube data is stored
    data_dir (str): folder you want to load
    ch (str): digital channel
        
    Returns:
        tuple: Tuple containing:
            | **on_times (n_times):** times at which sync line turned on
            | **off_times (n_times):** times at which sync line turned off    
    '''
    
    # Load ecube digital data
    digital_data, metadata = load_ecube_digital(path, data_dir)
    FS = metadata['samplerate']
    
    # Convert 64bit information into each channel data
    digital_data_ch = convert_digital_to_channels(digital_data)
        
    # Get on_times and off_times in digital data
    on_times,_ = detect_edges(digital_data_ch[:,ch], FS, rising=True, falling=False)
    off_times,_ = detect_edges(digital_data_ch[:,ch], FS, rising=False, falling=True)

    return on_times, off_times

#####################
# Preprocessed data #
#####################
def get_interp_kinematics(exp_data, exp_metadata, datatype='cursor', samplerate=1000, **kwargs):
    '''
    Gets interpolated kinematic data from preprocessed experiment data to the desired 
    sampling rate. Cursor kinematics are returned in screen coordinates, while other 
    kinematics are returned in their original coordinate system (e.g. hand kinematics 
    in optitrack coordinates).

    Args:
        exp_data (dict): A dictionary containing the experiment data.
        exp_metadata (dict): A dictionary containing the experiment metadata.
        datatype (str, optional): The type of kinematic data to interpolate. 
            For 'hand' kinematics, interp the 'clean_hand_position' experiment data
            For 'cursor' kinematics, interp the x and z position of the 'cursor' task data
            For other kinematics, try to interp exp_data['task'][datatype]
        samplerate (float, optional): The desired output sampling rate in Hz. 
            Defaults to 1000.
        **kwargs: Additional keyword arguments to pass to sample_timestamped_data()

    Returns:
        data_time (ns, ...): Kinematic data interpolated and filtered 
            to the desired sampling rate.

    Examples:
        
        Cursor kinematics

        .. code-block:: python
        
            exp_data, exp_metadata = load_preproc_exp_data(preproc_dir, 'test',  3498, '2021-12-13')
            cursor_interp = get_interp_kinematics(exp_data, exp_metadata, datatype='cursor', samplerate=100)

            plt.figure()
            visualization.plot_trajectories([cursor_interp], [-10, 10, -10, 10])
        
        .. image:: _images/get_interp_cursor_centerout.png

        Hand kinematics
       
        .. code-block:: python

            hand_interp = get_interp_kinematics(exp_data, exp_metadata, datatype='hand', samplerate=100)
            ax = plt.axes(projection='3d')
            visualization.plot_trajectories([hand_interp], [-10, 10, -10, 10, -10, 10])

        .. image:: _images/get_interp_hand_centerout.png

        Target positions

        .. code-block:: python
            
            targets_interp = get_interp_kinematics(exp_data, exp_metadata, datatype='targets', samplerate=100)
            time = np.arange(len(targets_interp))/100
            plt.plot(time, targets_interp[:,:,0]) # plot just the x coordinate
            plt.xlim(10, 20)
            plt.xlabel('time (s)')
            plt.ylabel('x position (cm)')

        .. image:: _images/get_interp_targets_centerout.png

        Cursor and target (reference) kinematics

        .. code-block:: python
            
            exp_data, exp_metadata = load_preproc_exp_data(data_dir, 'test', 8461, '2023-02-25')
            cursor_interp = get_interp_kinematics(exp_data, exp_metadata, datatype='cursor', samplerate=exp_metadata['fps'])
            ref_interp = get_interp_kinematics(exp_data, exp_metadata, datatype='reference', samplerate=exp_metadata['fps'])
            time = np.arange(exp_metadata['fps']*120)/exp_metadata['fps']
            plt.plot(time, cursor_interp[:int(exp_metadata['fps']*120),1], color='blueviolet', label='cursor') # plot just the y coordinate
            plt.plot(time, ref_interp[:int(exp_metadata['fps']*120),1], color='darkorange', label='ref')
            plt.xlabel('time (s)')
            plt.ylabel('y position (cm)'); plt.ylim(-10,10)
            plt.legend()

        .. image:: _images/get_interp_cursor_tracking.png

        User, reference, and disturbance kinematics

        .. code-block:: python
            
            user_interp = get_interp_kinematics(exp_data, exp_metadata, datatype='user', samplerate=exp_metadata['fps'])
            ref_interp = get_interp_kinematics(exp_data, exp_metadata, datatype='reference', samplerate=exp_metadata['fps'])
            dis_interp = get_interp_kinematics(exp_data, exp_metadata, datatype='disturbance', samplerate=exp_metadata['fps'])
            time = np.arange(exp_metadata['fps']*120)/exp_metadata['fps']
            plt.plot(time, user_interp[:int(exp_metadata['fps']*120),1], color='darkturquoise', label='user')
            plt.plot(time, ref_interp[:int(exp_metadata['fps']*120),1], color='darkorange', label='ref')
            plt.plot(time, dis_interp[:int(exp_metadata['fps']*120),1], color='tab:red', linestyle='--', label='dis')
            plt.xlabel('time (s)')
            plt.ylabel('y position (cm)'); plt.ylim(-10,10)
            plt.legend()

        .. image:: _images/get_interp_user_tracking.png

    Changes:
        2023-10-20: Added support for 'targets' datatype
        2024-01-29: Removed kinematic filtering below 15 Hz. See :func:`~aopy.precondition.filter_kinematics`.
    '''

    # Fetch the available timestamps
    try:
        clock = exp_data['clock']['timestamp_sync']
    except:
        clock = exp_data['clock']['timestamp_bmi3d']

    # Fetch the relevant BMI3D data
    if datatype == 'hand':
        data_cycles = exp_data['clean_hand_position'] # 3d hand position (optitrack coords: x,y,z) on each bmi3d cycle
    elif datatype == 'cursor':
        data_cycles = exp_data['task']['cursor'][:,[0,2]] # 2d cursor position (bmi3d coords: x,z) on each bmi3d cycle
    elif datatype == 'user':
        dis_on = int(json.loads(exp_metadata['sequence_params'])['disturbance']) # whether disturbance was turned on (0 or 1)
        data_cycles = exp_data['task']['cursor'][:,[0,2]] - exp_data['task']['current_disturbance'][:,[0,2]]*dis_on # 2d cursor position before disturbance added (bmi3d coords: x,z)
    elif datatype == 'reference':
        data_cycles =  exp_data['task']['current_target'][:,[0,2]] # 2d target position (bmi3d coords: x,z)
    elif datatype == 'disturbance':
        dis_on = int(json.loads(exp_metadata['sequence_params'])['disturbance']) # whether disturbance was turned on (0 or 1)
        data_cycles = exp_data['task']['current_disturbance'][:,[0,2]]*dis_on # 2d disturbance value (bmi3d coords: x,z)
    elif datatype == 'targets':
        data_cycles = get_target_events(exp_data, exp_metadata)
        clock = exp_data['events']['timestamp']
        kwargs['remove_nan'] = False # In this case we need to keep NaN values.
    elif datatype in exp_data['task'].dtype.names:
        data_cycles = exp_data['task'][datatype]
    else:
        raise ValueError(f"Unknown datatype {datatype}")
    
    # Interpolate
    data_time = sample_timestamped_data(data_cycles, clock, samplerate, append_time=10, **kwargs)

    return data_time

def get_velocity_segments(*args, norm=True, **kwargs):
    '''
    Estimates velocity from cursor position, then finds the trial segments for velocity using 
    :func:`~aopy.postproc.get_kinematic_segments()`.
    
    Args:
        *args: arguments for :func:`~aopy.postproc.get_kinematic_segments`
        norm (bool): if the output segments should be normalized. Set to false to output component velocities.
        **kwargs: parameters for :func:`~aopy.postproc.get_kinematic_segments`
        
    Returns:
        tuple: tuple containing:
            | **velocities (ntrial):** array of velocity estimates for each trial
            | **trial_segments (ntrial):** array of numeric code segments for each trial
    '''
    def preproc(pos, fs):
        time = np.arange(pos.shape[0])/fs
        return derivative(time, pos, norm=norm), fs
    return get_kinematic_segments(*args, **kwargs, preproc=preproc)

@lru_cache(maxsize=1)
def get_kinematics(preproc_dir, subject, te_id, date, samplerate, preproc=None, datatype='cursor', **kwargs):
    '''
    Return all kinds of kinematics from preprocessed data. Caches the data for faster loading.

    Note: 
        You can avoid the phase shift in downsampled data when using get_interp_kinematics by setting 
        upsamplerate=samplerate, so that it doesn't do any up/down sampling, only interpolation at the 
        same samplerate.

    Args:
        preproc_dir (str): base directory where the files live
        subject (str): Subject name
        te_id (int): Block number of Task entry object 
        date (str): Date of recording
        samplerate (float, optional): optionally choose the samplerate of the data in Hz. Default 1000.
        preproc (fn, optional): function mapping (position, fs) data to (kinematics, fs_new). For example,
            a smoothing function or an estimate of velocity from position
        datatype (str, optional): type of kinematics to load. Defaults to 'cursor'.   
        kwargs: additional keyword arguments to pass to get_interp_kinematics 

    Raises:
        ValueError: if the datatype is invalid

    Returns:
        tuple: tuple containing:
            | **kinematics (nt, nch):** kinematics from the given experiment after preprocessing
            | **samplerate (float):** the sampling rate of the kinematics after preprocessing
    '''
    data, metadata = load_preproc_exp_data(preproc_dir, subject, te_id, date)

    if 'eye' in datatype:
        eye_data, eye_metadata = load_preproc_eye_data(preproc_dir, subject, te_id, date)
        if datatype == 'eye_raw':
            eye_data = eye_data['raw_data']
        elif datatype == 'eye_closed_mask':
            eye_data = eye_data['eye_closed_mask']
        elif 'calibrated_data' in eye_data.keys():
            eye_data = eye_data['calibrated_data']
        else:
            raise ValueError(f"No calibrated eye data for {te_id}")
        
        time = np.arange(len(eye_data))/eye_metadata['samplerate']
        raw_kinematics, _ = interp_timestamps2timeseries(time, eye_data, samplerate)
    else:
        raw_kinematics = get_interp_kinematics(
            data, metadata, datatype, samplerate=samplerate, **kwargs
        )

    time = np.arange(len(raw_kinematics))/samplerate
    if preproc is not None:
        kinematics, samplerate = preproc(raw_kinematics, samplerate)
    else:
        kinematics = raw_kinematics

    return kinematics, samplerate

def get_kinematic_segment(preproc_dir, subject, te_id, date, start_time, end_time, samplerate, 
                          preproc=None, datatype='cursor', **kwargs):
    '''
    Return one segment of kinematics

    Args:
        preproc_dir (str): base directory where the files live
        subject (str): Subject name
        te_id (int): Block number of Task entry object 
        date (str): Date of recording
        start_time (float): time in the recording at which the desired segment starts
        end_time (float): time in the recording at which the desired segment ends
        samplerate (float, optional): optionally choose the samplerate of the data in Hz. Default 1000.
        preproc (fn, optional): function mapping (position, fs) data to (kinematics, fs_new). For example,
            a smoothing function or an estimate of velocity from position
        datatype (str, optional): type of kinematics to load. Defaults to 'cursor'.    
        kwargs: additional keyword arguments to pass to get_kinematics

    Returns:
        tuple: tuple containing:
            | **segment (nt, nch):** single kinematic segment from the given experiment after preprocessing
            | **samplerate (float):** the sampling rate of the kinematics after preprocessing
    '''
    kinematics, samplerate = get_kinematics(preproc_dir, subject, te_id, date, samplerate, preproc, datatype, **kwargs)
    assert kinematics is not None

    return get_data_segment(kinematics, start_time, end_time, samplerate), samplerate

def get_kinematic_segments(preproc_dir, subject, te_id, date, trial_start_codes, trial_end_codes, 
                           trial_filter=lambda x:True, preproc=None, datatype='cursor',
                           samplerate=1000, **kwargs):
    '''
    Loads x,y,z cursor, hand, or eye trajectories for each "trial" from a preprocessed HDF file. Trials can
    be specified by numeric start and end codes. Trials can also be filtered so that only successful
    trials are included, for example. The filter is applied to numeric code segments for each trial. 
    Finally, the cursor data can be preprocessed by a supplied function to, for example, convert 
    position to velocity estimates. The preprocessing function is applied to the (time, position)
    cursor or eye data.

    See also:
        :func:`~aopy.data.bmi3d.get_kinematic_segment`, :func:`~aopy.data.bmi3d.get_kinematics`
    
    Example:
        subject = 'beignet'
        te_id = 4301
        date = '2021-01-01'
        trial_filter = lambda t: TRIAL_END not in t
        trajectories, segments = get_kinematic_segments(preproc_dir, subject, te_id, date,
                                                       [CURSOR_ENTER_CENTER_TARGET], 
                                                       [REWARD, TRIAL_END], 
                                                       trial_filter=trial_filter) 
    
    Args:
        preproc_dir (str): base directory where the files live
        subject (str): Subject name
        te_id (int): Block number of Task entry object 
        date (str): Date of recording
        trial_start_codes (list): list of numeric codes representing the start of a trial
        trial_end_codes (list): list of numeric codes representing the end of a trial
        trial_filter (fn, optional): function mapping trial segments to boolean values. Any trials
            for which the filter returns False will not be included in the output
        preproc (fn, optional): function mapping (position, samplerate) data to kinematics. For example,
            a smoothing function or an estimate of velocity from position
        datatype (str, optional): type of kinematics to load. Defaults to 'cursor'.    
        samplerate (float, optional): optionally choose the samplerate of the data in Hz. Default 1000.
        kwargs: additional keyword arguments to pass to get_kinematics
    
    Returns:
        tuple: tuple containing:
            | **trajectories (ntrial):** array of filtered cursor trajectories for each trial
            | **trial_segments (ntrial):** array of numeric code segments for each trial   

    Note:
        The sampling rate of the returned data might be different from the requested sampling rate if the
        preprocessing function does any modification to the length of the data.

    Modified September 2023 to include optional sampling rate argument     
    '''
    data, metadata = load_preproc_exp_data(preproc_dir, subject, te_id, date)
    event_codes = data['events']['code']
    event_times = data['events']['timestamp']
    trial_segments, trial_times = get_trial_segments(event_codes, event_times, 
                                                                  trial_start_codes, trial_end_codes)
    segments = [
        get_kinematic_segment(preproc_dir, subject, te_id, date, t[0], t[1], samplerate, preproc, datatype, **kwargs)[0] 
        for t in trial_times
    ]
    trajectories = np.array(segments, dtype='object')
    trial_segments = np.array(trial_segments, dtype='object')
    success_trials = [trial_filter(t) for t in trial_segments]
    
    return trajectories[success_trials], trial_segments[success_trials]

def get_lfp_segments(preproc_dir, subject, te_id, date, trial_start_codes, trial_end_codes, 
                           trial_filter=lambda x:True):
    '''
    Loads lfp segments (different length for each trial) from a preprocessed HDF file. Trials can
    be specified by numeric start and end codes. Trials can also be filtered so that only successful
    trials are included, for example. The filter is applied to numeric code segments for each trial. 
        
    Args:
        preproc_dir (str): path to the preprocessed directory
        preproc_dir (str): base directory where the files live
        subject (str): Subject name
        te_id (int): Block number of Task entry object 
        date (str): Date of recording
        trial_start_codes (list): list of numeric codes representing the start of a trial
        trial_end_codes (list): list of numeric codes representing the end of a trial
        trial_filter (fn, optional): function mapping trial segments to boolean values. Any trials
            for which the filter returns False will not be included in the output
    
    Returns:
        tuple: tuple containing:
            | **lfp_segments (ntrial):** array of filtered lfp segments for each trial
            | **trial_segments (ntrial):** array of numeric code segments for each trial

    '''
    data, metadata = load_preproc_exp_data(preproc_dir, subject, te_id, date)
    lfp_data, lfp_metadata = load_preproc_lfp_data(preproc_dir, subject, te_id, date)
    samplerate = lfp_metadata['lfp_samplerate']

    event_codes = data['events']['code']
    event_times = data['events']['timestamp']

    trial_segments, trial_times = get_trial_segments(event_codes, event_times, 
                                                                  trial_start_codes, trial_end_codes)
    lfp_segments = np.array(get_data_segments(lfp_data, trial_times, samplerate), dtype='object')
    trial_segments = np.array(trial_segments, dtype='object')
    success_trials = [trial_filter(t) for t in trial_segments]
    
    return lfp_segments[success_trials], trial_segments[success_trials]


def get_lfp_aligned(preproc_dir, subject, te_id, date, trial_start_codes, trial_end_codes, 
                           time_before, time_after, trial_filter=lambda x:True):
    '''
    Loads lfp data (same length for each trial) from a preprocessed HDF file. Trials can
    be specified by numeric start and end codes. Trials can also be filtered so that only successful
    trials are included, for example. The filter is applied to numeric code segments for each trial. 

    Args:
        preproc_dir (str): base directory where the files live
        subject (str): Subject name
        te_id (int): Block number of Task entry object 
        date (str): Date of recording
        trial_start_codes (list): list of numeric codes representing the start of a trial
        trial_end_codes (list): list of numeric codes representing the end of a trial
        time_before (float): time before the trial start to include in the aligned lfp (in seconds)
        time_after (float): time after the trial end to include in the aligned lfp (in seconds)
        trial_filter (fn, optional): function mapping trial segments to boolean values. Any trials
            for which the filter returns False will not be included in the output
    
    Returns:
        (ntrials, nt, nch): aligned lfp data output from `func:aopy.preproc.trial_align_data`


    '''
    data, metadata = load_preproc_exp_data(preproc_dir, subject, te_id, date)
    lfp_data, lfp_metadata = load_preproc_lfp_data(preproc_dir, subject, te_id, date)
    samplerate = lfp_metadata['lfp_samplerate']

    event_codes = data['events']['code']
    event_times = data['events']['timestamp']

    trial_segments, trial_times = get_trial_segments(event_codes, event_times, 
                                                     trial_start_codes, trial_end_codes)
    trial_start_times = [t[0] for t in trial_times]
    assert len(trial_start_times) > 0, "No trials found"
    trial_aligned_data = trial_align_data(lfp_data, trial_start_times, time_before, time_after, samplerate) # (nt, nch, ntrial)
    success_trials = [trial_filter(t) for t in trial_segments]
    
    return trial_aligned_data[:,:,success_trials]

def get_ts_data_trial(preproc_dir, subject, te_id, date, trigger_time, time_before, time_after,
                      channels=None, datatype='lfp'):
    '''
    Simple wrapper around load_hdf_ts_trial for lfp or broadband data.
    
    Args:
        preproc_dir (str): base directory where the files live
        subject (str): Subject name
        te_id (int): Block number of Task entry object 
        date (str): Date of recording
        trigger_time (float): time (in seconds) in the recording at which the desired segment starts
        time_before (float): time (in seconds) to include before the trigger times
        time_after (float): time (in seconds) to include after the trigger times
        channels (int array, optional): which channel indices to load
        datatype (str, optional): choice of 'lfp' or 'broadband' data to load. Defaults to 'lfp'.    

    Returns:
        tuple: tuple containing:
            | **segment (nt, nch):** data segment from the given preprocessed file
            | **samplerate (float):** sampling rate of the returned data
    '''
    if datatype == 'lfp':
        data_group='/'
        data_name='lfp_data'
        metadata_group='lfp_metadata'
        samplerate_key='lfp_samplerate'
    elif datatype == 'broadband':
        data_group='/'
        data_name='broadband_data'
        metadata_group='broadband_metadata'
        samplerate_key='samplerate'
    filename = get_preprocessed_filename(subject, te_id, date, datatype)
    preproc_dir = os.path.join(preproc_dir, subject)

    try:
        samplerate = load_hdf_data(preproc_dir, filename, samplerate_key, metadata_group, cached=True)
        data = load_hdf_ts_trial(preproc_dir, filename, data_group, data_name, 
                                 samplerate, trigger_time, time_before, time_after, channels=channels)
    except FileNotFoundError as e:
        print(f"No data found in {preproc_dir} for subject {subject} on {date} ({te_id})")
        raise e

    return data, samplerate

def get_ts_data_segment(preproc_dir, subject, te_id, date, start_time, end_time,
                        channels=None, datatype='lfp'):
    '''
    Simple wrapper around load_hdf_ts_segment for lfp or broadband data.
    
    Args:
        preproc_dir (str): base directory where the files live
        subject (str): Subject name
        te_id (int): Block number of Task entry object 
        date (str): Date of recording
        trigger_time (float): time (in seconds) in the recording at which the desired segment starts
        time_before (float): time (in seconds) to include before the trigger times
        time_after (float): time (in seconds) to include after the trigger times
        channels (int array, optional): which channel indices to load
        datatype (str, optional): choice of 'lfp' or 'broadband' data to load. Defaults to 'lfp'.    

    Returns:
        tuple: tuple containing:
            | **segment (nt, nch):** data segment from the given preprocessed file
            | **samplerate (float):** sampling rate of the returned data
    '''
    if datatype == 'lfp':
        data_group='/'
        data_name='lfp_data'
        metadata_group='lfp_metadata'
        samplerate_key='lfp_samplerate'
    elif datatype == 'broadband':
        data_group='/'
        data_name='broadband_data'
        metadata_group='broadband_metadata'
        samplerate_key='samplerate'
    filename = get_preprocessed_filename(subject, te_id, date, datatype)
    preproc_dir = os.path.join(preproc_dir, subject)

    try:
        samplerate = load_hdf_data(preproc_dir, filename, samplerate_key, metadata_group)
        data = load_hdf_ts_segment(preproc_dir, filename, data_group, data_name, 
                                   samplerate, start_time, end_time, channels=channels)
    except FileNotFoundError as e:
        print(f"No data found in {preproc_dir} for subject {subject} on {date} ({te_id})")
        raise e

    return data, samplerate

    
def get_target_locations(preproc_dir, subject, te_id, date, target_indices):
    '''
    Loads the x,y,z location of targets in a preprocessed HDF file given by their index. Requires
    that the preprocessed `exp_data` includes a `trials` structured array containing `index` and 
    `target` fields (the default behavior of `:func:~aopy.preproc.proc_exp`)
    
    Args:
        preproc_dir (str): base directory where the files live
        subject (str): Subject name
        te_id (int): Block number of Task entry object 
        date (str): Date of recording
        target_indices (ntarg): a list of which targets to fetch
        
    Returns:
        ndarray: (ntarg x 3) array of coordinates of the given targets
    '''
    data, metadata = load_preproc_exp_data(preproc_dir, subject, te_id, date)
    try:
        trials = data['trials']
    except:
        trials = data['bmi3d_trials']
    locations = np.nan*np.zeros((len(target_indices), 3))
    for i in range(len(target_indices)):
        trial_idx = np.where(trials['index'] == target_indices[i])[0]
        if len(trial_idx) > 0:
            locations[i,:] = trials['target'][trial_idx[0]][[0,2,1]] # use x,y,z format
        else:
            raise ValueError(f"Target index {target_indices[i]} not found")
    return np.round(locations,4)

def get_trajectory_frequencies(preproc_dir, subject, te_id, date):
    '''
    For continuous tracking tasks, get the set of frequencies (in Hz) used to 
    generate the trajectories that were preesented on each trial of the experiment, 
    using :func:`~aopy.preproc.bmi3d.get_ref_dis_frequencies`. 

    Args:
        preproc_dir (str): base directory where the files live
        subject (str): Subject name
        te_id (int): Block number of Task entry object 
        date (str): Date of recording

    Returns:
        tuple: Tuple containing:
            | **freq_r (list of arrays):** (ntrial) list of (nfreq,) frequencies used to generate reference trajectory
            | **freq_d (list of arrays):** (ntrial) list of (nfreq,) frequencies used to generate disturbance trajectory
    '''
    data, metadata = load_preproc_exp_data(preproc_dir, subject, te_id, date)
    freq_r, freq_d = get_ref_dis_frequencies(data, metadata)
    return freq_r, freq_d

def get_source_files(preproc_dir, subject, te_id, date):
    '''
    Retrieves the dictionary of source files from a preprocessed file

    Args:
        preproc_dir (str): base directory where the files live
        subject (str): Subject name
        te_id (int): Block number of Task entry object 
        date (str): Date of recording

    Returns:
        tuple: tuple containing:
            | ** files (dict):** dictionary of (source, filepath) files that are associated with the given experiment
            | ** data_dir (str):** directory where the source files were located
    '''
    exp_data, exp_metadata = load_preproc_exp_data(preproc_dir, subject, te_id, date)
    return exp_metadata['source_files'], exp_metadata['source_dir']

def tabulate_behavior_data(preproc_dir, subjects, ids, dates, trial_start_codes, 
                           trial_end_codes, reward_codes, penalty_codes, metadata=[],
                           df=None, event_code_type='code', return_bad_entries=False):
    '''
    Concatenate trials from across experiments. Experiments are given as lists of 
    subjects, task entry ids, and dates. Each list must be the same length. Trials 
    are defined by intervals between the given trial start and end codes. 

    Args:
        preproc_dir (str): base directory where the files live
        subjects (list of str): Subject name for each recording
        ids (list of int): Block number of Task entry object for each recording
        dates (list of str): Date for each recording
        trial_start_codes (list): list of numeric codes representing the start of a trial
        trial_end_codes (list): list of numeric codes representing the end of a trial
        reward_codes (list): list of numeric codes representing rewards
        penalty_codes (list): list of numeric codes representing penalties
        metadata (list, optional): list of metadata keys that should be included in the df
        df (DataFrame, optional): pandas DataFrame object to append. Defaults to None.
        event_code_type (str, optional): type of event codes to use. Defaults to 'code'. Other
            choices include 'event' and 'data'.
        return_bad_entries (bool, optional): If True, returns the list of task entries that could 
            not be loaded. Defaults to False.

    Returns:
        pd.DataFrame: pandas DataFrame containing the concatenated trial data with columns:
            | **subject (str):** subject name
            | **te_id (str):** task entry id
            | **date (str):** date of recording
            | **event_codes (ntrial):** numeric code segments for each trial (specified by `event_code_type`)
            | **event_times (ntrial):** time segments for each trial
            | **reward (ntrial):** boolean values indicating whether each trial was rewarded
            | **penalty (ntrial):** boolean values indicating whether each trial was penalized
            | **%metadata_key% (ntrial):** requested metadata values for each key requested
    '''
    bad_entries = []
    if df is None:
        df = pd.DataFrame()

    entries = list(zip(subjects, dates, ids))
    for subject, date, te in tqdm(entries): 

        # Load data from bmi3d hdf 
        try:
            exp_data, exp_metadata = load_preproc_exp_data(preproc_dir, subject, te, date)
        except:
            print(f"Entry {subject} {date} {te} could not be loaded.")
            traceback.print_exc()
            bad_entries.append([subject,date,te])
            continue
        event_codes = exp_data['events'][event_code_type]
        event_times = exp_data['events']['timestamp']

        # Trial aligned event codes and event times
        tr_seg, tr_t = get_trial_segments_and_times(event_codes, event_times, trial_start_codes, trial_end_codes)
        reward = [np.any(np.isin(reward_codes, ec)) for ec in tr_seg]
        penalty = [np.any(np.isin(penalty_codes, ec)) for ec in tr_seg]
        
        # Build a dataframe for this task entry
        exp = {
            'subject': subject,
            'te_id': te, 
            'date': date, 
            'event_codes': tr_seg,
            'event_times': tr_t, 
            'reward': reward,
            'penalty': penalty,
        }

        # Add requested metadata
        for key in metadata:
            if key in exp_metadata:
                exp[key] = [exp_metadata[key] for _ in range(len(tr_seg))]
            else:
                exp[key] = None
                print(f"Entry {subject} {date} {te} does not have metadata {key}.")

        # Concatenate with existing dataframes
        df = pd.concat([df,pd.DataFrame(exp)], ignore_index=True)
    
    if return_bad_entries:
        return df, bad_entries
    else:
        return df

def tabulate_behavior_data_flash(preproc_dir, subjects, ids, dates, metadata=[], 
                                      df=None):
    '''
    Wrapper around tabulate_behavior_data() specifically for flash experiments. 
    Uses the task event names (b'TARGET_ON', b'REWARD', and b'TRIAL_END', specifically)
    to find start and end times for flash experiments.

    Args:
        preproc_dir (str): base directory where the files live
        subjects (list of str): Subject name for each recording
        ids (list of int): Block number of Task entry object for each recording
        dates (list of str): Date for each recording
        metadata (list, optional): list of metadata keys that should be included in the df
        df (DataFrame, optional): pandas DataFrame object to append. Defaults to None.

    Returns:
        pd.DataFrame: pandas DataFrame containing the concatenated trial data with columns:
            | **subject (str):** subject name
            | **te_id (str):** task entry id
            | **date (str):** date of recording
            | **event_names (ntrial):** event name segments for each trial
            | **event_times (ntrial):** time segments for each trial
            | **%metadata_key% (ntrial):** requested metadata values for each key requested
            | **flash_start_time (ntrial):** time the flash started
            | **flash_end_time (ntrial):** time the flash ended
'''
    # Use default "trial" definition
    trial_end_codes = [b'TRIAL_END']
    trial_start_codes = [b'TARGET_ON']
    reward_codes = [b'REWARD']
    penalty_codes = []

    # Concatenate base trial data
    new_df = tabulate_behavior_data(
        preproc_dir, subjects, ids, dates, trial_start_codes, trial_end_codes, 
        reward_codes, penalty_codes, metadata, df=None, event_code_type='event')
    if len(new_df) == 0:
        warnings.warn("No trials found")
        return df
    
    # Remove any unrewarded trials and then get rid of the 'reward' column
    new_df.drop(new_df.index[~new_df['reward']], inplace=True)
    new_df.drop(['reward', 'penalty'], axis=1, inplace=True)
    new_df.reset_index()
    
    # Add trial segment timing
    new_df['flash_start_time'] = np.nan
    new_df['flash_end_time'] = np.nan
    new_df['prev_trial_end_time'] = np.nan
    new_df['trial_end_time'] = np.nan
    for i in range(len(new_df)):
        event_times = new_df.loc[i, 'event_times']
        
        # Trial end times
        if i > 0 and new_df.loc[i-1, 'event_times'][-1] < event_times[0]:
            new_df.loc[i, 'prev_trial_end_time'] = new_df.loc[i-1, 'event_times'][-1]
        else:
            new_df.loc[i, 'prev_trial_end_time'] = 0.
        if i < len(new_df)-1 and new_df.loc[i+1, 'event_times'][0] > event_times[-1]:
            new_df.loc[i, 'trial_end_time'] = new_df.loc[i+1, 'event_times'][0]
        else:
            new_df.loc[i, 'trial_end_time'] = event_times[-1]

        # Flash starts when trial starts
        new_df.loc[i, 'flash_start_time'] = event_times[0]
        new_df.loc[i, 'flash_end_time'] = event_times[2]
            
    df = pd.concat([df, new_df], ignore_index=True)
    return df
    
def tabulate_behavior_data_center_out(preproc_dir, subjects, ids, dates, metadata=[], 
                                      df=None):
    '''
    Wrapper around tabulate_behavior_data() specifically for center-out experiments. 
    Makes use of the task codes saved in `/config/task_codes.yaml` to automatically 
    assign event codes for trial start, trial end, reward, penalty, and targets. 

    Args:
        preproc_dir (str): base directory where the files live
        subjects (list of str): Subject name for each recording
        ids (list of int): Block number of Task entry object for each recording
        dates (list of str): Date for each recording
        metadata (list, optional): list of metadata keys that should be included in the df
        df (DataFrame, optional): pandas DataFrame object to append. Defaults to None.

    Returns:
        pd.DataFrame: pandas DataFrame containing the concatenated trial data with columns:
            | **subject (str):** subject name
            | **te_id (str):** task entry id
            | **date (str):** date of recording
            | **event_codes (ntrial):** numeric code segments for each trial
            | **event_times (ntrial):** time segments for each trial
            | **reward (ntrial):** boolean values indicating whether each trial was rewarded
            | **penalty (ntrial):** boolean values indicating whether each trial was penalized
            | **%metadata_key% (ntrial):** requested metadata values for each key requested
            | **target_idx (ntrial):** index of the target that was presented
            | **target_location (ntrial):** location of the target that was presented
            | **trial_start_time (ntrial):** time at which the trial started
            | **trial_end_time (ntrial):** time at which the trial ended
            | **trial_initiated (ntrial):** boolean values indicating whether the trial was initiated
            | **hold_start_time (ntrial):** time at which the hold period started
            | **hold_completed (ntrial):** boolean values indicating whether the hold period was completed
            | **delay_start_time (ntrial):** time at which the delay period started
            | **delay_completed (ntrial):** boolean values indicating whether the delay period was completed
            | **go_cue_time (ntrial):** time at which the go cue was presented
            | **reach_completed (ntrial):** boolean values indicating whether the reach was completed
            | **reach_end_time (ntrial):** time at which the reach was completed
            | **reward_start_time (ntrial):** time at which the reward was presented
            | **reward_end_time (ntrial):** time at which the reward was completed
            | **penalty_start_time (ntrial):** time at which the penalty was presented
            | **penalty_end_time (ntrial):** time at which the penalty was completed
            | **penalty_event (ntrial):** numeric code for the penalty event
    '''
    # Use default "trial" definition
    task_codes = load_bmi3d_task_codes()
    trial_end_codes = [task_codes['TRIAL_END']]
    trial_start_codes = [task_codes['CENTER_TARGET_ON']]
    reward_codes = [task_codes['REWARD']]
    penalty_codes = [task_codes['HOLD_PENALTY'], task_codes['TIMEOUT_PENALTY']]
    target_codes = task_codes['PERIPHERAL_TARGET_ON']

    # Concatenate base trial data
    new_df = tabulate_behavior_data(
        preproc_dir, subjects, ids, dates, trial_start_codes, trial_end_codes, 
        reward_codes, penalty_codes, metadata, df=None)
    if len(new_df) == 0:
        warnings.warn("No trials found")
        return df

    # Add target info
    target_idx = [
        code[np.isin(code, target_codes)][0] - target_codes[0] + 1 # add 1 for center target
        if np.sum(np.isin(code, target_codes)) > 0 else 0 
        for code 
        in new_df['event_codes']
    ]
    target_location = [
        np.squeeze(get_target_locations(preproc_dir, s, te, d, [t_idx]))
        for s, te, d, t_idx 
        in zip(new_df['subject'], new_df['te_id'], new_df['date'], target_idx)
    ]
    new_df['target_idx'] = target_idx
    new_df['target_location'] = target_location

    # Add trial segment timing
    new_df['prev_trial_end_time'] = np.nan
    new_df['trial_end_time'] = np.nan
    new_df['center_target_on_time'] = np.nan
    new_df['trial_initiated'] = False
    new_df['hold_start_time'] = np.nan
    new_df['hold_completed'] = False
    new_df['delay_start_time'] = np.nan
    new_df['delay_completed'] = False
    new_df['go_cue_time'] = np.nan
    new_df['reach_completed'] = False
    new_df['reach_end_time'] = np.nan
    new_df['reward_start_time'] = np.nan
    new_df['penalty_start_time'] = np.nan
    new_df['penalty_event'] = np.nan
    for i in range(len(new_df)):
        event_codes = new_df.loc[i, 'event_codes']
        event_times = new_df.loc[i, 'event_times']

        # Trial end times
        if i > 0 and new_df.loc[i-1, 'event_times'][-1] < event_times[0]:
            new_df.loc[i, 'prev_trial_end_time'] = new_df.loc[i-1, 'event_times'][-1]
        else:
            new_df.loc[i, 'prev_trial_end_time'] = 0.
        new_df.loc[i, 'trial_end_time'] = event_times[-1]

        # Center target appears
        new_df.loc[i, 'center_target_on_time'] = event_times[0]
        
        # Trial initiated if cursor enters the center target
        hold_times = event_times[np.isin(event_codes, [task_codes['CURSOR_ENTER_CENTER_TARGET']])]
        new_df.loc[i, 'trial_initiated'] = len(hold_times) > 0
        if new_df.loc[i, 'trial_initiated']:
            new_df.loc[i, 'hold_start_time'] = hold_times[0]

        # Hold completed if peripheral target turns on (start of delay)
        delay_times = event_times[np.isin(event_codes, task_codes['PERIPHERAL_TARGET_ON'])]
        new_df.loc[i, 'hold_completed'] = len(delay_times) > 0
        if new_df.loc[i, 'hold_completed']:
            new_df.loc[i, 'delay_start_time'] = delay_times[0]

        # Delay completed when center target turns off (go cue)
        go_cue_times = event_times[np.isin(event_codes, task_codes['CENTER_TARGET_OFF'])]
        new_df.loc[i, 'delay_completed'] = len(go_cue_times) > 0
        if new_df.loc[i, 'delay_completed']:
            new_df.loc[i, 'go_cue_time'] = go_cue_times[0]

        # Reach completed if cursor enters target (regardless of whether the trial was successful)
        reach_times = event_times[np.isin(event_codes, task_codes['CURSOR_ENTER_PERIPHERAL_TARGET'])]
        new_df.loc[i, 'reach_completed'] = len(reach_times) > 0
        if new_df.loc[i, 'reach_completed']:
            new_df.loc[i, 'reach_end_time'] = reach_times[0]

        # Reward start times
        reward_times = event_times[np.isin(event_codes, task_codes['REWARD'])]      
        if len(reward_times) > 0:
            new_df.loc[i, 'reward_start_time'] = reward_times[0]

        # Penalty start times
        penalty_idx = np.isin(event_codes, penalty_codes)
        penatly_codes = event_codes[penalty_idx]
        penalty_times = event_times[penalty_idx]
        if len(penalty_times) > 0:
            new_df.loc[i, 'penalty_start_time'] = penalty_times[0]
            new_df.loc[i, 'penalty_event'] = penatly_codes[0]

    df = pd.concat([df, new_df], ignore_index=True)
    return df

def tabulate_behavior_data_out(preproc_dir, subjects, ids, dates, metadata=[], 
                               df=None):
    '''
    Wrapper around tabulate_behavior_data() specifically for out experiments (similar to
    center-out but without a trial-initiating center target). Makes use of the task codes 
    saved in `/config/task_codes.yaml` to automatically assign event codes for trial start, 
    trial end, reward, penalty, and targets. 

    Args:
        preproc_dir (str): base directory where the files live
        subjects (list of str): Subject name for each recording
        ids (list of int): Block number of Task entry object for each recording
        dates (list of str): Date for each recording
        metadata (list, optional): list of metadata keys that should be included in the df
        df (DataFrame, optional): pandas DataFrame object to append. Defaults to None.

    Returns:
        pd.DataFrame: pandas DataFrame containing the concatenated trial data with columns:
            | **subject (str):** subject name
            | **te_id (str):** task entry id
            | **date (str):** date of recording
            | **event_codes (ntrial):** numeric code segments for each trial
            | **event_times (ntrial):** time segments for each trial
            | **reward (ntrial):** boolean values indicating whether each trial was rewarded
            | **penalty (ntrial):** boolean values indicating whether each trial was penalized
            | **%metadata_key% (ntrial):** requested metadata values for each key requested
            | **target_idx (ntrial):** index of the target that was presented
            | **target_location (ntrial):** location of the target that was presented
            | **trial_start_time (ntrial):** time at which the trial started
            | **trial_end_time (ntrial):** time at which the trial ended
            | **reach_completed (ntrial):** boolean values indicating whether the reach was completed
            | **reach_end_time (ntrial):** time at which the reach was completed
            | **reward_start_time (ntrial):** time at which the reward was presented
            | **reward_end_time (ntrial):** time at which the reward was completed
            | **penalty_start_time (ntrial):** time at which the penalty was presented
            | **penalty_end_time (ntrial):** time at which the penalty was completed
            | **penalty_event (ntrial):** numeric code for the penalty event
    '''
    # Use default "trial" definition
    task_codes = load_bmi3d_task_codes()
    trial_end_codes = [task_codes['TRIAL_END']]
    trial_start_codes = task_codes['PERIPHERAL_TARGET_ON']
    reward_codes = [task_codes['REWARD']]
    penalty_codes = [task_codes['HOLD_PENALTY'], task_codes['TIMEOUT_PENALTY']]
    target_codes = task_codes['PERIPHERAL_TARGET_ON']

    # Concatenate base trial data
    new_df = tabulate_behavior_data(
        preproc_dir, subjects, ids, dates, trial_start_codes, trial_end_codes, 
        reward_codes, penalty_codes, metadata, df=None)
    if len(new_df) == 0:
        warnings.warn("No trials found")
        return df

    # Add target info
    target_idx = [
        code[np.isin(code, target_codes)][0] - target_codes[0] + 1 # add 1 for center target
        if np.sum(np.isin(code, target_codes)) > 0 else 0 
        for code 
        in new_df['event_codes']
    ]
    target_location = [
        np.squeeze(get_target_locations(preproc_dir, s, te, d, [t_idx]))
        for s, te, d, t_idx 
        in zip(new_df['subject'], new_df['te_id'], new_df['date'], target_idx)
    ]
    new_df['target_idx'] = target_idx
    new_df['target_location'] = target_location

    # Add trial segment timing
    new_df['prev_trial_end_time'] = np.nan
    new_df['trial_end_time'] = np.nan
    new_df['target_on_time'] = np.nan
    new_df['reach_completed'] = False
    new_df['reach_end_time'] = np.nan
    new_df['reward_start_time'] = np.nan
    new_df['reward_end_time'] = np.nan
    new_df['penalty_start_time'] = np.nan
    new_df['penalty_end_time'] = np.nan
    new_df['penalty_event'] = np.nan
    for i in range(len(new_df)):
        event_codes = new_df.loc[i, 'event_codes']
        event_times = new_df.loc[i, 'event_times']

        # Trial end times
        if i > 0 and new_df.loc[i-1, 'event_times'][-1] < event_times[0]:
            new_df.loc[i, 'prev_trial_end_time'] = new_df.loc[i-1, 'event_times'][-1]
        else:
            new_df.loc[i, 'prev_trial_end_time'] = 0.
        new_df.loc[i, 'trial_end_time'] = event_times[-1]

        # Target appears
        new_df.loc[i, 'target_on_time'] = event_times[0]

        # Reach completed if cursor enters target (regardless of whether the trial was successful)
        reach_times = event_times[np.isin(event_codes, task_codes['CURSOR_ENTER_PERIPHERAL_TARGET'])]
        new_df.loc[i, 'reach_completed'] = len(reach_times) > 0
        if new_df.loc[i, 'reach_completed']:
            new_df.loc[i, 'reach_end_time'] = reach_times[0]

        # Reward start times
        reward_times = event_times[np.isin(event_codes, task_codes['REWARD'])]      
        if len(reward_times) > 0:
            new_df.loc[i, 'reward_start_time'] = reward_times[0]

        # Penalty start times
        penalty_idx = np.isin(event_codes, penalty_codes)
        penatly_codes = event_codes[penalty_idx]
        penalty_times = event_times[penalty_idx]
        if len(penalty_times) > 0:
            new_df.loc[i, 'penalty_start_time'] = penalty_times[0]
            new_df.loc[i, 'penalty_event'] = penatly_codes[0]

    df = pd.concat([df, new_df], ignore_index=True)
    return df

def tabulate_behavior_data_tracking_task(preproc_dir, subjects, ids, dates, metadata=[], df=None):
    '''
    Wrapper around tabulate_behavior_data() specifically for tracking task experiments. 
    Makes use of the task codes saved in `/config/task_codes.yaml` to automatically 
    assign event codes for trial start, trial end, reward, penalty.
    
    Args:
        preproc_dir (str): base directory where the files live
        subjects (list of str): Subject name for each recording
        ids (list of int): Block number of Task entry object for each recording
        dates (list of str): Date for each recording
        metadata (list, optional): list of metadata keys that should be included in the df
        df (DataFrame, optional): pandas DataFrame object to append. Defaults to None.

    Returns:
        pd.DataFrame: pandas DataFrame containing the concatenated trial data with columns:
            | **subject (str):** subject name
            | **te_id (str):** task entry id
            | **date (str):** date of recording
            | **event_codes (ntrial):** numeric code segments for each trial
            | **event_times (ntrial):** time segments for each trial
            | **reward (ntrial):** boolean values indicating whether each trial was rewarded
            | **penalty (ntrial):** boolean values indicating whether each trial was penalized
            | **%metadata_key% (ntrial):** requested metadata values for each key requested
            | **sequence_params (ntrial):** string of params used to generate all trajectories in the same task entry
            | **ref_freqs (ntrial):** array of frequencies used to generate reference trajectory for each trial
            | **dis_freqs (ntrial):** array of frequencies used to generate disturbance trajectory for each trial
            | **trial_initiated (ntrial):** boolean values indicating whether the trial was initiated (i.e. hold was attempted)
            | **hold_start_time (ntrial):** time at which the hold period started
            | **hold_completed (ntrial):** boolean values indicating whether the hold period was completed
            | **tracking_start_time (ntrial):** time at which the hold period ended and tracking started
            | **tracking_end_time (ntrial):** time at which tracking ended (whether with a reward or tracking out penalty)
            | **trajectory_start_time (ntrial):** time at which the ref & dis trajectories started (excluding the ramp up period)
            | **trajectory_end_time (ntrial):** time at which the ref & dis trajectories ended (excluding the ramp down period if the trial was rewarded)
    '''
    # Use default "trial" definition
    task_codes = load_bmi3d_task_codes()
    trial_start_codes = [task_codes['CENTER_TARGET_ON']]
    trial_end_codes = [task_codes['TRIAL_END'], task_codes['PAUSE_START'], task_codes['PAUSE']]
    reward_codes = [task_codes['REWARD']]
    penalty_codes = [task_codes['TIMEOUT_PENALTY'], task_codes['HOLD_PENALTY'], task_codes['OTHER_PENALTY']]
    
    # Concatenate base trial data
    if 'sequence_params' not in metadata:
        metadata.append('sequence_params')
    new_df, bad_entries = tabulate_behavior_data(
        preproc_dir, subjects, ids, dates, trial_start_codes, trial_end_codes, 
        reward_codes, penalty_codes, metadata, df=None, return_bad_entries=True)
    
    # Add frequency content of reference & disturbance trajectories
    ref_freqs = []
    dis_freqs = []
    for subject, te, date in zip(subjects, ids, dates):
        if [subject,date,te] in bad_entries:
            continue
        r, d = get_trajectory_frequencies(preproc_dir, subject, te, date)
        ref_freqs.extend(r)
        dis_freqs.extend(d)
    new_df['ref_freqs'] = ref_freqs
    new_df['dis_freqs'] = dis_freqs
    
    # Add trial segment timing
    new_df['trial_initiated'] = np.zeros(len(new_df), dtype='bool')
    new_df['hold_start_time'] = np.nan*np.zeros(len(new_df))
    new_df['hold_completed'] = np.zeros(len(new_df), dtype='bool')
    new_df['tracking_start_time'] = np.nan*np.zeros(len(new_df))
    new_df['tracking_end_time'] = np.nan*np.zeros(len(new_df))
    new_df['trajectory_start_time'] = np.nan*np.zeros(len(new_df))
    new_df['trajectory_end_time'] = np.nan*np.zeros(len(new_df))
    ramp = [
        json.loads(params)['ramp']
        if 'ramp' in json.loads(params) else 0
        for params
        in new_df['sequence_params']
    ]
    ramp_down = [
        json.loads(params)['ramp_down']
        if 'ramp_down' in json.loads(params) else 0
        for params
        in new_df['sequence_params']
    ]
    for i in range(len(new_df)):
        event_codes = new_df.iloc[i]['event_codes']
        event_times = new_df.iloc[i]['event_times']

        # Trial initiated when hold begins
        _, initiation_times = get_trial_segments(event_codes, event_times,
                                            [task_codes['CENTER_TARGET_ON']], [task_codes['TRIAL_START']])
        new_df.loc[i,'trial_initiated'] = len(initiation_times) > 0 # if False, TIMEOUT_PENALTY
        if new_df.loc[i,'trial_initiated']:
            new_df.loc[i,'hold_start_time'] = initiation_times[0][-1]

        # Hold completed when tracking begins (first time cursor enters target)
        tracking_events, tracking_times = get_trial_segments_and_times(event_codes, event_times,
                                            [task_codes['TRIAL_START']], [task_codes['REWARD'],task_codes['OTHER_PENALTY']])
        new_df.loc[i,'hold_completed'] = len(tracking_times) > 0 # if False, HOLD_PENALTY
        if new_df.loc[i,'hold_completed']:
            new_df.loc[i,'tracking_start_time'] = tracking_times[0][1] # TRIAL_START, CURSOR_ENTER_TARGET ... REWARD/OTHER_PENALTY
            new_df.loc[i,'tracking_end_time'] = tracking_times[0][-1]

            # Trajectory excludes ramp period at beginning of tracking (regardless of tracking success)
            new_df.loc[i,'trajectory_start_time'] = new_df.loc[i,'tracking_start_time']+ramp[i]
            new_df.loc[i,'trajectory_end_time'] = new_df.loc[i,'tracking_end_time']

            # If tracking successful, trajectory excludes ramp period at end of tracking
            if new_df.loc[i,'reward']:
                new_df.loc[i,'trajectory_end_time'] = new_df.loc[i,'tracking_end_time']-ramp_down[i]

    df = pd.concat([df, new_df], ignore_index=True)
    return df

def tabulate_kinematic_data(preproc_dir, subjects, te_ids, dates, start_times, end_times, 
                            samplerate=1000, preproc=None, datatype='cursor', **kwargs):
    '''
    Grab kinematics data from trials across arbitrary preprocessed files.

    Args:
        preproc_dir (str): base directory where the files live
        subjects (list of str): Subject name for each recording
        ids (list of int): Block number of Task entry object for each recording
        dates (list of str): Date for each recording
        start_times (list of float): times in the recording at which the desired segments starts
        end_times (list of float): times in the recording at which the desired segments ends
        samplerate (float, optional): optionally choose the samplerate of the data in Hz. Default 1000.
        preproc (fn, optional): function mapping (position, fs) data to (kinematics, fs_new). For example,
            a smoothing function or an estimate of velocity from position
        datatype (str, optional): type of kinematics to tabulate. Defaults to 'cursor'.  
        kwargs (dict, optional): optional keyword arguments to pass to :func:`~aopy.preproc.get_kinematic_segment`  

    Returns:
        (ntrial,): list of tensors of (nt, nch) kinematics from each trial
    '''

    assert len(subjects) == len(te_ids) == len(dates) == len(start_times) == len(end_times)

    segments = [get_kinematic_segment(preproc_dir, s, t, d, ts, te, samplerate, preproc, datatype, **kwargs)[0] 
                for s, t, d, ts, te in zip(subjects, te_ids, dates, start_times, end_times)]
    trajectories = np.array(segments, dtype='object')
    return trajectories

def tabulate_stim_data(preproc_dir, subjects, ids, dates, metadata=['stimulation_site'], 
                       debug=True, df=None, **kwargs):
    '''
    Concatenate stimulation data from across experiments. Experiments are given as lists of 
    subjects, task entry ids, and dates. Each list must be the same length. 
    
    Args:
        preproc_dir (str): base directory where the files live
        subjects (list of str): Subject name for each recording
        ids (list of int): Block number of Task entry object for each recording
        dates (list of str): Date for each recording
        metadata (list, optional): list of metadata keys that should be included in the df. By default,
            only `stimulation_site` is included.
        debug (bool, optional): Passed to :func:`~aopy.preproc.laser.find_stim_times`, if True prints
            an laser sensor alignment plot for each trial. Defaults to True.
        df (DataFrame, optional): pandas DataFrame object to append. Defaults to None.
        kwargs (dict, optional): optional keyword arguments to pass to :func:`~aopy.preproc.laser.find_stim_times`

    Returns:
        pd.DataFrame: pandas DataFrame containing the concatenated trial data
            | **subject (str):** subject name
            | **te_id (str):** task entry id
            | **date (str):** date of stimulation
            | **stimulation_site (int):** site of stimulation
            | **%metadata_key% (ntrial):** requested metadata values for each key requested
            | **trial_time (float):** time of stimulation within recording
            | **trial_width (float):** width of stimulation pulse
            | **trial_gain (float):** fraction of maximum laser power setting
            | **trial_power (float):** power (in mW) of stimulation pulse at the fiber output

    Note:
        Only supports single-site stimulation.
    '''
    if df is None:
        df = pd.DataFrame()

    entries = list(zip(subjects, dates, ids))
    for subject, date, te in tqdm(entries): 

        # Load data from bmi3d hdf 
        try:
            exp_data, exp_metadata = load_preproc_exp_data(preproc_dir, subject, te, date)
        except:
            print(f"Entry {subject} {date} {te} could not be loaded.")
            traceback.print_exc()
            continue

        # Find laser trial times 
        try:
            trial_times, trial_widths, trial_gains, trial_powers = preproc.bmi3d.get_laser_trial_times(
                preproc_dir, subject, te, date, debug=debug, **kwargs)
        except:
            print(f"Problem extracting stimulation trials from entry {subject} {date} {te}")
            traceback.print_exc()
            continue
                        
        # Tabulate everything together
        exp = {
            'subject': subject,
            'te_id': te, 
            'date': date, 
            'trial_time': trial_times,
            'trial_width': trial_widths, 
            'trial_gain': trial_gains,
            'trial_power': trial_powers,
        }

        # Add requested metadata
        for key in metadata:
            if key == 'stimulation_site' and 'qwalor_switch_rdy_dch' in exp_metadata:
                
                # Switched laser with multiple stim sites
                exp['stimulation_site'] = preproc.bmi3d.get_switched_stimulation_sites(
                    preproc_dir, subject, te, date, trial_times, debug=debug
                )

            elif key in exp_metadata:
                exp[key] = [exp_metadata[key] for _ in range(len(trial_times))]
            else:
                exp[key] = None
                print(f"Entry {subject} {date} {te} does not have metadata {key}.")


        # Concatenate with existing dataframe
        df = pd.concat([df,pd.DataFrame(exp)], ignore_index=True)
    
    return df

def tabulate_ts_data(preproc_dir, subjects, te_ids, dates, trigger_times, time_before, time_after, 
                     channels=None, datatype='lfp'):
    '''
    Grab rectangular timeseries data from trials across arbitrary preprocessed files.
    
    Args:
        preproc_dir (str): base directory where the files live
        subjects (list of str): Subject name for each recording
        ids (list of int): Block number of Task entry object for each recording
        dates (list of str): Date for each recording
        trigger_times (list of float): times in the recording at which the desired trials start
        time_before (float): time (in seconds) to include before the trigger times
        time_after (float): time (in seconds) to include after the trigger times
        channels (list of int, optional): list of channel indices to include. Defaults to None.
        datatype (str, optional): choice of 'lfp' or 'broadband' data to load. Defaults to 'lfp'.    
        
    Returns:
        tuple: tuple containing:
            | **data (nt, nch, ntr):** tensor of data from each channel and trial
            | **samplerate (float):** sampling rate of the data
    '''

    assert len(subjects) == len(te_ids) == len(dates) == len(trigger_times)
    
    # Get the first trial
    segment_1, samplerate = get_ts_data_trial(
        preproc_dir, subjects[0], te_ids[0], dates[0], trigger_times[0], 
        time_before, time_after, channels=channels, datatype=datatype
    )
        
    # Construct the tensor using the first trial as a template
    if segment_1.ndim == 1:
        segment_1 = np.expand_dims(segment_1, 1)
    nt, nch = segment_1.shape
    segments = np.zeros((nt, nch, len(trigger_times)), like=segment_1)
    segments[:,:,0] = segment_1
    
    # Add the remaining trial
    idx = 1
    for s, t, d, tr in list(zip(subjects, te_ids, dates, trigger_times))[1:]:
        segments[:,:,idx] = get_ts_data_trial(preproc_dir, s, t, d, tr, 
                                              time_before, time_after, 
                                              channels=channels, datatype=datatype)[0]
        idx += 1
        
    return segments, samplerate

# Also, tabulate_ts_data has some errors in the docstring!
def tabulate_ts_segments(preproc_dir, subjects, te_ids, dates, start_times, end_times,
                         channels=None, datatype='lfp'):
    '''
    Grab nonrectangular timeseries data from trials across arbitrary preprocessed files.
    
    Args:
        preproc_dir (str): base directory where the files live
        subjects (list of str): Subject name for each recording
        ids (list of int): Block number of Task entry object for each recording
        dates (list of str): Date for each recording
        start_times (list of float): times in the recording at which the desired segments start
        end_times (list of float): times in the recording at which the desired segments end
        channels (list of int, optional): list of channel indices to include. Defaults to None.
        datatype (str, optional): choice of 'lfp' or 'broadband' data to load. Defaults to 'lfp'.    
        
    Returns:
        tuple: tuple containing:
            | **data (list of (nt, nch)):** list of data segments
            | **samplerate (float):** sampling rate of the data
    '''

    assert len(subjects) == len(te_ids) == len(dates) == len(start_times) == len(end_times)
    
    # Fetch the segments
    segments = []
    for s, t, d, st, et in list(zip(subjects, te_ids, dates, start_times, end_times)):
        segment, samplerate = get_ts_data_segment(preproc_dir, s, t, d, st, et, 
                                                  channels=channels, datatype=datatype)
        segments.append(segment)
        
    return segments, samplerate
        
def load_bmi3d_task_codes(filename='task_codes.yaml'):
    '''
    Load the default BMI3D task codes. File-specific codes can be found in exp_metadata['event_sync_dict']

    Args:
        filename (str, optional): filename of the task codes to load. Defaults to 'task_codes.yaml'.

    Returns:
        dict: (name, code) task code dictionary
    '''
    config_dir = files('aopy').joinpath('config')
    params_file = as_file(config_dir.joinpath(filename))
    with params_file as f:
        task_codes = yaml_read(f)[0]
    return task_codes
