"""
Author: linxiaohong linxiaohong@neuracle.cn
Date: 2023-07-17 14:14:20
LastEditors: linxiaohong linxiaohong@neuracle.cn
LastEditTime: 2023-07-19 14:02:01
FilePath: Albatross/backend/tests/utils/core.py
Description: tests/core 中的测试共用的工具函数

Copyright (c) 2023 by Neuracle, All Rights Reserved.
"""
import mne


def get_epochs(raw, picks, event_name=None, tmin=0, tmax=1):
    events, event_id = mne.events_from_annotations(raw)
    if event_name is None:
        event_id_pick = event_id
    else:
        event_id_pick = {event_name: event_id[event_name]}
    epochs = mne.Epochs(raw,
                        events,
                        event_id_pick,
                        tmin,
                        tmax,
                        picks=picks,
                        baseline=None,
                        preload=True)
    return epochs


def crop_by_annotation(raw, annot):
    onset = annot["onset"] - raw.first_time
    if -raw.info["sfreq"] / 2 < onset < 0:
        onset = 0
    raw_crop = raw.copy().crop(onset, onset + annot["duration"])
    return raw_crop