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

sys.path.insert(0, realpath(join(dirname(__file__), '..')))
import numpy as np
import load_local_neo_odml_elephant

import unittest
import reachgraspio.reachgraspio as rg

rio = rg.ReachGraspIO("../../datasets/i140703-001",
                      odml_directory="../../datasets",
                      nsx_to_load='all')
block1 = rio.read_block(lazy=True)

rio = rg.ReachGraspIO("../../datasets/l101210-001",
                      odml_directory="../../datasets",
                      nsx_to_load='all')
block2 = rio.read_block(lazy=True)

class RGIOTestCase(unittest.TestCase):
    def setUp(self):
        self.blocks = [block1, block2]

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

    def test_channel_infos_present(self):
        for block in self.blocks:
            for seg in 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 block in self.blocks:
            for group in 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 block in self.blocks:
            # run this test only if neuronal signals are present
            if not [a for seg in block.segments for a in seg.analogsignals if
                    a.annotations['neural_signal']]:
                return
            for group in 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_consecutive_trial_ids(self):
        for block in self.blocks:
            for seg in block.segments:
                for ev in seg.events:
                    if 'trial_id' in ev.array_annotations:
                        # exclude invalid trial_ids:
                        trial_ids = ev.array_annotations['trial_id']
                        valid_ids = trial_ids[trial_ids != -1]
                        self.assertTrue(all(np.diff(valid_ids) >= 0))

    def test_rejection_annotations_present(self):
        for block in self.blocks:
            for seg in block.segments:
                for anasig in seg.analogsignals:
                    if anasig.annotations['neural_signal']:
                        for key in ['file_origin', 'connector_ID', 'connector_pinID',
                                    '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)

                for st in seg.spiketrains:
                    for key in ['electrode_reject_IFC', 'electrode_reject_LFC',
                                'electrode_reject_HFC','nev_dig_factor', 'nb_sorted_units',
                                'nev_hi_freq_order', 'nev_hi_freq_type', 'nev_lo_freq_order',
                                'nev_lo_freq_type']:
                        self.assertIn(key, st.annotations)

    def test_event_annotations(self):
        for block in self.blocks:
            for seg in 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):
        for block in self.blocks:
            self.assertIn('conditions', block.annotations)
            self.assertGreater(len(block.annotations['conditions']), 0)

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


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

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