import unittest
from dataloaders import neo
import mne
import numpy as np


class TestDataloader(unittest.TestCase):
    def test_load_sample_data(self):
        root_path = './tests/data'
        sessions = {'cylinder': ['eeg-data/2', 'eeg-data/6'], 'ball': ['eeg-data/4', 'eeg-data/5']}
        event_id = {'rest': 0, 'cylinder': 1, 'ball': 2}

        raw = neo.raw_preprocessing(root_path, sessions, unify_label=True)
        events, event_id = mne.events_from_annotations(raw, event_id=event_id)
        events, events_cnt = np.unique(events[:, -1], return_counts=True)
        self.assertTrue(np.allclose(events_cnt, (300, 150, 150)))

    def test_load_session(self):
        root_path = './tests/data'
        sessions = {'cylinder': ['eeg-data/2', 'eeg-data/6'], 'ball': ['eeg-data/4']}
        raws = neo.load_sessions(root_path, sessions)
        # test if interleaved
        sess_f = tuple(f for f, r in raws)
        self.assertEqual(len(raws), 3)
        self.assertTupleEqual(sess_f, ('cylinder', 'ball', 'cylinder'))


    def test_event_parser(self):
        # fixed length
        fs = 100
        test_event = np.array([[0, 0, 4], [100, 0, 4], [600, 0, 3], [700, 0, 3], [1000, 0, 4], [1100, 0, 4]])
        gt = np.array([[0, 400, 0], [600, 400, 2], [1000, 400, 0]])
        ret = neo.reconstruct_events(test_event, fs, 'ball', trial_duration=4)
        self.assertTrue(np.allclose(ret, gt))
        # varing length
        gt = np.array([[0, 600, 0], [600, 400, 2], [1000, 100, 0]])
        ret = neo.reconstruct_events(test_event, fs, 'ball', trial_duration=None)
        self.assertTrue(np.allclose(ret, gt))
        # change indices
        gt = np.array([[0, 400, 2], [600, 400, 0], [1000, 400, 2]])
        ret = neo.reconstruct_events(test_event, fs, 'ball', mov_trial_ind=[4], rest_trial_ind=[2, 3], trial_duration=4)
        self.assertTrue(np.allclose(ret, gt))
        # use original indices
        gt = np.array([[0, 400, 4], [600, 400, 3], [1000, 400, 4]])
        ret = neo.reconstruct_events(test_event, fs, None, trial_duration=4, use_original_label=True, mov_trial_ind=[3], rest_trial_ind=[4])
        self.assertTrue(np.allclose(ret, gt))
        # use original indices, 
        gt = np.array([[0, 400, 0], [600, 400, 1], [1000, 400, 0]])
        test_event = np.array([[0, 0, 0], [100, 0, 0], [600, 0, 1], [700, 0, 1], [1000, 0, 0], [1100, 0, 0]])
        ret = neo.reconstruct_events(test_event, fs, None, trial_duration=4, use_original_label=True, mov_trial_ind=[1], rest_trial_ind=[0])
        self.assertTrue(np.allclose(ret, gt))