import unittest
from pathlib import Path
from numpy.testing import assert_array_equal, assert_

from neo.rawio.tdtrawio import TdtRawIO
from neo.test.rawiotest.common_rawio_test import BaseTestRawIO


class TestTdtRawIO(BaseTestRawIO, unittest.TestCase, ):
    rawioclass = TdtRawIO
    entities_to_download = [
        'tdt'
    ]
    entities_to_test = [
        'tdt/aep_05',
        'tdt/aep_05/Block-1/aep_05_Block-1.Tdx'
    ]

    def test_invalid_dirname(self):
        invalid_name = 'random_non_existant_tdt_filename'
        assert not Path(invalid_name).exists()

        with self.assertRaises(ValueError):
            TdtRawIO(invalid_name)

    def test_compare_load_multi_single_block(self):
        dirname = self.get_local_path('tdt/aep_05')
        filename = self.get_local_path('tdt/aep_05/Block-1/aep_05_Block-1.Tdx')

        io_single = TdtRawIO(filename)
        io_multi = TdtRawIO(dirname)

        io_single.parse_header()
        io_multi.parse_header()

        self.assertEqual(io_single.tdt_block_mode, 'single')
        self.assertEqual(io_multi.tdt_block_mode, 'multi')

        self.assertEqual(io_single.block_count(), 1)
        self.assertEqual(io_multi.block_count(), 1)

        self.assertEqual(io_single.segment_count(0), 1)
        self.assertEqual(io_multi.segment_count(0), 2)

        # compare header infos
        assert_array_equal(io_single.header['signal_streams'], io_multi.header['signal_streams'])
        assert_array_equal(io_single.header['signal_channels'], io_multi.header['signal_channels'])
        assert_array_equal(io_single.header['event_channels'], io_multi.header['event_channels'])

        # not all spiking channels are present in first tdt block (segment)
        for spike_channel in io_single.header['spike_channels']:
            self.assertIn(spike_channel, io_multi.header['spike_channels'])

        # check that extracted signal chunks are identical
        assert_array_equal(io_single.get_analogsignal_chunk(0, 0, 0, 100, 0),
                           io_multi.get_analogsignal_chunk(0, 0, 0, 100, 0))


if __name__ == "__main__":
    unittest.main()
