import unittest

import scipy as sp

from connectome.model.dynamicsorn.fast_cython_sorn import propagate_normalize_intrinsic_plasticity, copy_activity

longtype = sp.longlong


def zero_input(nr_neurons):
    return sp.zeros(nr_neurons, dtype=float)


class TestSparseCythonPropagate(unittest.TestCase):
    def test_propagate_all_zero(self):
        current_activity = sp.ones(3, dtype=longtype) * 42
        past_activity = sp.zeros(3, dtype=longtype)
        adjacency = sp.ones((3,3), dtype=float)
        row_sum = adjacency.sum(axis=1)
        delta_row_sum = sp.zeros(3, dtype=float)
        firing_threshold = sp.ones(3, dtype=float)
        target_firing_rate = .1
        eta_intrinsic_plasticity = .1
        propagate_normalize_intrinsic_plasticity(current_activity, past_activity, adjacency,
                                                 delta_row_sum, row_sum, firing_threshold,
                                                 target_firing_rate, eta_intrinsic_plasticity, zero_input, 3)
        self.assertTrue((current_activity == sp.zeros(3, dtype=longtype)).all())

    def test_propagate_inh_cancel_exc(self):
        current_activity = sp.ones(3, dtype=longtype) * 42
        past_activity = sp.asarray([1, 0, 1], dtype=longtype)
        adjacency = sp.ones((3,3), dtype=float)
        adjacency[:,2] *= -1
        row_sum = adjacency[:2,:2].sum(axis=1)
        delta_row_sum = sp.zeros(2, dtype=float)
        firing_threshold = sp.ones(2, dtype=float) * .01
        target_firing_rate = .1
        eta_intrinsic_plasticity = .1
        propagate_normalize_intrinsic_plasticity(current_activity, past_activity, adjacency,
                                                 delta_row_sum, row_sum, firing_threshold,
                                                 target_firing_rate, eta_intrinsic_plasticity, zero_input, 2)
        self.assertTrue((current_activity == sp.zeros(3)).all())

    def test_propagate_exc_to_all(self):
        current_activity = sp.ones(3, dtype=longtype) * 42
        past_activity = sp.asarray([1, 0, 0], dtype=longtype)
        adjacency = sp.ones((3,3), dtype=float)
        adjacency[:,2] *= -1
        row_sum = adjacency[:2,:2].sum(axis=1)
        delta_row_sum = sp.zeros(2, dtype=float)
        firing_threshold = sp.ones(2, dtype=float) * .01
        target_firing_rate = .1
        eta_intrinsic_plasticity = .1
        propagate_normalize_intrinsic_plasticity(current_activity, past_activity, adjacency,
                                                 delta_row_sum, row_sum, firing_threshold,
                                                 target_firing_rate, eta_intrinsic_plasticity, zero_input, 2)
        self.assertTrue((current_activity == sp.ones(3)).all())

    def test_propagate_inh_to_all(self):
        current_activity = sp.ones(3, dtype=longtype) * 42
        past_activity = sp.asarray([0, 0, 1], dtype=longtype)
        adjacency = sp.ones((3,3), dtype=float)
        adjacency[:,2] *= -1
        row_sum = adjacency[:2,:2].sum(axis=1)
        delta_row_sum = sp.zeros(2, dtype=float)
        firing_threshold = sp.ones(2, dtype=float) * .01
        target_firing_rate = .1
        eta_intrinsic_plasticity = .1
        propagate_normalize_intrinsic_plasticity(current_activity, past_activity, adjacency,
                                                 delta_row_sum, row_sum, firing_threshold,
                                                 target_firing_rate, eta_intrinsic_plasticity, zero_input, 2)
        self.assertTrue((current_activity == sp.zeros(3, dtype=longtype)).all())

    def test_one_cancel_other_not(self):
        current_activity = sp.ones(3, dtype=longtype) * 42
        past_activity = sp.asarray([0, 1, 1], dtype=longtype)
        adjacency = sp.asarray([[0,1,-1],
                                [0,1,0],
                                [0,0,0]], dtype=float)
        row_sum = adjacency[:2,:2].sum(axis=1)
        delta_row_sum = sp.zeros(2, dtype=float)
        firing_threshold = sp.ones(2, dtype=float) * .01
        target_firing_rate = .1
        eta_intrinsic_plasticity = .1
        propagate_normalize_intrinsic_plasticity(current_activity, past_activity, adjacency,
                                                 delta_row_sum, row_sum, firing_threshold,
                                                 target_firing_rate, eta_intrinsic_plasticity, zero_input, 2)
        self.assertTrue((current_activity == sp.asarray([0, 1, 0], dtype=int)).all())

    def test_propagate_all_zero_intrinsic_plasticity(self):
        current_activity = sp.zeros(3, dtype=longtype) * 42
        past_activity = sp.zeros(3, dtype=longtype)
        adjacency = sp.ones((3,3), dtype=float)
        row_sum = adjacency.sum(axis=1)
        delta_row_sum = sp.zeros(3, dtype=float)
        firing_threshold = sp.ones(3, dtype=float)
        target_firing_rate = .1
        eta_intrinsic_plasticity = .1
        propagate_normalize_intrinsic_plasticity(current_activity, past_activity, adjacency,
                                                 delta_row_sum, row_sum, firing_threshold,
                                                 target_firing_rate, eta_intrinsic_plasticity, zero_input, 3)
        self.assertTrue((firing_threshold == sp.ones(3, dtype=float) - eta_intrinsic_plasticity * target_firing_rate).all())

    def test_activate_all_neurons(self):
        current_activity = sp.zeros(3, dtype=longtype)
        past_activity = sp.ones(3, dtype=longtype)
        adjacency = sp.ones((3,3), dtype=float)
        row_sum = adjacency.sum(axis=1)
        delta_row_sum = sp.zeros(3, dtype=float)
        firing_threshold = sp.ones(3, dtype=float) * .2
        target_firing_rate = .1
        eta_intrinsic_plasticity = .1
        propagate_normalize_intrinsic_plasticity(current_activity, past_activity, adjacency,
                                                 delta_row_sum, row_sum, firing_threshold,
                                                 target_firing_rate, eta_intrinsic_plasticity, zero_input, 3)
        self.assertTrue((current_activity == sp.asarray([1, 1, 1], dtype=int)).all())

    def test_activate_all_neurons_but_no_firing_due_to_high_threshold(self):
        current_activity = sp.zeros(3, dtype=longtype)
        past_activity = sp.ones(3, dtype=longtype)
        adjacency = sp.ones((3,3), dtype=float)
        row_sum = adjacency.sum(axis=1)
        delta_row_sum = sp.zeros(3, dtype=float)
        firing_threshold = sp.ones(3, dtype=float) * 1.01
        target_firing_rate = .1
        eta_intrinsic_plasticity = .1
        propagate_normalize_intrinsic_plasticity(current_activity, past_activity, adjacency,
                                                 delta_row_sum, row_sum, firing_threshold,
                                                 target_firing_rate, eta_intrinsic_plasticity, zero_input, 3)
        self.assertTrue((current_activity == sp.asarray([0, 0, 0], dtype=int)).all())

    def test_activate_single_neuron(self):
        current_activity = sp.zeros(3, dtype=longtype)
        past_activity = sp.zeros(3, dtype=longtype)
        past_activity[0] = 1
        adjacency = sp.zeros((3,3), dtype=float)
        adjacency[1, 0] = 10
        row_sum = adjacency.sum(axis=1)
        delta_row_sum = sp.zeros(3, dtype=float)
        firing_threshold = sp.ones(3, dtype=float) * .5
        target_firing_rate = .1
        eta_intrinsic_plasticity = .1
        propagate_normalize_intrinsic_plasticity(current_activity, past_activity, adjacency,
                                                 delta_row_sum, row_sum, firing_threshold,
                                                 target_firing_rate, eta_intrinsic_plasticity, zero_input, 3)
        self.assertTrue((current_activity == sp.asarray([0, 1, 0], dtype=int)).all())

    def test_row_sum_zero_initialized(self):
        current_activity = sp.zeros(3, dtype=longtype)
        past_activity = sp.zeros(3, dtype=longtype)
        past_activity[0] = 1
        adjacency = sp.rand(3,3)
        row_sum = sp.zeros(3)
        delta_row_sum = adjacency.sum(axis=1) - row_sum
        firing_threshold = sp.ones(3, dtype=float) * .5
        target_firing_rate = .1
        eta_intrinsic_plasticity = .1
        propagate_normalize_intrinsic_plasticity(current_activity, past_activity, adjacency,
                                                 delta_row_sum, row_sum, firing_threshold,
                                                 target_firing_rate, eta_intrinsic_plasticity, zero_input, 3)
        self.assertLess(sp.absolute(row_sum - adjacency.sum(axis=1)).sum(), 1e-5)

    def test_row_sum_erroneously_initialized(self):
        current_activity = sp.zeros(3, dtype=longtype)
        past_activity = sp.zeros(3, dtype=longtype)
        past_activity[0] = 1
        adjacency = sp.rand(3,3)
        row_sum = adjacency.sum(axis=1) * 3
        delta_row_sum = adjacency.sum(axis=1) - row_sum
        firing_threshold = sp.ones(3, dtype=float) * .5
        target_firing_rate = .1
        eta_intrinsic_plasticity = .1
        propagate_normalize_intrinsic_plasticity(current_activity, past_activity, adjacency,
                                                 delta_row_sum, row_sum, firing_threshold,
                                                 target_firing_rate, eta_intrinsic_plasticity, zero_input, 3)
        self.assertLess(sp.absolute(row_sum - adjacency.sum(axis=1)).sum(), 1e-5)

    def test_row_sum_erroneously_initialized_rows_afterwards_normalized(self):
        current_activity = sp.zeros(3, dtype=longtype)
        past_activity = sp.zeros(3, dtype=longtype)
        past_activity[0] = 1
        adjacency = sp.rand(3,3)
        row_sum = adjacency.sum(axis=1) * 3
        delta_row_sum = adjacency.sum(axis=1) - row_sum
        firing_threshold = sp.ones(3, dtype=float) * .5
        target_firing_rate = .1
        eta_intrinsic_plasticity = .1
        propagate_normalize_intrinsic_plasticity(current_activity, past_activity, adjacency,
                                                 delta_row_sum, row_sum, firing_threshold,
                                                 target_firing_rate, eta_intrinsic_plasticity, zero_input, 3)
        self.assertLess(sp.absolute(sp.ones(3) - adjacency.sum(axis=1)).sum(), 1e-5)

    def test_row_sum_erroneously_initialized_delta_zero(self):
        current_activity = sp.zeros(3, dtype=longtype)
        past_activity = sp.zeros(3, dtype=longtype)
        past_activity[0] = 1
        adjacency = sp.rand(3,3)
        row_sum = adjacency.sum(axis=1) * 3
        delta_row_sum = adjacency.sum(axis=1) - row_sum
        firing_threshold = sp.ones(3, dtype=float) * .5
        target_firing_rate = .1
        eta_intrinsic_plasticity = .1
        propagate_normalize_intrinsic_plasticity(current_activity, past_activity, adjacency,
                                                 delta_row_sum, row_sum, firing_threshold,
                                                 target_firing_rate, eta_intrinsic_plasticity, zero_input, 3)
        self.assertLess(sp.absolute(delta_row_sum).sum(), 1e-5)

    def test_row_sum_correctly_initialized(self):
        current_activity = sp.zeros(3, dtype=longtype)
        past_activity = sp.zeros(3, dtype=longtype)
        past_activity[0] = 1
        adjacency = sp.rand(3,3)
        row_sum = adjacency.sum(axis=1)
        delta_row_sum = adjacency.sum(axis=1) - row_sum
        firing_threshold = sp.ones(3, dtype=float) * .5
        target_firing_rate = .1
        eta_intrinsic_plasticity = .1
        propagate_normalize_intrinsic_plasticity(current_activity, past_activity, adjacency,
                                                 delta_row_sum, row_sum, firing_threshold,
                                                 target_firing_rate, eta_intrinsic_plasticity, zero_input, 3)
        self.assertLess(sp.absolute(row_sum - adjacency.sum(axis=1)).sum(), 1e-5)

    def test_intrinsic_plasticity_all_activated(self):
        current_activity = sp.zeros(3, dtype=longtype)
        past_activity = sp.ones(3, dtype=longtype)
        adjacency = sp.ones((3,3), dtype=float)
        row_sum = adjacency.sum(axis=1)
        delta_row_sum = sp.zeros(3, dtype=float)
        firing_threshold = sp.ones(3, dtype=float) * .5
        target_firing_rate = .1
        eta_intrinsic_plasticity = .1
        propagate_normalize_intrinsic_plasticity(current_activity, past_activity, adjacency,
                                                 delta_row_sum, row_sum, firing_threshold,
                                                 target_firing_rate, eta_intrinsic_plasticity, zero_input, 3)
        self.assertTrue((firing_threshold == .5 * sp.ones(3, dtype=float) + eta_intrinsic_plasticity * (1- target_firing_rate)).all())

    def test_intrinsic_plasticity_single_neuron_active(self):
        current_activity = sp.zeros(3, dtype=longtype)
        past_activity = sp.zeros(3, dtype=longtype)
        past_activity[0] = 1
        adjacency = sp.zeros((3,3), dtype=float)
        adjacency[1, 0] = 10
        row_sum = adjacency.sum(axis=1)
        delta_row_sum = sp.zeros(3, dtype=float)
        firing_threshold = sp.ones(3, dtype=float) * .5
        target_firing_rate = .1
        eta_intrinsic_plasticity = .1
        propagate_normalize_intrinsic_plasticity(current_activity, past_activity, adjacency,
                                                 delta_row_sum, row_sum, firing_threshold,
                                                 target_firing_rate, eta_intrinsic_plasticity, zero_input, 3)
        self.assertTrue((firing_threshold == sp.asarray([.5 - eta_intrinsic_plasticity * target_firing_rate,
                                                         .5 + eta_intrinsic_plasticity * (1 - target_firing_rate),
                                                         .5 - eta_intrinsic_plasticity * target_firing_rate], dtype=float) ).all())

    def test_nr_synapses(self):
        current_activity = sp.zeros(3, dtype=longtype)
        past_activity = sp.zeros(3, dtype=longtype)
        past_activity[0] = 1
        adjacency = sp.zeros((3,3), dtype=float)
        adjacency[1, 0] = 10
        adjacency[2, 0] = 10
        row_sum = adjacency.sum(axis=1)
        delta_row_sum = sp.zeros(3, dtype=float)
        firing_threshold = sp.ones(3, dtype=float) * .5
        target_firing_rate = .1
        eta_intrinsic_plasticity = .1
        nr_synapses = propagate_normalize_intrinsic_plasticity(current_activity, past_activity, adjacency,
                                                               delta_row_sum, row_sum, firing_threshold,
                                                               target_firing_rate, eta_intrinsic_plasticity, zero_input, 3)
        self.assertEqual(2, nr_synapses)


class TestCopy(unittest.TestCase):
    def test_copy_activity(self):
        current = sp.asarray([1, 0, 1], dtype=longtype)
        current_copy = current.copy()
        past = sp.asarray([4,4,4], dtype=longtype)
        copy_activity(current, past)
        self.assertTrue((current == past).all())
        self.assertTrue((current == current_copy).all())


if __name__ == "__main__":
    unittest.main()