import unittest
import numpy as np
import brainbox.task.passive as passive
from iblutil.numerical import ismember2d


class TestPassive(unittest.TestCase):
    def test_rf_map(self):
        """

        """
        # Simulate fake rfmap data
        test_frames = np.full((60, 15, 15), 128, dtype='uint8')
        # Test on and off individually
        test_frames[10:20, 8, 8] = 0
        test_frames[25:35, 10, 13] = 255
        # Test that interleaved are detected correctly
        test_frames[40:50, 4, 9] = 0
        test_frames[42:52, 6, 10] = 255
        test_frames[42:55, 11, 4] = 0
        test_frames[50:60, 8, 8] = 0

        test_times = np.arange(60)
        rf_map = {}
        rf_map['times'] = test_times
        rf_map['frames'] = test_frames

        rf_map_times, rf_map_pos, rf_stim_frames = passive.get_on_off_times_and_positions(rf_map)

        self.assertTrue(np.all(rf_map_times == test_times))
        self.assertEqual(rf_map_pos.shape, (15 * 15, 2))
        self.assertEqual(len(rf_stim_frames['on']), 15 * 15)
        self.assertEqual(len(rf_stim_frames['off']), 15 * 15)

        # Off is for the 0 ones
        idx = ismember2d(rf_map_pos, np.array([[8, 8]]))[0]
        self.assertTrue(np.all(rf_stim_frames['off'][idx][0][0] == [10, 50]))
        idx = ismember2d(rf_map_pos, np.array([[4, 9]]))[0]
        self.assertEqual(rf_stim_frames['off'][idx][0][0], 40)
        idx = ismember2d(rf_map_pos, np.array([[11, 4]]))[0]
        self.assertEqual(rf_stim_frames['off'][idx][0][0], 42)

        # On is for the 255 ones
        idx = ismember2d(rf_map_pos, np.array([[10, 13]]))[0]
        self.assertEqual(rf_stim_frames['on'][idx][0][0], 25)
        idx = ismember2d(rf_map_pos, np.array([[6, 10]]))[0]
        self.assertEqual(rf_stim_frames['on'][idx][0][0], 42)

        # Next test that the firing rate function works
        # Basically just make one square responsive
        spike_times = np.arange(25, 35, 0.01)
        spike_depths = 500 * np.ones_like(spike_times)

        rf_map_avg, depths = passive.get_rf_map_over_depth(rf_map_times, rf_map_pos,
                                                           rf_stim_frames, spike_times,
                                                           spike_depths, x_lim=[0, 60])
        non_zero = np.where(rf_map_avg['on'] != 0)
        self.assertEqual(np.argmin(np.abs(depths - 500)), non_zero[0][0])
        self.assertTrue(np.all(non_zero[1] == 10))
        self.assertTrue(np.all(non_zero[2] == 13))

        self.assertTrue(np.all(rf_map_avg['off'] == 0))

        rf_svd = passive.get_svd_map(rf_map_avg)
        # Make sure that the one responsive element is non-zero
        self.assertTrue(rf_svd['on'][non_zero[0][0]][non_zero[1][0], non_zero[2][0]] != 0)
        # But that all the rest are zero
        rf_svd['on'][non_zero[0][0]][non_zero[1][0], non_zero[2][0]] = 0
        self.assertTrue(np.all(np.isclose(np.vstack(rf_svd['on']), 0)))
        self.assertTrue(np.all(np.vstack(rf_svd['off']) == 0))

    def test_stim_aligned(self):

        # Make random times
        aud_stim = {}
        aud_stim['valveOn'] = np.array([10, 20, 30])
        spike_times = np.r_[np.arange(8, 9.6, 0.01), np.arange(9.6, 15, 0.002),
                            np.arange(18, 19.6, 0.005), np.arange(19.6, 25, 0.002),
                            np.arange(28, 29.6, 0.01), np.arange(29.6, 35, 0.002)]
        spike_depths = np.zeros_like(spike_times)

        stim_activity = passive.get_stim_aligned_activity(aud_stim, spike_times, spike_depths,
                                                          z_score_flag=False, x_lim=[0, 40])

        self.assertCountEqual(stim_activity.keys(), ['valveOn'])
        # The first may be a bit different due to overlap with noise floor
        self.assertTrue(np.all(stim_activity['valveOn'][0][1:] == 5))
        # make sure the rest of the depths are all zero
        self.assertTrue(np.all(stim_activity['valveOn'][1:] == 0))
