import numpy as np
import unittest
from conmorph.connections import RecurrentInhibition

class TestRecurrentInhibition(unittest.TestCase):
    def setUp(self):
        self.n_in = 2
        self.n_ex = 3

        self.connections = [(0,2, 0), (0,1,1) ] #(ex_1_idx, ex_2_idx, in_idx)

        in_ex_conn = np.zeros((self.n_in, self.n_ex))
        for conn in self.connections:
            ex_1, ex_2, in_idx = conn
            in_ex_conn[in_idx, ex_1] = 1
            in_ex_conn[in_idx, ex_2] = 1

        self.ex_poss = np.array([[1,0,0],[0,1,0],[0,0,1]])

        self.rg = RecurrentInhibition(in_ex_conn, self.ex_poss)

    def test_recurrent_inhibition_graph_has_the_right_number_of_vertices(self):
        self.assertTrue(len(list(self.rg.g.vertices())) == self.n_ex)

    def test_recurrent_inhibition_graph_has_the_right_number_of_edges(self):
        self.assertEqual(len(list(self.rg.g.edges())), len(self.connections))

    def test_edges_conserve_the_interneuron_identity(self):
        for conn in self.connections:
            ex_1, ex_2, in_idx = conn
            edges = self.rg.g.edge(ex_1, ex_2, all_edges=True)
            self.assertTrue(len(edges)>=1)
            self.assertTrue(in_idx in [self.rg.g.edge_properties["interneuron_idx"][edge] for edge in edges])

    def test_vertices_are_located_in_space(self):
        for ex_idx in range(self.n_ex):
            ex = self.rg.g.vertex(ex_idx)
            np.testing.assert_array_equal(self.rg.g.vertex_properties["position"][ex],self.ex_poss[ex_idx,:])

    def test_get_distance_to_neighbors(self):
        expected_distance = np.linalg.norm(self.ex_poss[:,0]-self.ex_poss[:,1])
        excitatory_index = 0
        distances = self.rg.get_distance_to_neighbors(excitatory_index)
        for actual_distance in distances:
            np.testing.assert_array_equal(expected_distance, actual_distance)

    def test_get_distances_to_neighbors(self):
        expected_distance = np.linalg.norm(self.ex_poss[:,0]-self.ex_poss[:,1])
        distances = self.rg.get_distances_to_neighbors()
        for actual_distance in distances:
            np.testing.assert_array_equal(expected_distance, actual_distance)

    def test_get_second_order_neighbors(self):
        tested_vertex = 1
        expected_second_order_neighbors = [2]
        actual_second_order_neighbors = self.rg.get_second_order_neighbors(tested_vertex)
        self.assertListEqual(expected_second_order_neighbors, actual_second_order_neighbors )

    def test_get_distance_to_second_order_neighbors(self):
        expected_distance = np.linalg.norm(self.ex_poss[:, 0] - self.ex_poss[:, 1])
        distances = self.rg.get_distances_to_second_order_neighbors()
        self.assertEqual(len(distances), 2)
        self.assertEqual(distances[0], expected_distance)