import sys
from os.path import join, realpath, dirname

sys.path.insert(0, join(dirname(realpath(__file__)), '..'))
print(sys.path)
import load_local_neo_odml_elephant

import unittest
import reachgraspio.reachgraspio as rg

class RGIOTestCase(unittest.TestCase):
    def setUp(self):
        rio = rg.ReachGraspIO("../../datasets/i140703-001", odml_directory="../../datasets",
                              nsx_to_load='all')
        self.block = rio.read_block(lazy=True)

    def test_group_units_present(self):
        unit_groups = [g for g in self.block.groups if 'Unit' in g.name]
        self.assertGreater(len(unit_groups), 0)
        self.assertGreaterEqual(len(unit_groups), len(self.block.segments[0].spiketrains))

    def test_channel_infos_present(self):
        for seg in self.block.segments:
            for anasig in seg.analogsignals:
                if anasig.annotations['neural_signal']:
                    self.assertIn('connector_aligned_ids', anasig.array_annotations)
                    self.assertIn('coordinates_x', anasig.array_annotations)
                    self.assertIn('coordinates_y', anasig.array_annotations)

    def test_group_unit_annotations(self):
        for group in self.block.groups:
            if 'Unit' in group.name:
                self.assertIn('unit_id', group.annotations)
                self.assertIn('connector_aligned_id', group.annotations)
                self.assertIn('sua', group.annotations)
                self.assertIn('mua', group.annotations)
                self.assertIn('noise', group.annotations)

                # To be investigated
                # if group.annotations['unit_id'] > 0:
                #     print(group.annotations)
                #     self.assertIn('spike_duration', group.annotations)
                #     self.assertIn('spike_amplitude', group.annotations)
                #     self.assertIn('spike_count', group.annotations)

    def test_group_unit_linking(self):
        for group in self.block.groups:
            if 'Unit' in group.name:
                asig_annotations = group.channelviews[0].obj.array_annotations
                idx = group.channelviews[0].index
                self.assertEqual(group.annotations['channel_id'],
                                 asig_annotations['channel_ids'][idx])

    def test_rejection_annotations_present(self):
        for seg in self.block.segments:
            for anasig in seg.analogsignals:
                if anasig.annotations['neural_signal']:
                    for key in ['file_origin', 'connector_ID', 'connector_pinID', 'nev_dig_factor',
                                'nb_sorted_units', 'nev_hi_freq_order', 'nev_hi_freq_type',
                                'nev_lo_freq_order', 'nev_lo_freq_type', 'nsx_hi_freq_order',
                                'nsx_lo_freq_order', 'nsx_hi_freq_type', 'nsx_lo_freq_type',
                                'description', 'nsx', 'electrode_reject_IFC',
                                'electrode_reject_LFC', 'electrode_reject_HFC']:
                        self.assertIn(key, anasig.array_annotations)

    def test_event_annotations(self):
        for seg in self.block.segments:
            for ev in seg.events:
                if ev.name in ['DigitalTrialEvents', 'AnalogTrialEvents', 'TrialEvents']:
                    for key in ['trial_event_labels', 'trial_timestamp_id', 'trial_id',
                                'belongs_to_trialtype', 'performance_in_trial',
                                'performance_in_trial_str', 'trial_reject_HFC', 'trial_reject_LFC',
                                'trial_reject_IFC']:
                        self.assertIn(key, ev.array_annotations)

            ev_names = [ev.name for ev in seg.events]
            for key in ['AnalogTrialEvents', 'DigitalTrialEvents', 'TrialEvents']:
                self.assertIn(key, ev_names)

    def test_block_annotation(self):
        self.assertIn('conditions', self.block.annotations)
        self.assertGreater(len(self.block.annotations['conditions']), 0)

    def test_connector_aligned_ids(self):
        for seg in self.block.segments:
            for anasig in seg.analogsignals:
                for coords in ['coordinates_x', 'coordinates_y']:
                    if coords in anasig.array_annotations:
                        print(np.unique(anasig.array_annotations[coords]))
                        self.assertEqual(len(np.unique(anasig.array_annotations[coords])), 10)


def suite():
    suite = unittest.makeSuite(RGIOTestCase, 'test')
    return suite

if __name__ == "__main__":
    runner = unittest.TextTestRunner(verbosity=2)
    runner.run(suite())