""" WPLI 单元测试 """
import os
from core.sig_chain.sig_reader import Reader
from core.mi.eeg_wpli import Wpli

TEST_DATA_PATH = "tests/data/"
BDF_FILE_PATH = os.path.join(TEST_DATA_PATH, "eeg_raw_data.bdf")
WPLI_FILE_PATH = os.path.join(TEST_DATA_PATH, "wpli.png")


def setup_module():
    if not os.path.exists(TEST_DATA_PATH):
        os.makedirs(TEST_DATA_PATH)


def teardown_module():
    if os.path.exists(WPLI_FILE_PATH):
        os.remove(WPLI_FILE_PATH)


def test_main():
    # WPLI/CSP
    # 左右手
    ch_names = [
        "Fz", "Fp1", "F3", "F7", "C3", "T3", "T5", "P3", "O1", "Cz", "Oz", "Pz",
        "O2", "P4", "T6", "T4", "C4", "F8", "F4", "Fp2"
    ]

    reader = Reader()
    raw = reader.read(BDF_FILE_PATH, tuple(ch_names))
    raw.annotations.rename({
        "trainSuccess": "mi",
        "trainFailed": "mi",
        "restState": "rest"
    })

    wpli = Wpli()
    _, mi_epochs, _ = wpli.get_epochs(raw, tuple(ch_names))
    con_wpli, con_wpli_data = wpli.process(mi_epochs, raw.info["sfreq"])
    wpli.draw_image(con_wpli, con_wpli_data, save_path=WPLI_FILE_PATH)