import numpy as np
from unittest import TestCase
import numpy.testing as npt

from distancematrix.util import diag_indices_of
from distancematrix.consumer.distance_matrix import DistanceMatrix


class TestContextualMatrixProfile(TestCase):

    def setUp(self):
        self.dist_matrix = np.array([
            [8.67, 1.10, 1.77, 1.26, 1.91, 4.29, 6.32, 4.24, 4.64, 5.06, 6.41, 4.07, 4.67, 9.32, 5.09],
            [4.33, 4.99, 0.14, 2.79, 2.10, 6.26, 9.40, 4.14, 5.53, 4.26, 8.21, 5.91, 6.83, 9.26, 6.19],
            [0.16, 9.05, 1.35, 4.78, 7.01, 4.36, 5.24, 8.81, 7.90, 5.84, 8.90, 7.88, 3.37, 4.70, 6.94],
            [0.94, 8.70, 3.87, 6.29, 0.32, 1.79, 5.80, 2.61, 1.43, 6.32, 1.62, 0.20, 2.28, 7.11, 2.15],
            [9.90, 4.51, 2.11, 2.83, 5.52, 8.55, 6.90, 0.24, 1.58, 4.26, 8.75, 3.71, 9.93, 8.33, 0.38],
            [7.30, 5.84, 9.63, 1.95, 3.76, 3.61, 9.42, 5.56, 5.09, 7.07, 1.90, 4.78, 1.06, 0.69, 3.67],
            [2.17, 8.37, 3.99, 4.28, 4.37, 2.86, 8.61, 3.39, 8.37, 6.95, 6.57, 1.79, 7.40, 4.41, 7.64],
            [6.26, 0.29, 6.44, 8.84, 1.24, 2.52, 6.25, 3.07, 5.55, 3.19, 8.16, 5.32, 9.01, 0.39, 9.],
            [4.67, 8.88, 3.05, 3.06, 2.36, 8.34, 4.91, 5.46, 9.25, 9.78, 0.03, 5.64, 5.10, 3.58, 6.92],
            [1.01, 0.91, 6.28, 7.79, 0.68, 5.50, 6.72, 5.11, 0.80, 9.30, 9.77, 4.71, 3.26, 7.29, 6.26]])

    def mock_initialise(self, dm):
        dm.initialise(1, self.dist_matrix.shape[0], self.dist_matrix.shape[1])

    def test_process_diagonal(self):
        dm = DistanceMatrix()
        self.mock_initialise(dm)

        for diag in range(-self.dist_matrix.shape[0] + 1, self.dist_matrix.shape[1]):
            diag_ind = diag_indices_of(self.dist_matrix, diag)
            dm.process_diagonal(diag, np.atleast_2d(self.dist_matrix[diag_ind]))

        npt.assert_equal(dm.distance_matrix, self.dist_matrix)

    def test_process_diagonal_partial_calculation(self):
        dm = DistanceMatrix()
        self.mock_initialise(dm)

        correct = np.full_like(self.dist_matrix, np.nan, dtype=float)

        for diag in range(-8, self.dist_matrix.shape[1], 3):
            diag_ind = diag_indices_of(self.dist_matrix, diag)
            dm.process_diagonal(diag, np.atleast_2d(self.dist_matrix[diag_ind]))
            correct[diag_ind] = self.dist_matrix[diag_ind]

        npt.assert_equal(dm.distance_matrix, correct)

    def test_process_column(self):
        dm = DistanceMatrix()
        self.mock_initialise(dm)

        for column in range(0, self.dist_matrix.shape[1]):
            dm.process_column(column, np.atleast_2d(self.dist_matrix[:, column]))

        npt.assert_equal(dm.distance_matrix, self.dist_matrix)

    def test_process_column_partial_calculation(self):
        dm = DistanceMatrix()
        self.mock_initialise(dm)

        correct = np.full_like(self.dist_matrix, np.nan, dtype=float)

        for column in [2, 3, 4, 5, 10, 11, 12]:
            dm.process_column(column, np.atleast_2d(self.dist_matrix[:, column]))
            correct[:, column] = self.dist_matrix[:, column]

        npt.assert_equal(dm.distance_matrix, correct)

    def test_streaming_process_column(self):
        dm = DistanceMatrix()
        dm.initialise(1, 5, 5)

        dm.process_column(0, np.atleast_2d(self.dist_matrix[0, 0]))
        dm.process_column(1, np.atleast_2d(self.dist_matrix[:2, 1]))
        expected = np.full((5, 5), np.nan)
        expected[0, 0] = self.dist_matrix[0, 0]
        expected[:2, 1] = self.dist_matrix[:2, 1]
        npt.assert_equal(dm.distance_matrix, expected)

        for column in range(0, 5):
            dm.process_column(column, np.atleast_2d(self.dist_matrix[:5, :5][:, column]))
        npt.assert_equal(dm.distance_matrix, self.dist_matrix[:5, :5])

        dm.shift_query(1)
        dm.shift_series(3)

        correct = np.full((5, 5), np.nan)
        correct[0:4, 0:2] = self.dist_matrix[1:5, 3:5]
        npt.assert_equal(dm.distance_matrix, correct)

        for column in range(0, 5):
            dm.process_column(column, np.atleast_2d(self.dist_matrix[1:6, 3:8][:, column]))
        npt.assert_equal(dm.distance_matrix, self.dist_matrix[1:6, 3:8])

        dm.shift_query(2)
        dm.shift_series(1)
        dm.process_column(4, np.atleast_2d(self.dist_matrix[3:8, 8]))

        correct = np.full((5, 5), np.nan)
        correct[0:3, 0:4] = self.dist_matrix[3:6, 4:8]
        correct[:, 4] = self.dist_matrix[3:8, 8]
        npt.assert_equal(dm.distance_matrix, correct)

    def test_streaming_process_diagonal(self):
        dm = DistanceMatrix()
        dm.initialise(1, 5, 5)

        dm.process_diagonal(0, np.atleast_2d(self.dist_matrix[0, 0]))
        diag_ind = diag_indices_of(self.dist_matrix[:3, :3], 1)
        dm.process_diagonal(1, np.atleast_2d(np.atleast_2d(self.dist_matrix[diag_ind])))
        expected = np.full((5, 5), np.nan)
        expected[0, 0] = self.dist_matrix[0, 0]
        expected[0, 1] = self.dist_matrix[0, 1]
        expected[1, 2] = self.dist_matrix[1, 2]
        npt.assert_equal(dm.distance_matrix, expected)

        for diag in range(-4,5):
            diag_ind = diag_indices_of(self.dist_matrix[:5, :5], diag)
            dm.process_diagonal(diag, np.atleast_2d(self.dist_matrix[diag_ind]))

        npt.assert_equal(dm.distance_matrix, self.dist_matrix[:5, :5])

        dm.shift_query(2)
        dm.shift_series(1)
        expected = self.dist_matrix[2:7, 1:6].copy()
        expected[-2:, :] = np.nan
        expected[:, -1:] = np.nan
        npt.assert_equal(dm.distance_matrix, expected)

        for diag in range(-4,5):
            diag_ind = diag_indices_of(self.dist_matrix[:5, :5], diag)
            dm.process_diagonal(diag, np.atleast_2d(self.dist_matrix[diag_ind]))
        npt.assert_equal(dm.distance_matrix, self.dist_matrix[:5, :5])

