import functools
import shutil
import tempfile
import unittest
import unittest.mock
from pathlib import Path

import numpy as np
import pandas as pd

import one.alf.io as alfio
from ibllib.io.extractors import training_trials, biased_trials, camera
from ibllib.io import raw_data_loaders as raw
from ibllib.io.extractors.base import BaseExtractor


def wheelMoves_fixture(func):
    """Decorator to save some dummy wheelMoves ALF files for extraction tests"""
    @functools.wraps(func)
    def wrapper(obj=None):
        # Save some wheelMoves ALF files
        attr_list = ['training_lt5',
                     'training_ge5',
                     'biased_lt5',
                     'biased_ge5']
        alf_paths = [getattr(obj, p)['path'] / 'alf' for p in attr_list]
        n_trials = [getattr(obj, p)['ntrials'] for p in attr_list]
        for p, n in zip(alf_paths, n_trials):
            p.mkdir()
            np.save(str(p / '_ibl_wheelMoves.intervals.npy'), np.zeros((n, 2)))
            np.save(str(p / '_ibl_wheelMoves.peakAmplitude.npy'), np.zeros(n))

        # Run method
        func(obj)

        # Teardown; delete the files
        for p in alf_paths:
            shutil.rmtree(p)
    return wrapper


class TestExtractTrialData(unittest.TestCase):

    def setUp(self):
        self.main_path = Path(__file__).parent
        self.training_lt5 = {'path': self.main_path / 'data' / 'session_training_lt5'}
        self.biased_lt5 = {'path': self.main_path / 'data' / 'session_biased_lt5'}
        self.training_ge5 = {'path': self.main_path / 'data' / 'session_training_ge5'}
        self.biased_ge5 = {'path': self.main_path / 'data' / 'session_biased_ge5'}
        self.training_lt5['ntrials'] = len(raw.load_data(self.training_lt5['path']))
        self.biased_lt5['ntrials'] = len(raw.load_data(self.biased_lt5['path']))
        self.training_ge5['ntrials'] = len(raw.load_data(self.training_ge5['path']))
        self.biased_ge5['ntrials'] = len(raw.load_data(self.biased_ge5['path']))
        # turn off logging for unit testing as we will purposedly go into warning/error cases
        self.wheel_ge5_path = self.main_path / 'data' / 'wheel_ge5'
        self.wheel_lt5_path = self.main_path / 'data' / 'wheel_lt5'
        # Save some dummy wheel moves data for trial firstMovement_times extraction

    def test_get_feedbackType(self):
        # TRAINING SESSIONS
        ft = training_trials.FeedbackType(
            self.training_lt5['path']).extract()[0]
        self.assertEqual(ft.size, self.training_lt5['ntrials'])
        # check if no 0's in feedbackTypes
        self.assertFalse(ft[ft == 0].size > 0)
        # -- version >= 5.0.0
        ft = training_trials.FeedbackType(
            self.training_ge5['path']).extract()[0]
        self.assertEqual(ft.size, self.training_ge5['ntrials'])
        # check if no 0's in feedbackTypes
        self.assertFalse(ft[ft == 0].size > 0)

        # BIASED SESSIONS
        ft = biased_trials.FeedbackType(
            self.biased_lt5['path']).extract()[0]
        self.assertEqual(ft.size, self.biased_lt5['ntrials'])
        # check if no 0's in feedbackTypes
        self.assertFalse(ft[ft == 0].size > 0)
        # -- version >= 5.0.0
        ft = biased_trials.FeedbackType(
            self.biased_ge5['path']).extract()[0]
        self.assertEqual(ft.size, self.biased_ge5['ntrials'])
        # check if no 0's in feedbackTypes
        self.assertFalse(ft[ft == 0].size > 0)

    def test_get_contrastLR(self):
        # TRAINING SESSIONS
        cl, cr = training_trials.ContrastLR(
            self.training_lt5['path']).extract()[0]
        self.assertTrue(all([np.sign(x) >= 0 for x in cl if ~np.isnan(x)]))
        self.assertTrue(all([np.sign(x) >= 0 for x in cr if ~np.isnan(x)]))
        self.assertTrue(sum(np.isnan(cl)) + sum(np.isnan(cr)) == len(cl))
        self.assertTrue(sum(~np.isnan(cl)) + sum(~np.isnan(cr)) == len(cl))
        # -- version >= 5.0.0
        cl, cr = training_trials.ContrastLR(
            self.training_ge5['path']).extract()[0]
        self.assertTrue(all([np.sign(x) >= 0 for x in cl if ~np.isnan(x)]))
        self.assertTrue(all([np.sign(x) >= 0 for x in cr if ~np.isnan(x)]))
        self.assertTrue(sum(np.isnan(cl)) + sum(np.isnan(cr)) == len(cl))
        self.assertTrue(sum(~np.isnan(cl)) + sum(~np.isnan(cr)) == len(cl))

        # BIASED SESSIONS
        cl, cr = biased_trials.ContrastLR(
            self.biased_lt5['path']).extract()[0]
        self.assertTrue(all([np.sign(x) >= 0 for x in cl if ~np.isnan(x)]))
        self.assertTrue(all([np.sign(x) >= 0 for x in cr if ~np.isnan(x)]))
        self.assertTrue(sum(np.isnan(cl)) + sum(np.isnan(cr)) == len(cl))
        self.assertTrue(sum(~np.isnan(cl)) + sum(~np.isnan(cr)) == len(cl))
        # -- version >= 5.0.0
        cl, cr = biased_trials.ContrastLR(
            self.biased_ge5['path']).extract()[0]
        self.assertTrue(all([np.sign(x) >= 0 for x in cl if ~np.isnan(x)]))
        self.assertTrue(all([np.sign(x) >= 0 for x in cr if ~np.isnan(x)]))
        self.assertTrue(sum(np.isnan(cl)) + sum(np.isnan(cr)) == len(cl))
        self.assertTrue(sum(~np.isnan(cl)) + sum(~np.isnan(cr)) == len(cl))

    def test_get_probabilityLeft(self):
        # TRAINING SESSIONS
        pl = training_trials.ProbabilityLeft(
            self.training_lt5['path']).extract()[0]
        self.assertTrue(isinstance(pl, np.ndarray))
        # -- version >= 5.0.0
        pl = training_trials.ProbabilityLeft(
            self.training_ge5['path']).extract()[0]
        self.assertTrue(isinstance(pl, np.ndarray))

        # BIASED SESSIONS
        pl = biased_trials.ProbabilityLeft(
            self.biased_lt5['path']).extract()[0]
        self.assertTrue(isinstance(pl, np.ndarray))
        # Test if only probs that are in prob set
        md = raw.load_settings(self.biased_lt5['path'])
        if md:
            probs = md['BLOCK_PROBABILITY_SET']
            probs.append(0.5)
            self.assertTrue(sum([x in probs for x in pl]) == len(pl))
        # -- version >= 5.0.0
        pl = biased_trials.ProbabilityLeft(
            self.biased_ge5['path']).extract()[0]
        self.assertTrue(isinstance(pl, np.ndarray))
        # Test if only probs that are in prob set
        md = raw.load_settings(self.biased_ge5['path'])
        probs = md['BLOCK_PROBABILITY_SET']
        probs.append(0.5)
        self.assertTrue(sum([x in probs for x in pl]) == len(pl))

    def test_get_choice(self):
        # TRAINING SESSIONS
        choice = training_trials.Choice(
            session_path=self.training_lt5['path']).extract(save=False)[0]
        self.assertTrue(isinstance(choice, np.ndarray))
        data = raw.load_data(self.training_lt5['path'])
        trial_nogo = np.array(
            [~np.isnan(t['behavior_data']['States timestamps']['no_go'][0][0])
             for t in data])
        if any(trial_nogo):
            self.assertTrue(all(choice[trial_nogo]) == 0)
        # -- version >= 5.0.0
        choice = training_trials.Choice(
            session_path=self.training_ge5['path']).extract(save=False)[0]
        self.assertTrue(isinstance(choice, np.ndarray))
        data = raw.load_data(self.training_ge5['path'])
        trial_nogo = np.array(
            [~np.isnan(t['behavior_data']['States timestamps']['no_go'][0][0])
             for t in data])
        if any(trial_nogo):
            self.assertTrue(all(choice[trial_nogo]) == 0)

        # BIASED SESSIONS
        choice = biased_trials.Choice(
            session_path=self.biased_lt5['path']).extract(save=False)[0]
        self.assertTrue(isinstance(choice, np.ndarray))
        data = raw.load_data(self.biased_lt5['path'])
        trial_nogo = np.array(
            [~np.isnan(t['behavior_data']['States timestamps']['no_go'][0][0])
             for t in data])
        if any(trial_nogo):
            self.assertTrue(all(choice[trial_nogo]) == 0)
        # -- version >= 5.0.0
        choice = biased_trials.Choice(
            session_path=self.biased_ge5['path']).extract(save=False)[0]
        self.assertTrue(isinstance(choice, np.ndarray))
        data = raw.load_data(self.biased_ge5['path'])
        trial_nogo = np.array(
            [~np.isnan(t['behavior_data']['States timestamps']['no_go'][0][0])
             for t in data])
        if any(trial_nogo):
            self.assertTrue(all(choice[trial_nogo]) == 0)

    def test_get_repNum(self):
        # TODO: Test its sawtooth
        # TRAINING SESSIONS
        rn = training_trials.RepNum(
            self.training_lt5['path']).extract()[0]
        self.assertTrue(isinstance(rn, np.ndarray))
        for i in range(3):
            self.assertTrue(i in rn)
        # -- version >= 5.0.0
        rn = training_trials.RepNum(
            self.training_ge5['path']).extract()[0]
        self.assertTrue(isinstance(rn, np.ndarray))
        for i in range(4):
            self.assertTrue(i in rn)

        # BIASED SESSIONS have no repeted trials

    def test_get_rewardVolume(self):
        # TRAINING SESSIONS
        rv = training_trials.RewardVolume(
            self.training_lt5['path']).extract()[0]
        self.assertTrue(isinstance(rv, np.ndarray))
        # -- version >= 5.0.0
        rv = training_trials.RewardVolume(
            self.training_ge5['path']).extract()[0]
        self.assertTrue(isinstance(rv, np.ndarray))

        # BIASED SESSIONS
        rv = biased_trials.RewardVolume(
            self.biased_lt5['path']).extract()[0]
        self.assertTrue(isinstance(rv, np.ndarray))
        # Test if all non zero rewards are of the same value
        self.assertTrue(all([x == max(rv) for x in rv if x != 0]))
        # -- version >= 5.0.0
        rv = biased_trials.RewardVolume(
            self.biased_ge5['path']).extract()[0]
        self.assertTrue(isinstance(rv, np.ndarray))
        # Test if all non zero rewards are of the same value
        self.assertTrue(all([x == max(rv) for x in rv if x != 0]))

    def test_get_feedback_times_ge5(self):
        # TRAINING SESSIONS
        ft = training_trials.FeedbackTimes(
            self.training_ge5['path']).extract()[0]
        self.assertTrue(isinstance(ft, np.ndarray))

        # BIASED SESSIONS
        ft = biased_trials.FeedbackTimes(
            self.biased_ge5['path']).extract()[0]
        self.assertTrue(isinstance(ft, np.ndarray))

    def test_get_feedback_times_lt5(self):
        # TRAINING SESSIONS
        ft = training_trials.FeedbackTimes(
            self.training_lt5['path']).extract()[0]
        self.assertTrue(isinstance(ft, np.ndarray))

        # BIASED SESSIONS
        ft = biased_trials.FeedbackTimes(
            self.biased_lt5['path']).extract()[0]
        self.assertTrue(isinstance(ft, np.ndarray))

    def test_get_stimOnTrigger_times(self):
        # TRAINING SESSIONS
        sott = training_trials.StimOnTriggerTimes(
            self.training_lt5['path']).extract()[0]
        self.assertTrue(isinstance(sott, np.ndarray))
        # -- version >= 5.0.0
        sott = training_trials.StimOnTriggerTimes(
            self.training_ge5['path']).extract()[0]
        self.assertTrue(isinstance(sott, np.ndarray))
        # BIASED SESSIONS
        sott = biased_trials.StimOnTriggerTimes(
            self.biased_lt5['path']).extract()[0]
        self.assertTrue(isinstance(sott, np.ndarray))
        # -- version >= 5.0.0
        sott = biased_trials.StimOnTriggerTimes(
            self.biased_ge5['path']).extract()[0]
        self.assertTrue(isinstance(sott, np.ndarray))

    def test_get_stimOn_times_lt5(self):
        # TRAINING SESSIONS
        st = training_trials.StimOnTimes_deprecated(
            self.training_lt5['path']).extract()[0]
        self.assertTrue(isinstance(st, np.ndarray))

        # BIASED SESSIONS
        st = biased_trials.StimOnTimes_deprecated(
            self.biased_lt5['path']).extract()[0]
        self.assertTrue(isinstance(st, np.ndarray))

    def test_get_stimOn_times_ge5(self):
        # TRAINING SESSIONS
        st = training_trials.StimOnTimes_deprecated(
            self.training_ge5['path']).extract()[0]
        self.assertTrue(isinstance(st, np.ndarray))

        # BIASED SESSIONS
        st = biased_trials.StimOnTimes_deprecated(
            self.biased_ge5['path']).extract()[0]
        self.assertTrue(isinstance(st, np.ndarray))

    def test_stimOnOffFreeze_times(self):
        # TRAINING SESSIONS
        st = training_trials.StimOnOffFreezeTimes(
            self.training_lt5['path']).extract()[0]
        self.assertTrue(isinstance(st[0], np.ndarray))

        # BIASED SESSIONS
        st = biased_trials.StimOnOffFreezeTimes(
            self.biased_lt5['path']).extract()[0]
        self.assertTrue(isinstance(st[0], np.ndarray))

        # TRAINING SESSIONS
        st = training_trials.StimOnOffFreezeTimes(
            self.training_ge5['path']).extract()[0]
        self.assertTrue(isinstance(st[0], np.ndarray))

        # BIASED SESSIONS
        st = biased_trials.StimOnOffFreezeTimes(
            self.biased_ge5['path']).extract()[0]
        self.assertTrue(isinstance(st[0], np.ndarray))

    def test_get_intervals(self):
        # TRAINING SESSIONS
        di = training_trials.Intervals(
            self.training_lt5['path']).extract()[0]
        self.assertTrue(isinstance(di, np.ndarray))
        self.assertFalse(np.isnan(di).all())
        # -- version >= 5.0.0
        di = training_trials.Intervals(
            self.training_ge5['path']).extract()[0]
        self.assertTrue(isinstance(di, np.ndarray))
        self.assertFalse(np.isnan(di).all())

        # BIASED SESSIONS
        di = biased_trials.Intervals(
            self.training_lt5['path']).extract()[0]
        self.assertTrue(isinstance(di, np.ndarray))
        self.assertFalse(np.isnan(di).all())
        # -- version >= 5.0.0
        di = biased_trials.Intervals(
            self.training_ge5['path']).extract()[0]
        self.assertTrue(isinstance(di, np.ndarray))
        self.assertFalse(np.isnan(di).all())

    def test_get_response_times(self):
        # TRAINING SESSIONS
        rt = training_trials.ResponseTimes(
            self.training_lt5['path']).extract()[0]
        self.assertTrue(isinstance(rt, np.ndarray))
        # -- version >= 5.0.0
        rt = training_trials.ResponseTimes(
            self.training_ge5['path']).extract()[0]
        self.assertTrue(isinstance(rt, np.ndarray))

        # BIASED SESSIONS
        rt = biased_trials.ResponseTimes(
            self.biased_lt5['path']).extract()[0]
        self.assertTrue(isinstance(rt, np.ndarray))
        # -- version >= 5.0.0
        rt = biased_trials.ResponseTimes(
            self.biased_ge5['path']).extract()[0]
        self.assertTrue(isinstance(rt, np.ndarray))

    def test_get_goCueTrigger_times(self):
        # TRAINING SESSIONS
        data = raw.load_data(self.training_lt5['path'])
        gct = np.array([tr['behavior_data']['States timestamps']
                        ['closed_loop'][0][0] for tr in data])
        self.assertTrue(isinstance(gct, np.ndarray))
        # -- version >= 5.0.0
        gct = training_trials.GoCueTriggerTimes(
            self.training_ge5['path']).extract()[0]
        self.assertTrue(isinstance(gct, np.ndarray))

        # BIASED SESSIONS
        data = raw.load_data(self.biased_lt5['path'])
        gct = np.array([tr['behavior_data']['States timestamps']
                        ['closed_loop'][0][0] for tr in data])
        self.assertTrue(isinstance(gct, np.ndarray))
        # -- version >= 5.0.0
        gct = biased_trials.GoCueTriggerTimes(
            self.biased_ge5['path']).extract()[0]
        self.assertTrue(isinstance(gct, np.ndarray))

    def test_get_goCueOnset_times(self):
        # TRAINING SESSIONS
        gcot = training_trials.GoCueTimes(
            self.training_lt5['path']).extract()[0]
        self.assertTrue(isinstance(gcot, np.ndarray))
        self.assertTrue(np.all(np.isnan(gcot)))
        self.assertTrue(gcot.size != 0 or gcot.size == 4)
        # -- version >= 5.0.0
        gcot = training_trials.GoCueTimes(
            self.training_ge5['path']).extract()[0]
        self.assertTrue(isinstance(gcot, np.ndarray))
        self.assertFalse(np.any(np.isnan(gcot)))
        self.assertTrue(gcot.size != 0 or gcot.size == 12)

        # BIASED SESSIONS
        gcot = biased_trials.GoCueTimes(
            self.biased_lt5['path']).extract()[0]
        self.assertTrue(isinstance(gcot, np.ndarray))
        self.assertFalse(np.any(np.isnan(gcot)))
        self.assertTrue(gcot.size != 0 or gcot.size == 4)
        # -- version >= 5.0.0
        gcot = biased_trials.GoCueTimes(
            self.biased_ge5['path']).extract()[0]
        self.assertTrue(isinstance(gcot, np.ndarray))
        self.assertFalse(np.any(np.isnan(gcot)))
        self.assertTrue(gcot.size != 0 or gcot.size == 8)

    def test_get_included_trials_lt5(self):
        # TRAINING SESSIONS
        it = training_trials.IncludedTrials(
            self.training_lt5['path']).extract()[0]
        self.assertTrue(isinstance(it, np.ndarray))
        # BIASED SESSIONS
        it = biased_trials.IncludedTrials(
            self.biased_lt5['path']).extract()[0]
        self.assertTrue(isinstance(it, np.ndarray))

    def test_get_included_trials_ge5(self):
        # TRAINING SESSIONS
        it = training_trials.IncludedTrials(
            self.training_ge5['path']).extract()[0]
        self.assertTrue(isinstance(it, np.ndarray))
        # BIASED SESSIONS
        it = biased_trials.IncludedTrials(
            self.biased_ge5['path']).extract()[0]
        self.assertTrue(isinstance(it, np.ndarray))

    def test_get_included_trials(self):
        # TRAINING SESSIONS
        it = training_trials.IncludedTrials(
            self.training_lt5['path']).extract(settings={'IBLRIG_VERSION_TAG': '4.9.9'})[0]
        self.assertTrue(isinstance(it, np.ndarray))
        # -- version >= 5.0.0
        it = training_trials.IncludedTrials(
            self.training_ge5['path']).extract()[0]
        self.assertTrue(isinstance(it, np.ndarray))

        # BIASED SESSIONS
        it = biased_trials.IncludedTrials(
            self.biased_lt5['path']).extract(settings={'IBLRIG_VERSION_TAG': '4.9.9'})[0]
        self.assertTrue(isinstance(it, np.ndarray))
        # -- version >= 5.0.0
        it = biased_trials.IncludedTrials(
            self.biased_ge5['path']).extract()[0]
        self.assertTrue(isinstance(it, np.ndarray))

    @wheelMoves_fixture
    def test_extract_all(self):
        # TRAINING SESSIONS
        # Expect an error raised because no wheel moves were present in test data
        with self.assertRaises(ValueError) as ex:
            training_trials.extract_all(
                self.training_lt5['path'], settings={'IBLRIG_VERSION_TAG': '4.9.9'}, save=True)
            self.assertIn('_ibl_wheelMoves.intervals.npy appears to be empty', str(ex.exception))
        # -- version >= 5.0.0
        out, files = training_trials.extract_all(self.training_ge5['path'], save=True)
        self.assertEqual(19, len(out))
        self.assertTrue(all(map(Path.exists, files)))

        # BIASED SESSIONS
        # The new trials extractor additionally extracts the wheel data and this fails for the < 5.0
        # test data so we will stub the wheel extractor
        with unittest.mock.patch('ibllib.io.extractors.biased_trials.Wheel') as Wheel:
            Wheel.var_names = tuple()
            Wheel().extract.return_value = ({}, [])
            out, files = biased_trials.extract_all(
                self.biased_lt5['path'], settings={'IBLRIG_VERSION_TAG': '4.9.9'}, save=True)
            self.assertEqual(15, len(out))
            self.assertTrue(all(map(Path.exists, files)))
        # -- version >= 5.0.0
        out, files = biased_trials.extract_all(self.biased_ge5['path'], save=True)
        self.assertEqual(19, len(out))
        self.assertTrue(all(map(Path.exists, files)))

    def test_encoder_positions_clock_reset(self):
        # TRAINING SESSIONS
        # only for training?
        path = self.training_lt5['path'] / "raw_behavior_data"
        path = next(path.glob("_iblrig_encoderPositions.raw*.ssv"), None)
        dy = raw._load_encoder_positions_file_lt5(path)
        dat = np.array([849736, 1532230, 1822449, 1833514, 1841566, 1848206, 1853979, 1859144])
        self.assertTrue(np.all(np.diff(dy['re_ts']) > 0))
        self.assertTrue(all(dy['re_ts'][6:] - 2 ** 32 - dat == 0))

    def test_encoder_positions_clock_errors(self):
        # here we test for 2 kinds of file corruption that happen
        # 1/2 the first sample time is corrupt and absurdly high and should be discarded
        # 2/2 2 samples are swapped and need to be swapped backk
        path = self.biased_lt5['path'] / "raw_behavior_data"
        path = next(path.glob("_iblrig_encoderPositions.raw*.ssv"), None)
        dy = raw._load_encoder_positions_file_lt5(path)
        self.assertTrue(np.all(np.diff(np.array(dy.re_ts)) > 0))
        # -- version >= 5.0.0
        path = self.biased_ge5['path'] / "raw_behavior_data"
        path = next(path.glob("_iblrig_encoderPositions.raw*.ssv"), None)
        dy = raw._load_encoder_positions_file_ge5(path)
        self.assertTrue(np.all(np.diff(np.array(dy.re_ts)) > 0))

    def test_wheel_folders(self):
        # the wheel folder contains other errors in bpod output that had to be addressed
        for wf in self.wheel_lt5_path.glob('_iblrig_encoderPositions*.raw*.ssv'):
            df = raw._load_encoder_positions_file_lt5(wf)
            self.assertTrue(np.all(np.diff(np.array(df.re_ts)) > 0))
        for wf in self.wheel_lt5_path.glob('_iblrig_encoderEvents*.raw*.ssv'):
            df = raw._load_encoder_events_file_lt5(wf)
            self.assertTrue(np.all(np.diff(np.array(df.re_ts)) > 0))
        for wf in self.wheel_ge5_path.glob('_iblrig_encoderPositions*.raw*.ssv'):
            df = raw._load_encoder_positions_file_ge5(wf)
            self.assertTrue(np.all(np.diff(np.array(df.re_ts)) > 0))
        for wf in self.wheel_ge5_path.glob('_iblrig_encoderEvents*.raw*.ssv'):
            df = raw._load_encoder_events_file_ge5(wf)
            self.assertTrue(np.all(np.diff(np.array(df.re_ts)) > 0))

    def test_load_encoder_positions(self):
        raw.load_encoder_positions(self.training_lt5['path'],
                                   settings={'IBLRIG_VERSION_TAG': '4.9.9'})
        raw.load_encoder_positions(self.training_ge5['path'])
        raw.load_encoder_positions(self.biased_lt5['path'],
                                   settings={'IBLRIG_VERSION_TAG': '4.9.9'})
        raw.load_encoder_positions(self.biased_ge5['path'])

    def test_load_encoder_events(self):
        raw.load_encoder_events(self.training_lt5['path'],
                                settings={'IBLRIG_VERSION_TAG': '4.9.9'})
        raw.load_encoder_events(self.training_ge5['path'])
        raw.load_encoder_events(self.biased_lt5['path'],
                                settings={'IBLRIG_VERSION_TAG': '4.9.9'})
        raw.load_encoder_events(self.biased_ge5['path'])

    def test_size_outputs(self):
        # check the output dimensions
        # VERSION >= 5.0.0
        from ibllib.io.extractors.bpod_trials import extract_all
        extract_all(self.training_ge5['path'])
        trials = alfio.load_object(self.training_ge5['path'] / 'alf', object='trials')
        self.assertTrue(alfio.check_dimensions(trials) == 0)
        extract_all(self.biased_ge5['path'])
        trials = alfio.load_object(self.biased_ge5['path'] / 'alf', object='trials')
        self.assertTrue(alfio.check_dimensions(trials) == 0)
        # VERSION < 5.0.0
        # for these test data there are no wheel moves so let's mock the output
        mock_data = {
            'intervals': np.array([[0, 1], ]),
            'peakAmplitude': np.array([1, 1]),
            'peakVelocity_times': np.array([1, 1])}
        function_name = 'ibllib.io.extractors.training_wheel.extract_wheel_moves'
        # Training
        with unittest.mock.patch(function_name, return_value=mock_data):
            extract_all(self.training_lt5['path'])
        trials = alfio.load_object(self.training_lt5['path'] / 'alf', object='trials')
        self.assertTrue(alfio.check_dimensions(trials) == 0)
        # Biased
        with unittest.mock.patch(function_name, return_value=mock_data):
            extract_all(self.biased_lt5['path'])
        trials = alfio.load_object(self.biased_lt5['path'] / 'alf', object='trials')
        self.assertTrue(alfio.check_dimensions(trials) == 0)

    def tearDown(self):
        for f in self.main_path.rglob('_ibl_log.*.log'):
            f.unlink()
        [x.unlink() for x in self.training_lt5['path'].rglob('alf/*') if x.is_file()]
        [x.unlink() for x in self.biased_lt5['path'].rglob('alf/*') if x.is_file()]
        [x.unlink() for x in self.training_ge5['path'].rglob('alf/*') if x.is_file()]
        [x.unlink() for x in self.biased_ge5['path'].rglob('alf/*') if x.is_file()]
        [x.rmdir() for x in self.training_lt5['path'].rglob('alf/') if x.is_dir()]
        [x.rmdir() for x in self.biased_lt5['path'].rglob('alf/') if x.is_dir()]
        [x.rmdir() for x in self.training_ge5['path'].rglob('alf/') if x.is_dir()]
        [x.rmdir() for x in self.biased_ge5['path'].rglob('alf/') if x.is_dir()]


