text: # -*- coding: utf-8 -*-
  """
  Tests of neo.io.blackrockio
  """
  
  # needed for python 3 compatibility
  from __future__ import absolute_import
  
  import time
  import warnings
  
  import unittest
  
  import numpy as np
  import quantities as pq
  
  from neo.test.iotest.common_io_test import BaseTestIO
  from neo.core import *
  
  from neo.io.neuralynxio import NeuralynxIO
  from neo.io.neuralynxio import NeuralynxIO as NewNeuralynxIO
  from neo.io.neuralynxio_v1 import NeuralynxIO as OldNeuralynxIO
  from neo import AnalogSignal
  
  
  class CommonNeuralynxIOTest(BaseTestIO, unittest.TestCase, ):
      ioclass = NeuralynxIO
      files_to_test = [
          'Cheetah_v5.5.1/original_data',
          'Cheetah_v5.6.3/original_data',
          'Cheetah_v5.7.4/original_data',
      ]
      files_to_download = [
          'Cheetah_v5.5.1/original_data/CheetahLogFile.txt',
          'Cheetah_v5.5.1/original_data/CheetahLostADRecords.txt',
          'Cheetah_v5.5.1/original_data/Events.nev',
          'Cheetah_v5.5.1/original_data/STet3a.nse',
          'Cheetah_v5.5.1/original_data/STet3b.nse',
          'Cheetah_v5.5.1/original_data/Tet3a.ncs',
          'Cheetah_v5.5.1/original_data/Tet3b.ncs',
          'Cheetah_v5.5.1/plain_data/STet3a.txt',
          'Cheetah_v5.5.1/plain_data/STet3b.txt',
          'Cheetah_v5.5.1/plain_data/Tet3a.txt',
          'Cheetah_v5.5.1/plain_data/Tet3b.txt',
          'Cheetah_v5.5.1/plain_data/Events.txt',
          'Cheetah_v5.5.1/README.txt',
          'Cheetah_v5.6.3/original_data/CheetahLogFile.txt',
          'Cheetah_v5.6.3/original_data/CheetahLostADRecords.txt',
          'Cheetah_v5.6.3/original_data/Events.nev',
          'Cheetah_v5.6.3/original_data/CSC1.ncs',
          'Cheetah_v5.6.3/original_data/CSC2.ncs',
          'Cheetah_v5.6.3/original_data/TT1.ntt',
          'Cheetah_v5.6.3/original_data/TT2.ntt',
          'Cheetah_v5.6.3/original_data/VT1.nvt',
          'Cheetah_v5.6.3/plain_data/Events.txt',
          'Cheetah_v5.6.3/plain_data/CSC1.txt',
          'Cheetah_v5.6.3/plain_data/CSC2.txt',
          'Cheetah_v5.6.3/plain_data/TT1.txt',
          'Cheetah_v5.6.3/plain_data/TT2.txt',
          'Cheetah_v5.6.3/original_data/VT1.nvt',
          'Cheetah_v5.7.4/original_data/CSC1.ncs',
          'Cheetah_v5.7.4/original_data/CSC2.ncs',
          'Cheetah_v5.7.4/original_data/CSC3.ncs',
          'Cheetah_v5.7.4/original_data/CSC4.ncs',
          'Cheetah_v5.7.4/original_data/CSC5.ncs',
          'Cheetah_v5.7.4/original_data/Events.nev',
          'Cheetah_v5.7.4/plain_data/CSC1.txt',
          'Cheetah_v5.7.4/plain_data/CSC2.txt',
          'Cheetah_v5.7.4/plain_data/CSC3.txt',
          'Cheetah_v5.7.4/plain_data/CSC4.txt',
          'Cheetah_v5.7.4/plain_data/CSC5.txt',
          'Cheetah_v5.7.4/plain_data/Events.txt',
          'Cheetah_v5.7.4/README.txt']
  
  
  class TestCheetah_v551(CommonNeuralynxIOTest, unittest.TestCase):
      cheetah_version = '5.5.1'
      files_to_test = []
  
      def test_read_block(self):
          """Read data in a certain time range into one block"""
          dirname = self.get_filename_path('Cheetah_v5.5.1/original_data')
          nio = NeuralynxIO(dirname=dirname, use_cache=False)
  
          block = nio.read_block()
  
          # Everything put in one segment
          self.assertEqual(len(block.segments), 2)
          seg = block.segments[0]
          self.assertEqual(len(seg.analogsignals), 1)
          self.assertEqual(seg.analogsignals[0].shape[-1], 2)
  
          self.assertEqual(seg.analogsignals[0].sampling_rate, 32. * pq.kHz)
          self.assertEqual(len(seg.spiketrains), 2)
  
          # Testing different parameter combinations
          block = nio.read_block(load_waveforms=True)
          self.assertEqual(len(block.segments[0].analogsignals), 1)
          self.assertEqual(len(block.segments[0].spiketrains), 2)
          self.assertEqual(block.segments[0].spiketrains[0].waveforms.shape[0],
                           block.segments[0].spiketrains[0].shape[0])
          self.assertGreater(len(block.segments[0].events), 0)
  
          self.assertEqual(len(block.channel_indexes[-1].units[0].spiketrains), 2)  # 2 segment
  
          block = nio.read_block(load_waveforms=True, units_group_mode='all-in-one')
          self.assertEqual(len(block.channel_indexes[-1].units), 2)  # 2 units
  
          block = nio.read_block(load_waveforms=True, units_group_mode='split-all')
          self.assertEqual(len(block.channel_indexes[-1].units), 1)  # 1 units by ChannelIndex
  
      def test_read_segment(self):
          dirname = self.get_filename_path('Cheetah_v5.5.1/original_data')
          nio = NeuralynxIO(dirname=dirname, use_cache=False)
  
          # read first segment entirely
          seg = nio.read_segment(seg_index=0, time_slice=None)
          self.assertEqual(len(seg.analogsignals), 1)
          self.assertEqual(seg.analogsignals[0].shape[-1], 2)
          self.assertEqual(seg.analogsignals[0].sampling_rate, 32 * pq.kHz)
          self.assertEqual(len(seg.spiketrains), 2)
  
          # Testing different parameter combinations
          seg = nio.read_segment(seg_index=0, load_waveforms=True)
          self.assertEqual(len(seg.analogsignals), 1)
          self.assertEqual(len(seg.spiketrains), 2)
          self.assertTrue(len(seg.spiketrains[0].waveforms) > 0)
          self.assertTrue(len(seg.events) > 0)
  
  
  class TestCheetah_v563(CommonNeuralynxIOTest, unittest.TestCase):
      cheetah_version = '5.6.3'
      files_to_test = []
  
      def test_read_block(self):
          """Read data in a certain time range into one block"""
          dirname = self.get_filename_path('Cheetah_v5.6.3/original_data')
          nio = NeuralynxIO(dirname=dirname, use_cache=False)
  
          block = nio.read_block()
  
          # There are two segments due to gap in recording
          self.assertEqual(len(block.segments), 2)
          for seg in block.segments:
              self.assertEqual(len(seg.analogsignals), 1)
              self.assertEqual(seg.analogsignals[0].shape[-1], 2)
              self.assertEqual(seg.analogsignals[0].sampling_rate, 2. * pq.kHz)
              self.assertEqual(len(seg.spiketrains), 8)
  
          # Testing different parameter combinations
          block = nio.read_block(load_waveforms=True)
          self.assertEqual(len(block.segments[0].analogsignals), 1)
          self.assertEqual(len(block.segments[0].spiketrains), 8)
          self.assertEqual(block.segments[0].spiketrains[0].waveforms.shape[0],
                           block.segments[0].spiketrains[0].shape[0])
          # this is tetrode data, containing 32 samples per waveform
          self.assertEqual(block.segments[0].spiketrains[0].waveforms.shape[1], 4)
          self.assertEqual(block.segments[0].spiketrains[0].waveforms.shape[-1], 32)
          self.assertGreater(len(block.segments[0].events), 0)
  
          self.assertEqual(len(block.channel_indexes[-1].units[0].spiketrains), 2)
  
          block = nio.read_block(load_waveforms=True, units_group_mode='all-in-one')
          self.assertEqual(len(block.channel_indexes[-1].units), 8)
  
          block = nio.read_block(load_waveforms=True, units_group_mode='split-all')
          self.assertEqual(len(block.channel_indexes[-1].units), 1)  # 1 units by ChannelIndex
  
      def test_read_segment(self):
          dirname = self.get_filename_path('Cheetah_v5.5.1/original_data')
          nio = NeuralynxIO(dirname=dirname, use_cache=False)
  
          # read first segment entirely
          seg = nio.read_segment(seg_index=0, time_slice=None)
          self.assertEqual(len(seg.analogsignals), 1)
          self.assertEqual(seg.analogsignals[0].shape[-1], 2)
          self.assertEqual(seg.analogsignals[0].sampling_rate, 32 * pq.kHz)
          self.assertEqual(len(seg.spiketrains), 2)
  
          # Testing different parameter combinations
          seg = nio.read_segment(seg_index=0, load_waveforms=True)
          self.assertEqual(len(seg.analogsignals), 1)
          self.assertEqual(len(seg.spiketrains), 2)
          self.assertTrue(len(seg.spiketrains[0].waveforms) > 0)
          self.assertTrue(len(seg.events) > 0)
  
  
  class TestCheetah_v574(CommonNeuralynxIOTest, unittest.TestCase):
      cheetah_version = '5.7.4'
      files_to_test = []
  
      def test_read_block(self):
          dirname = self.get_filename_path('Cheetah_v5.7.4/original_data')
          nio = NeuralynxIO(dirname=dirname, use_cache=False)
  
          block = nio.read_block()
  
          # Everything put in one segment
          seg = block.segments[0]
          self.assertEqual(len(seg.analogsignals), 1)
          self.assertEqual(seg.analogsignals[0].shape[-1], 5)
  
          self.assertEqual(seg.analogsignals[0].sampling_rate, 32 * pq.kHz)
          self.assertEqual(len(seg.spiketrains), 0)  # no nse files available
  
          # Testing different parameter combinations
          block = nio.read_block(load_waveforms=True)
          self.assertEqual(len(block.segments[0].analogsignals), 1)
          self.assertEqual(len(block.segments[0].spiketrains), 0)
          self.assertGreater(len(block.segments[0].events), 0)
  
          block = nio.read_block(signal_group_mode='split-all')
          self.assertEqual(len(block.channel_indexes), 5)
  
          block = nio.read_block(signal_group_mode='group-by-same-units')
          self.assertEqual(len(block.channel_indexes), 1)
  
  
  class TestData(CommonNeuralynxIOTest, unittest.TestCase):
      def test_ncs(self):
          for session in self.files_to_test[1:2]:  # in the long run this should include all files
              dirname = self.get_filename_path(session)
              nio = NeuralynxIO(dirname=dirname, use_cache=False)
              block = nio.read_block()
  
              for anasig_id, anasig in enumerate(block.segments[0].analogsignals):
                  chid = anasig.channel_index.channel_ids[anasig_id]
                  chname = anasig.channel_index.channel_names[anasig_id].decode('UTF-8') # need to decode, unless keyerror
                  chuid = (chname, chid)
                  filename = nio.ncs_filenames[chuid][:-3] + 'txt'
                  filename = filename.replace('original_data', 'plain_data')
                  plain_data = np.loadtxt(filename)[:, 5:].flatten()  # first columns are meta info
                  overlap = 512 * 500
                  gain_factor_0 = plain_data[0] / anasig.magnitude[0, 0]
                  np.testing.assert_allclose(plain_data[:overlap],
                                             anasig.magnitude[:overlap, 0] * gain_factor_0,
                                             rtol=0.01)
  
  
  class TestGaps(CommonNeuralynxIOTest, unittest.TestCase):
      def test_gap_handling_v551(self):
          dirname = self.get_filename_path('Cheetah_v5.5.1/original_data')
          nio = NeuralynxIO(dirname=dirname, use_cache=False)
  
          block = nio.read_block()
  
          # known gap values
          n_gaps = 1
          # so 2 segments, 2 anasigs by Channelindex, 2 SpikeTrain by Units
          self.assertEqual(len(block.segments), n_gaps + 1)
          self.assertEqual(len(block.channel_indexes[0].analogsignals), n_gaps + 1)
          self.assertEqual(len(block.channel_indexes[-1].units[0].spiketrains), n_gaps + 1)
  
      def test_gap_handling_v563(self):
          dirname = self.get_filename_path('Cheetah_v5.6.3/original_data')
          nio = NeuralynxIO(dirname=dirname, use_cache=False)
          block = nio.read_block()
  
          # known gap values
          n_gaps = 1
          # so 2 segments, 2 anasigs by Channelindex, 2 SpikeTrain by Units
          self.assertEqual(len(block.segments), n_gaps + 1)
          self.assertEqual(len(block.channel_indexes[0].analogsignals), n_gaps + 1)
          self.assertEqual(len(block.channel_indexes[-1].units[0].spiketrains), n_gaps + 1)
  
  
  def compare_old_and_new_neuralynxio():
      base = '/tmp/files_for_testing_neo/neuralynx/'
      dirname = base + 'Cheetah_v5.5.1/original_data/'
      # ~ dirname = base+'Cheetah_v5.7.4/original_data/'
  
      t0 = time.perf_counter()
      newreader = NewNeuralynxIO(dirname)
      t1 = time.perf_counter()
      bl1 = newreader.read_block(load_waveforms=True)
      t2 = time.perf_counter()
      print('newreader header', t1 - t0, 's')
      print('newreader data', t2 - t1, 's')
      print('newreader toal', t2 - t0, 's')
      for seg in bl1.segments:
          print('seg', seg.index)
          for anasig in seg.analogsignals:
              print(' AnalogSignal', anasig.name, anasig.shape, anasig.t_start)
          for st in seg.spiketrains:
              print(' SpikeTrain', st.name, st.shape, st.waveforms.shape, st[:5])
          for ev in seg.events:
              print(' Event', ev.name, ev.times.shape)
  
      print('*' * 10)
  
      t0 = time.perf_counter()
      oldreader = OldNeuralynxIO(sessiondir=dirname, use_cache='never')
      t1 = time.perf_counter()
      bl2 = oldreader.read_block(waveforms=True, events=True)
      t2 = time.perf_counter()
      print('oldreader header', t1 - t0, 's')
      print('oldreader data', t2 - t1, 's')
      print('oldreader toal', t2 - t0, 's')
      for seg in bl2.segments:
          print('seg', seg.index)
          for anasig in seg.analogsignals:
              print(' AnalogSignal', anasig.name, anasig.shape, anasig.t_start)
          for st in seg.spiketrains:
              print(' SpikeTrain', st.name, st.shape, st.waveforms.shape, st[:5])
          for ev in seg.events:
              print(' Event', ev.name, ev.times.shape)
  
      print('*' * 10)
      compare_neo_content(bl1, bl2)
  
  
  def compare_neo_content(bl1, bl2):
      print('*' * 5, 'Comparison of blocks', '*' * 5)
      object_types_to_test = [Segment, ChannelIndex, Unit, AnalogSignal,
                              SpikeTrain, Event, Epoch]
      for objtype in object_types_to_test:
          print('Testing {}'.format(objtype))
          children1 = bl1.list_children_by_class(objtype)
          children2 = bl2.list_children_by_class(objtype)
  
          if len(children1) != len(children2):
              warnings.warn('Number of {} is different in both blocks ({} != {}).'
                            ' Skipping comparison'
                            ''.format(objtype, len(children1), len(children2)))
              continue
  
          for child1, child2 in zip(children1, children2):
              compare_annotations(child1.annotations, child2.annotations)
              compare_attributes(child1, child2)
  
  
  def compare_annotations(anno1, anno2):
      if len(anno1) != len(anno2):
          warnings.warn('Different numbers of annotations! {} != {}\nSkipping further comparison of '
                        'this annotation list.'.format(anno1.keys(), anno2.keys()))
          return
      assert anno1.keys() == anno2.keys()
      for key in anno1.keys():
          anno1[key] = anno2[key]
  
  
  def compare_attributes(child1, child2):
      assert child1._all_attrs == child2._all_attrs
      for attr_id in range(len(child1._all_attrs)):
          attr_name = child1._all_attrs[attr_id][0]
          attr_dtype = child1._all_attrs[attr_id][1]
          if type(child1) == AnalogSignal and attr_name == 'signal':
              continue
          if type(child1) == SpikeTrain and attr_name == 'times':
              continue
          unequal = child1.__getattribute__(attr_name) != child2.__getattribute__(attr_name)
  
          if hasattr(unequal, 'any'):
              unequal = unequal.any()
          if unequal:
              warnings.warn('Attributes differ! {}.{}={} is not equal to {}.{}={}'
                            ''.format(child1.__class__.__name__, attr_name,
                                      child1.__getattribute__(attr_name),
                                      child2.__class__.__name__, attr_name,
                                      child2.__getattribute__(attr_name)))
  
  
  if __name__ == '__main__':
      unittest.main()
      # ~ compare_old_and_new_neuralynxio()
  