[((289, 1251), 'numpy.array', 'np.array', (['[[8.67, 1.1, 1.77, 1.26, 1.91, 4.29, 6.32, 4.24, 4.64, 5.06, 6.41, 4.07, \n    4.67, 9.32, 5.09], [4.33, 4.99, 0.14, 2.79, 2.1, 6.26, 9.4, 4.14, 5.53,\n    4.26, 8.21, 5.91, 6.83, 9.26, 6.19], [0.16, 9.05, 1.35, 4.78, 7.01, \n    4.36, 5.24, 8.81, 7.9, 5.84, 8.9, 7.88, 3.37, 4.7, 6.94], [0.94, 8.7, \n    3.87, 6.29, 0.32, 1.79, 5.8, 2.61, 1.43, 6.32, 1.62, 0.2, 2.28, 7.11, \n    2.15], [9.9, 4.51, 2.11, 2.83, 5.52, 8.55, 6.9, 0.24, 1.58, 4.26, 8.75,\n    3.71, 9.93, 8.33, 0.38], [7.3, 5.84, 9.63, 1.95, 3.76, 3.61, 9.42, 5.56,\n    5.09, 7.07, 1.9, 4.78, 1.06, 0.69, 3.67], [2.17, 8.37, 3.99, 4.28, 4.37,\n    2.86, 8.61, 3.39, 8.37, 6.95, 6.57, 1.79, 7.4, 4.41, 7.64], [6.26, 0.29,\n    6.44, 8.84, 1.24, 2.52, 6.25, 3.07, 5.55, 3.19, 8.16, 5.32, 9.01, 0.39,\n    9.0], [4.67, 8.88, 3.05, 3.06, 2.36, 8.34, 4.91, 5.46, 9.25, 9.78, 0.03,\n    5.64, 5.1, 3.58, 6.92], [1.01, 0.91, 6.28, 7.79, 0.68, 5.5, 6.72, 5.11,\n    0.8, 9.3, 9.77, 4.71, 3.26, 7.29, 6.26]]'], {}), '([[8.67, 1.1, 1.77, 1.26, 1.91, 4.29, 6.32, 4.24, 4.64, 5.06, 6.41,\n    4.07, 4.67, 9.32, 5.09], [4.33, 4.99, 0.14, 2.79, 2.1, 6.26, 9.4, 4.14,\n    5.53, 4.26, 8.21, 5.91, 6.83, 9.26, 6.19], [0.16, 9.05, 1.35, 4.78, \n    7.01, 4.36, 5.24, 8.81, 7.9, 5.84, 8.9, 7.88, 3.37, 4.7, 6.94], [0.94, \n    8.7, 3.87, 6.29, 0.32, 1.79, 5.8, 2.61, 1.43, 6.32, 1.62, 0.2, 2.28, \n    7.11, 2.15], [9.9, 4.51, 2.11, 2.83, 5.52, 8.55, 6.9, 0.24, 1.58, 4.26,\n    8.75, 3.71, 9.93, 8.33, 0.38], [7.3, 5.84, 9.63, 1.95, 3.76, 3.61, 9.42,\n    5.56, 5.09, 7.07, 1.9, 4.78, 1.06, 0.69, 3.67], [2.17, 8.37, 3.99, 4.28,\n    4.37, 2.86, 8.61, 3.39, 8.37, 6.95, 6.57, 1.79, 7.4, 4.41, 7.64], [6.26,\n    0.29, 6.44, 8.84, 1.24, 2.52, 6.25, 3.07, 5.55, 3.19, 8.16, 5.32, 9.01,\n    0.39, 9.0], [4.67, 8.88, 3.05, 3.06, 2.36, 8.34, 4.91, 5.46, 9.25, 9.78,\n    0.03, 5.64, 5.1, 3.58, 6.92], [1.01, 0.91, 6.28, 7.79, 0.68, 5.5, 6.72,\n    5.11, 0.8, 9.3, 9.77, 4.71, 3.26, 7.29, 6.26]])\n', (297, 1251), True, 'import numpy as np\n'), ((1505, 1521), 'distancematrix.consumer.distance_matrix.DistanceMatrix', 'DistanceMatrix', ([], {}), '()\n', (1519, 1521), False, 'from distancematrix.consumer.distance_matrix import DistanceMatrix\n'), ((1795, 1849), 'numpy.testing.assert_equal', 'npt.assert_equal', (['dm.distance_matrix', 'self.dist_matrix'], {}), '(dm.distance_matrix, self.dist_matrix)\n', (1811, 1849), True, 'import numpy.testing as npt\n'), ((1921, 1937), 'distancematrix.consumer.distance_matrix.DistanceMatrix', 'DistanceMatrix', ([], {}), '()\n', (1935, 1937), False, 'from distancematrix.consumer.distance_matrix import DistanceMatrix\n'), ((1990, 2041), 'numpy.full_like', 'np.full_like', (['self.dist_matrix', 'np.nan'], {'dtype': 'float'}), '(self.dist_matrix, np.nan, dtype=float)\n', (2002, 2041), True, 'import numpy as np\n'), ((2316, 2361), 'numpy.testing.assert_equal', 'npt.assert_equal', (['dm.distance_matrix', 'correct'], {}), '(dm.distance_matrix, correct)\n', (2332, 2361), True, 'import numpy.testing as npt\n'), ((2411, 2427), 'distancematrix.consumer.distance_matrix.DistanceMatrix', 'DistanceMatrix', ([], {}), '()\n', (2425, 2427), False, 'from distancematrix.consumer.distance_matrix import DistanceMatrix\n'), ((2612, 2666), 'numpy.testing.assert_equal', 'npt.assert_equal', (['dm.distance_matrix', 'self.dist_matrix'], {}), '(dm.distance_matrix, self.dist_matrix)\n', (2628, 2666), True, 'import numpy.testing as npt\n'), ((2736, 2752), 'distancematrix.consumer.distance_matrix.DistanceMatrix', 'DistanceMatrix', ([], {}), '()\n', (2750, 2752), False, 'from distancematrix.consumer.distance_matrix import DistanceMatrix\n'), ((2805, 2856), 'numpy.full_like', 'np.full_like', (['self.dist_matrix', 'np.nan'], {'dtype': 'float'}), '(self.dist_matrix, np.nan, dtype=float)\n', (2817, 2856), True, 'import numpy as np\n'), ((3058, 3103), 'numpy.testing.assert_equal', 'npt.assert_equal', (['dm.distance_matrix', 'correct'], {}), '(dm.distance_matrix, correct)\n', (3074, 3103), True, 'import numpy.testing as npt\n'), ((3163, 3179), 'distancematrix.consumer.distance_matrix.DistanceMatrix', 'DistanceMatrix', ([], {}), '()\n', (3177, 3179), False, 'from distancematrix.consumer.distance_matrix import DistanceMatrix\n'), ((3368, 3391), 'numpy.full', 'np.full', (['(5, 5)', 'np.nan'], {}), '((5, 5), np.nan)\n', (3375, 3391), True, 'import numpy as np\n'), ((3498, 3544), 'numpy.testing.assert_equal', 'npt.assert_equal', (['dm.distance_matrix', 'expected'], {}), '(dm.distance_matrix, expected)\n', (3514, 3544), True, 'import numpy.testing as npt\n'), ((3679, 3741), 'numpy.testing.assert_equal', 'npt.assert_equal', (['dm.distance_matrix', 'self.dist_matrix[:5, :5]'], {}), '(dm.distance_matrix, self.dist_matrix[:5, :5])\n', (3695, 3741), True, 'import numpy.testing as npt\n'), ((3815, 3838), 'numpy.full', 'np.full', (['(5, 5)', 'np.nan'], {}), '((5, 5), np.nan)\n', (3822, 3838), True, 'import numpy as np\n'), ((3902, 3947), 'numpy.testing.assert_equal', 'npt.assert_equal', (['dm.distance_matrix', 'correct'], {}), '(dm.distance_matrix, correct)\n', (3918, 3947), True, 'import numpy.testing as npt\n'), ((4084, 4148), 'numpy.testing.assert_equal', 'npt.assert_equal', (['dm.distance_matrix', 'self.dist_matrix[1:6, 3:8]'], {}), '(dm.distance_matrix, self.dist_matrix[1:6, 3:8])\n', (4100, 4148), True, 'import numpy.testing as npt\n'), ((4292, 4315), 'numpy.full', 'np.full', (['(5, 5)', 'np.nan'], {}), '((5, 5), np.nan)\n', (4299, 4315), True, 'import numpy as np\n'), ((4428, 4473), 'numpy.testing.assert_equal', 'npt.assert_equal', (['dm.distance_matrix', 'correct'], {}), '(dm.distance_matrix, correct)\n', (4444, 4473), True, 'import numpy.testing as npt\n'), ((4535, 4551), 'distancematrix.consumer.distance_matrix.DistanceMatrix', 'DistanceMatrix', ([], {}), '()\n', (4549, 4551), False, 'from distancematrix.consumer.distance_matrix import DistanceMatrix\n'), ((4673, 4717), 'distancematrix.util.diag_indices_of', 'diag_indices_of', (['self.dist_matrix[:3, :3]', '(1)'], {}), '(self.dist_matrix[:3, :3], 1)\n', (4688, 4717), False, 'from distancematrix.util import diag_indices_of\n'), ((4826, 4849), 'numpy.full', 'np.full', (['(5, 5)', 'np.nan'], {}), '((5, 5), np.nan)\n', (4833, 4849), True, 'import numpy as np\n'), ((5002, 5048), 'numpy.testing.assert_equal', 'npt.assert_equal', (['dm.distance_matrix', 'expected'], {}), '(dm.distance_matrix, expected)\n', (5018, 5048), True, 'import numpy.testing as npt\n'), ((5244, 5306), 'numpy.testing.assert_equal', 'npt.assert_equal', (['dm.distance_matrix', 'self.dist_matrix[:5, :5]'], {}), '(dm.distance_matrix, self.dist_matrix[:5, :5])\n', (5260, 5306), True, 'import numpy.testing as npt\n'), ((5490, 5536), 'numpy.testing.assert_equal', 'npt.assert_equal', (['dm.distance_matrix', 'expected'], {}), '(dm.distance_matrix, expected)\n', (5506, 5536), True, 'import numpy.testing as npt\n'), ((5731, 5793), 'numpy.testing.assert_equal', 'npt.assert_equal', (['dm.distance_matrix', 'self.dist_matrix[:5, :5]'], {}), '(dm.distance_matrix, self.dist_matrix[:5, :5])\n', (5747, 5793), True, 'import numpy.testing as npt\n'), ((1665, 1704), 'distancematrix.util.diag_indices_of', 'diag_indices_of', (['self.dist_matrix', 'diag'], {}), '(self.dist_matrix, diag)\n', (1680, 1704), False, 'from distancematrix.util import diag_indices_of\n'), ((2127, 2166), 'distancematrix.util.diag_indices_of', 'diag_indices_of', (['self.dist_matrix', 'diag'], {}), '(self.dist_matrix, diag)\n', (2142, 2166), False, 'from distancematrix.util import diag_indices_of\n'), ((3241, 3278), 'numpy.atleast_2d', 'np.atleast_2d', (['self.dist_matrix[0, 0]'], {}), '(self.dist_matrix[0, 0])\n', (3254, 3278), True, 'import numpy as np\n'), ((3309, 3349), 'numpy.atleast_2d', 'np.atleast_2d', (['self.dist_matrix[:2, (1)]'], {}), '(self.dist_matrix[:2, (1)])\n', (3322, 3349), True, 'import numpy as np\n'), ((4232, 4273), 'numpy.atleast_2d', 'np.atleast_2d', (['self.dist_matrix[3:8, (8)]'], {}), '(self.dist_matrix[3:8, (8)])\n', (4245, 4273), True, 'import numpy as np\n'), ((4615, 4652), 'numpy.atleast_2d', 'np.atleast_2d', (['self.dist_matrix[0, 0]'], {}), '(self.dist_matrix[0, 0])\n', (4628, 4652), True, 'import numpy as np\n'), ((5106, 5153), 'distancematrix.util.diag_indices_of', 'diag_indices_of', (['self.dist_matrix[:5, :5]', 'diag'], {}), '(self.dist_matrix[:5, :5], diag)\n', (5121, 5153), False, 'from distancematrix.util import diag_indices_of\n'), ((5594, 5641), 'distancematrix.util.diag_indices_of', 'diag_indices_of', (['self.dist_matrix[:5, :5]', 'diag'], {}), '(self.dist_matrix[:5, :5], diag)\n', (5609, 5641), False, 'from distancematrix.util import diag_indices_of\n'), ((1743, 1784), 'numpy.atleast_2d', 'np.atleast_2d', (['self.dist_matrix[diag_ind]'], {}), '(self.dist_matrix[diag_ind])\n', (1756, 1784), True, 'import numpy as np\n'), ((2205, 2246), 'numpy.atleast_2d', 'np.atleast_2d', (['self.dist_matrix[diag_ind]'], {}), '(self.dist_matrix[diag_ind])\n', (2218, 2246), True, 'import numpy as np\n'), ((2559, 2603), 'numpy.atleast_2d', 'np.atleast_2d', (['self.dist_matrix[:, (column)]'], {}), '(self.dist_matrix[:, (column)])\n', (2572, 2603), True, 'import numpy as np\n'), ((2944, 2988), 'numpy.atleast_2d', 'np.atleast_2d', (['self.dist_matrix[:, (column)]'], {}), '(self.dist_matrix[:, (column)])\n', (2957, 2988), True, 'import numpy as np\n'), ((3619, 3671), 'numpy.atleast_2d', 'np.atleast_2d', (['self.dist_matrix[:5, :5][:, (column)]'], {}), '(self.dist_matrix[:5, :5][:, (column)])\n', (3632, 3671), True, 'import numpy as np\n'), ((4022, 4076), 'numpy.atleast_2d', 'np.atleast_2d', (['self.dist_matrix[1:6, 3:8][:, (column)]'], {}), '(self.dist_matrix[1:6, 3:8][:, (column)])\n', (4035, 4076), True, 'import numpy as np\n'), ((4763, 4804), 'numpy.atleast_2d', 'np.atleast_2d', (['self.dist_matrix[diag_ind]'], {}), '(self.dist_matrix[diag_ind])\n', (4776, 4804), True, 'import numpy as np\n'), ((5192, 5233), 'numpy.atleast_2d', 'np.atleast_2d', (['self.dist_matrix[diag_ind]'], {}), '(self.dist_matrix[diag_ind])\n', (5205, 5233), True, 'import numpy as np\n'), ((5680, 5721), 'numpy.atleast_2d', 'np.atleast_2d', (['self.dist_matrix[diag_ind]'], {}), '(self.dist_matrix[diag_ind])\n', (5693, 5721), True, 'import numpy as np\n')]