class TestSyncWheelBpod(unittest.TestCase):

    def test_sync_bpod_bonsai_poor_quality_timestamps(self):
        sync_trials_robust = raw.sync_trials_robust
        drift_pol = np.array([11 * 1e-6, -20])  # bpod starts 20 secs before with 10 ppm drift
        np.random.seed(seed=784)
        t0_full = np.cumsum(np.random.rand(50)) + .001
        t1_full = np.polyval(drift_pol, t0_full) + t0_full
        t0 = t0_full.copy()
        t1 = t1_full.copy()

        t0_, t1_ = sync_trials_robust(t0, t1)
        assert np.allclose(t1_, np.polyval(drift_pol, t0_) + t0_)

        t0_, t1_ = sync_trials_robust(t0, t1[:-1])
        assert np.allclose(t1_, np.polyval(drift_pol, t0_) + t0_)

        t0_, t1_ = sync_trials_robust(t0, t1[1:])
        assert np.allclose(t1_, np.polyval(drift_pol, t0_) + t0_)

        t0_, t1_ = sync_trials_robust(t0[1:], t1)
        assert np.allclose(t1_, np.polyval(drift_pol, t0_) + t0_)

        t0_, t1_ = sync_trials_robust(t0[:-1], t1)
        assert np.allclose(t1_, np.polyval(drift_pol, t0_) + t0_)

        t0_, t1_ = sync_trials_robust(t0, np.delete(t1, 24))
        assert np.allclose(t1_, np.polyval(drift_pol, t0_) + t0_)

        t0_, t1_ = sync_trials_robust(np.delete(t0, 12), np.delete(t1, 24))
        assert np.allclose(t1_, np.polyval(drift_pol, t0_) + t0_)


