import unittest

import scipy as sp

from connectome.model.dynamicsorn.fast_cython_sorn import stdp_pruning


class TestSparseCythonSTDP(unittest.TestCase):
    def test_nothing_modified(self):
        current_active_neurons =  sp.flatnonzero(sp.zeros(4))
        past_active_neurons = sp.flatnonzero(sp.zeros(4))
        adjacency = sp.rand(4, 4)
        adjacency_original = adjacency.copy()
        delta_row_sum = sp.zeros(4)
        e = .5
        eta_equal = e
        eta_plus = e
        eta_minus = e
        pruning_threshold = .1
        stdp_pruning(current_active_neurons, past_active_neurons, adjacency, delta_row_sum,
                     eta_equal, eta_plus, eta_minus, pruning_threshold)
        self.assertTrue((adjacency_original==adjacency).all())

    def test_equal_time(self):
        past_active_neurons = sp.flatnonzero(sp.zeros(4))
        current_active_neurons = sp.flatnonzero(sp.asarray([0,0,1,1]))
        adjacency = sp.ones((4,4)).astype(float)
        adjacency_original = adjacency.copy()
        delta_row_sum = sp.zeros(4)
        e = .5
        eta_equal = e
        eta_plus = e
        eta_minus = -e
        pruning_threshold = 10
        expected_delta_adjacency = sp.asarray([[0, 0, 0, 0],
                                          [0, 0, 0, 0],
                                          [0, 0, 0, e],
                                          [0, 0, e, 0]])
        stdp_pruning(current_active_neurons, past_active_neurons, adjacency, delta_row_sum,
                     eta_equal, eta_plus, eta_minus, pruning_threshold)
        self.assertTrue((adjacency==adjacency_original + expected_delta_adjacency).all())
        self.assertTrue((delta_row_sum==sp.asarray([0,0, e, e])).all())

    def test_one_step_delay_plus(self):
        past_active_neurons = sp.flatnonzero(sp.asarray([0,0,1,0]))
        current_active_neurons = sp.flatnonzero(sp.asarray([0,0,0,1]))
        adjacency = sp.ones((4,4)).astype(float)
        adjacency[2,3] = 0
        adjacency_original = adjacency.copy()
        delta_row_sum = sp.zeros(4)
        e = .5
        eta_equal = e
        eta_plus = e
        eta_minus = -e
        pruning_threshold = 10
        expected_delta_adjacency = sp.asarray([[0, 0, 0, 0],
                                          [0, 0, 0, 0],
                                          [0, 0, 0, 0],
                                          [0, 0, e, 0]])
        stdp_pruning(current_active_neurons, past_active_neurons, adjacency, delta_row_sum,
                     eta_equal, eta_plus, eta_minus, pruning_threshold)
        self.assertTrue((delta_row_sum==sp.asarray([0, 0, 0, e])).all())
        self.assertTrue((adjacency==adjacency_original + expected_delta_adjacency).all())

    def test_one_step_delay_minus(self):
        past_active_neurons = sp.flatnonzero(sp.asarray([0,0,1,0]))
        current_active_neurons = sp.flatnonzero(sp.asarray([0,0,0,1]))
        adjacency = sp.ones((4,4)).astype(float)
        adjacency[3,2] = 0
        adjacency_original = adjacency.copy()
        delta_row_sum = sp.zeros(4)
        e = .5
        eta_equal = e
        eta_plus = e
        eta_minus = -e
        pruning_threshold = 0
        expected_delta_adjacency = sp.asarray([[0, 0, 0, 0],
                                          [0, 0, 0, 0],
                                          [0, 0, 0, -e],
                                          [0, 0, 0, 0]])
        stdp_pruning(current_active_neurons, past_active_neurons, adjacency, delta_row_sum,
                     eta_equal, eta_plus, eta_minus, pruning_threshold)
        self.assertTrue((delta_row_sum==sp.asarray([0, 0, -e, 0])).all())
        self.assertTrue((adjacency==adjacency_original + expected_delta_adjacency).all())

    def test_combined_plus_and_current(self):
        past_active_neurons =    sp.flatnonzero(sp.asarray([0,0,1,0])) #sp.flatnonzero(sp.asarray([0,0,1,0]))
        current_active_neurons = sp.flatnonzero(sp.asarray([0,1,0,1]))
        adjacency = sp.ones((4,4)).astype(float)
        adjacency[3,2] = 0
        adjacency_original = adjacency.copy()
        delta_row_sum = sp.zeros(4)
        e = .5
        eta_equal = e
        eta_plus = e
        eta_minus = -e
        pruning_threshold = -2
        expected_delta_adjacency = sp.asarray([[0, 0,     0,    0],
                                              [0, 0,    0+e,  e+0],
                                              [0, 0+0-e,  0,0+0-e],
                                              [0, e+0, 0,   0]])
        stdp_pruning(current_active_neurons, past_active_neurons, adjacency, delta_row_sum,
                     eta_equal, eta_plus, eta_minus, pruning_threshold)
        self.assertTrue((delta_row_sum==expected_delta_adjacency.sum(axis=1)).all())
        self.assertTrue((adjacency==adjacency_original + expected_delta_adjacency).all())

    def test_combined_plus_and_current_and_inh(self):
        past_active_neurons =    sp.flatnonzero(sp.asarray([0,0,1,0])) #sp.flatnonzero(sp.asarray([0,0,1,0]))
        current_active_neurons = sp.flatnonzero(sp.asarray([0,1,0,1]))
        adjacency = sp.ones((5,5)).astype(float)
        adjacency[3,2] = 0
        adjacency_original = adjacency.copy()
        delta_row_sum = sp.zeros(4)
        e = .5
        eta_equal = e
        eta_plus = e
        eta_minus = -e
        pruning_threshold = -2
        expected_delta_adjacency = sp.asarray([[0, 0,     0,    0, 0],
                                              [0, 0,    0+e,  e+0, 0],
                                              [0, 0+0-e,  0,0+0-e, 0],
                                              [0, e+0, 0,   0, 0],
                                               [0, 0, 0, 0, 0]])
        stdp_pruning(current_active_neurons, past_active_neurons, adjacency, delta_row_sum,
                     eta_equal, eta_plus, eta_minus, pruning_threshold)
        self.assertTrue((delta_row_sum==expected_delta_adjacency[:4,:4].sum(axis=1)).all())
        self.assertTrue((adjacency==adjacency_original + expected_delta_adjacency).all())