from brainbox.singlecell import acorr, calculate_peths
import unittest
import numpy as np


class TestPopulation(unittest.TestCase):
    def test_acorr_0(self):
        spike_times = np.array([0, 10, 10, 20])
        bin_size = 1
        winsize_bins = 2 * 3 + 1

        c_expected = np.zeros(7, dtype=np.int32)
        c_expected[3] = 1

        c = acorr(spike_times, bin_size=bin_size, window_size=winsize_bins)

        self.assertTrue(np.allclose(c, c_expected))

    def test_acorr_1(self):
        spike_times = np.array([0, 10, 10, 20], dtype=np.float64)
        bin_size = 1
        winsize_bins = 2 * 3 + 1

        c_expected = np.zeros(7, dtype=np.float64)
        c_expected[3] = 1

        c = acorr(spike_times, bin_size=bin_size, window_size=winsize_bins)

        self.assertTrue(np.allclose(c, c_expected))


class TestPeths(unittest.TestCase):
    def test_peths_synthetic(self):
        n_spikes = 20000
        n_clusters = 20
        n_events = 200
        record_length = 1654
        cluster_sel = [1, 2, 3, 6, 15, 16]
        np.random.seed(seed=42)
        spike_times = np.sort(np.random.rand(n_spikes, ) * record_length)
        spike_clusters = np.random.randint(0, n_clusters, n_spikes)
        event_times = np.sort(np.random.rand(n_events, ) * record_length)

        peth, fr = calculate_peths(spike_times, spike_clusters, cluster_ids=cluster_sel,
                                   align_times=event_times)
        self.assertTrue(peth.means.shape[0] == len(cluster_sel))
        self.assertTrue(np.all(peth.means.shape == peth.stds.shape))
        self.assertTrue(np.all(fr.shape == (n_events, len(cluster_sel), 28)))
        self.assertTrue(peth.tscale.size == 28)


def test_firing_rate():
    pass


if __name__ == "__main__":
    np.random.seed(0)
    unittest.main(exit=False)