class TestWheelLoaders(unittest.TestCase):

    def setUp(self) -> None:
        self.main_path = Path(__file__).parent

    def test_encoder_events_corrupt(self):
        path = self.main_path.joinpath('data', 'wheel', 'lt5')
        for file_events in path.rglob('_iblrig_encoderEvents.raw.*'):
            dy = raw._load_encoder_events_file_lt5(file_events)
            self.assertTrue(dy.size > 6)
        path = self.main_path.joinpath('data', 'wheel', 'ge5')
        for file_events in path.rglob('_iblrig_encoderEvents.raw.*'):
            dy = raw._load_encoder_events_file_ge5(file_events)
            self.assertTrue(dy.size > 6)

    def test_encoder_positions_corrupts(self):
        path = self.main_path.joinpath('data', 'wheel', 'ge5')
        for file_position in path.rglob('_iblrig_encoderPositions.raw.*'):
            dy = raw._load_encoder_positions_file_ge5(file_position)
            self.assertTrue(dy.size > 18)
        path = self.main_path.joinpath('data', 'wheel', 'lt5')
        for file_position in path.rglob('_iblrig_encoderPositions.raw.*'):
            dy = raw._load_encoder_positions_file_lt5(file_position)
            self.assertTrue(dy.size > 18)


class MockExtracor(BaseExtractor):
    save_names = (
        "some_file.csv",
        "some_file.tsv",
        "some_file.ssv",
        "some_file.npy",
    )
    var_names = (
        "csv",
        "ssv",
        "tsv",
        "npy",
    )

    def _extract(self, **kwargs) -> tuple:
        csv = pd.DataFrame([1, 2, 3])
        ssv = pd.DataFrame([1, 2, 3])
        tsv = pd.DataFrame([1, 2, 3])
        npy = np.array([1, 2, 3])

        return (csv, ssv, tsv, npy)


