from unittest import TestCase
import numpy as np

from interneuron_polarity.model.morphology.volume import get_second_order_distance_cube, mask_invalid_entries


class TestVolume(TestCase):

    def setUp(self):
        self.pos = np.array([[1,0,0], [2,0,0], [3,0,0]]).T
        self.number_of_positions = self.pos.shape[1]
    def test_second_order_distances(self):
        actual_second_order_distances = get_second_order_distance_cube(self.pos)

        expected_shape = tuple(self.number_of_positions for idx in range(3))
        self.assertTupleEqual(actual_second_order_distances.shape, expected_shape)
        self.assertEqual(actual_second_order_distances[0,1,2], 2)
        self.assertEqual(actual_second_order_distances[1,0,2], 3)

    def test_mask_invalid_second_order_distances(self):
        second_order_distances = get_second_order_distance_cube(self.pos)
        masked_second_order_distances = mask_invalid_entries(second_order_distances)
        expected_shape=(3,2,1)
        self.assertTupleEqual(expected_shape, masked_second_order_distances.shape)
        self.assertEqual(masked_second_order_distances[0,0,0], 2 )#distance from 0->1->2
        self.assertEqual(masked_second_order_distances[1,0,0], 3 )#distance from 0->1->2
