from unittest import TestCase
import numpy as np
import orcasong.extractors.bundles as bundles


class TestGetPlanePositions(TestCase):
    def setUp(self) -> None:
        self.positions = np.array(
            [
                [0, 0, 0],
                [0, 0, 5],
                [5, 3, 2],
                [0, 0, 2],
            ]
        )
        self.directions = np.array(
            [
                [0, 0, 1],
                [0, 0, 1],
                [0, 0, 1],
                [0, 1, -1],
            ]
        )

    def test_positions_flat_plane(self):
        plane_point = np.array([0, 0, 0])
        plane_normal = np.array([0, 0, 1])
        result = bundles.get_plane_positions(
            self.positions,
            self.directions,
            plane_point,
            plane_normal,
        )
        target = np.array(
            [
                [0, 0],
                [0, 0],
                [5, 3],
                [0, 2],
            ]
        )
        np.testing.assert_array_equal(result, target)

    def test_positions_flat_plane_no_plane_normal(self):
        plane_point = np.array([0, 0, 0])
        with self.assertRaises(ValueError):
            bundles.get_plane_positions(
                self.positions,
                self.directions,
                plane_point,
            )

    def test_positions_shifted_plane(self):
        plane_point = np.array([1, 0, 0])
        plane_normal = np.array([0, 0, 1])
        result = bundles.get_plane_positions(
            self.positions,
            self.directions,
            plane_point,
            plane_normal,
        )
        target = np.array(
            [
                [-1, 0],
                [-1, 0],
                [4, 3],
                [-1, 2],
            ]
        )
        np.testing.assert_array_equal(result, target)

    def test_positions_angled_plane(self):
        plane_point = np.array([0, 0, 0])
        plane_normal = np.array([1, 0, 1])
        result = bundles.get_plane_positions(
            self.positions,
            self.directions,
            plane_point,
            plane_normal,
        )
        target = np.array(
            [
                [0, 0],
                [0, 0],
                [np.sqrt(50), 3],
                [0, 2],
            ]
        )
        np.testing.assert_array_equal(result, target)

    def test_positions_angled_plane_2(self):
        result = bundles.get_plane_positions(
            positions=np.array([[-1, 0, 1]]),
            directions=np.array([[1, 0, -1]]),
            plane_point=np.array([0, 0, 0]),
            plane_normal=np.array([-1, 0, 1]),
        )
        np.testing.assert_array_equal(result, np.array([[0, 0]]))


class TestPairwiseDistances(TestCase):
    def setUp(self) -> None:
        self.positions_plane = np.array(
            [
                [0, -3],
                [4, 0],
                [0, 3],
            ]
        )

    def test_matrix(self):
        result = bundles.get_pairwise_distances(
            positions_plane=self.positions_plane,
            as_matrix=True,
        )
        target = np.array(
            [
                [0, 5, 6],
                [5, 0, 5],
                [6, 5, 0],
            ]
        )
        np.testing.assert_array_equal(result, target)

    def test_flat(self):
        result = bundles.get_pairwise_distances(
            positions_plane=self.positions_plane,
        )
        target = np.array([5, 6, 5])
        np.testing.assert_array_equal(result, target)