class TestBaseExtractorSavingMethods(unittest.TestCase):
    def setUp(self) -> None:
        self.tempdir = tempfile.TemporaryDirectory()
        self.session_path = self.tempdir.name
        # self.addClassCleanup(tempdir.cleanup)  # py3.8
        self.mock_extractor = MockExtracor(self.session_path)

    def test_saving_method(self):
        data, paths = self.mock_extractor.extract(save=True)
        self.assertTrue(all([x.exists() for x in paths]))

    def tearDown(self):
        self.tempdir.cleanup()


class TestCameraExtractors(unittest.TestCase):
    def test_groom_pin_state(self):
        # UNIT DATA
        fps = 60
        t_offset = 39.4
        ts = np.arange(0, 10, 1 / fps) + t_offset
        # Add drift
        ts += np.full_like(ts, 1e-4).cumsum()
        n_pulses = 2
        pulse_width = 0.3
        duty = 0.5
        gpio = {'indices': np.empty(n_pulses * 2, dtype=np.int32),
                'polarities': np.ones(n_pulses * 2, dtype=np.int32)}
        gpio['polarities'][1::2] = -1
        aud_offset = 40.
        audio = {'times': np.empty(n_pulses * 2),
                 'polarities': gpio['polarities']}
        for p in range(n_pulses):
            i = p * 2
            rise = (pulse_width * p) + duty * p + 1
            audio['times'][i] = aud_offset + rise
            audio['times'][i + 1] = audio['times'][i] + pulse_width
            rise += t_offset
            gpio['indices'][i] = np.where(ts > rise)[0][0]
            gpio['indices'][i + 1] = np.where(ts > rise + pulse_width)[0][0]

        gpio_, audio_, ts_ = camera.groom_pin_state(gpio, audio, ts)
        self.assertEqual(audio, audio_, 'Audio dict shouldn\'t be effected')
        np.testing.assert_array_almost_equal(ts_[:4], [40., 40.016667, 40.033333, 40.05])

        # Broken TTLs + extra TTL
        delay = 0.08
        pulse_width = 1e-5
        t = audio['times'][0] + delay
        audio['times'] = np.sort(np.append(audio['times'], [t, t + pulse_width, 80]))
        audio['polarities'] = np.ones(audio['times'].shape, dtype=np.int32)
        audio['polarities'][1::2] = -1

        gpio_, audio_, _ = camera.groom_pin_state(gpio, audio, ts, min_diff=5e-3)
        self.assertTrue(audio_['times'].size == gpio_['times'].size == 4)

        # One front shifted by a large amount
        audio['times'][4] -= 0.3
        gpio_, audio_, _ = camera.groom_pin_state(gpio, audio, ts, tolerance=.1, min_diff=5e-3)
        self.assertTrue(np.all(gpio_['times'] == audio_['times']))
        self.assertTrue(np.all(gpio_['times'] == np.array([41., 41.3])))

    def test_attribute_times(self, display=False):
        # Create two timestamp arrays at two different frequencies
        tsa = np.linspace(0, 60, 60 * 4)[:60]  # 240bpm
        tsb = np.linspace(0, 60, 60 * 3)[:45]  # 180bpm
        tsa = np.sort(np.append(tsa, .4))  # Add ambiguous front
        tsb = np.sort(np.append(tsb, .41))
        if display:
            from ibllib.plots import vertical_lines
            import matplotlib.pyplot as plt
            vertical_lines(tsb, linestyle=':', color='r', label='tsb')
            vertical_lines(tsa, linestyle=':', color='b', label='tsa')
            plt.legend()

        # Check with default args
        matches = camera.attribute_times(tsa, tsb)
        expected = np.array(
            [0, 1, 2, 4, 5, 6, 8, 9, 10, 12, 13, 14, 16, 17, 18, 20, 21,
             22, 24, 25, 26, 28, 29, 30, 32, 33, 34, 36, 37, 38, 40, 41, 42, 44,
             45, 46, 48, 49, -1, 52, 53, -1, 56, 57, -1, 60]
        )
        np.testing.assert_array_equal(matches, expected)
        self.assertEqual(matches.size, tsb.size)

        # Taking closest instead of first should change index of ambiguous front
        matches = camera.attribute_times(tsa, tsb, take='nearest')
        expected[np.r_[1:3]] = expected[1:3] + 1
        np.testing.assert_array_equal(matches, expected)

        # Lower tolerance
        matches = camera.attribute_times(tsa, tsb, tol=0.05)
        expected = np.array([0, 2, 5, 9, 13, 17, 21, 25, 29, 33, 37, 41, 45, 49, 53, 57])
        np.testing.assert_array_equal(matches[matches > -1], expected)

        # Remove injective assert
        matches = camera.attribute_times(tsa, tsb, injective=False, take='nearest')
        expected = np.array(
            [0, 2, 2, 4, 5, 6, 8, 9, 10, 12, 13, 14, 16, 17, 18, 20, 21, 22,
             24, 25, 26, 28, 29, 30, 32, 33, 34, 36, 37, 38, 40, 41, 42, 44, 45,
             46, 48, 49, -1, 52, 53, -1, 56, 57, -1, 60]
        )
        np.testing.assert_array_equal(matches, expected)

        # Check input validation
        with self.assertRaises(ValueError):
            camera.attribute_times(tsa, tsb, injective=False, take='closest')


if __name__ == "__main__":
    unittest.main(exit=False, verbosity=2)
