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])

[((12, 27, 22, 104), 'numpy.array', 'np.array', ({(12, 36, 22, 103): '[[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]])', True, 'import numpy as np\n'), ((28, 13, 28, 29), 'distancematrix.consumer.distance_matrix.DistanceMatrix', 'DistanceMatrix', ({}, {}), '()', False, 'from distancematrix.consumer.distance_matrix import DistanceMatrix\n'), ((35, 8, 35, 62), 'numpy.testing.assert_equal', 'npt.assert_equal', ({(35, 25, 35, 43): 'dm.distance_matrix', (35, 45, 35, 61): 'self.dist_matrix'}, {}), '(dm.distance_matrix, self.dist_matrix)', True, 'import numpy.testing as npt\n'), ((38, 13, 38, 29), 'distancematrix.consumer.distance_matrix.DistanceMatrix', 'DistanceMatrix', ({}, {}), '()', False, 'from distancematrix.consumer.distance_matrix import DistanceMatrix\n'), ((41, 18, 41, 69), 'numpy.full_like', 'np.full_like', (), '', True, 'import numpy as np\n'), ((48, 8, 48, 53), 'numpy.testing.assert_equal', 'npt.assert_equal', ({(48, 25, 48, 43): 'dm.distance_matrix', (48, 45, 48, 52): 'correct'}, {}), '(dm.distance_matrix, correct)', True, 'import numpy.testing as npt\n'), ((51, 13, 51, 29), 'distancematrix.consumer.distance_matrix.DistanceMatrix', 'DistanceMatrix', ({}, {}), '()', False, 'from distancematrix.consumer.distance_matrix import DistanceMatrix\n'), ((57, 8, 57, 62), 'numpy.testing.assert_equal', 'npt.assert_equal', ({(57, 25, 57, 43): 'dm.distance_matrix', (57, 45, 57, 61): 'self.dist_matrix'}, {}), '(dm.distance_matrix, self.dist_matrix)', True, 'import numpy.testing as npt\n'), ((60, 13, 60, 29), 'distancematrix.consumer.distance_matrix.DistanceMatrix', 'DistanceMatrix', ({}, {}), '()', False, 'from distancematrix.consumer.distance_matrix import DistanceMatrix\n'), ((63, 18, 63, 69), 'numpy.full_like', 'np.full_like', (), '', True, 'import numpy as np\n'), ((69, 8, 69, 53), 'numpy.testing.assert_equal', 'npt.assert_equal', ({(69, 25, 69, 43): 'dm.distance_matrix', (69, 45, 69, 52): 'correct'}, {}), '(dm.distance_matrix, correct)', True, 'import numpy.testing as npt\n'), ((72, 13, 72, 29), 'distancematrix.consumer.distance_matrix.DistanceMatrix', 'DistanceMatrix', ({}, {}), '()', False, 'from distancematrix.consumer.distance_matrix import DistanceMatrix\n'), ((77, 19, 77, 42), 'numpy.full', 'np.full', ({(77, 27, 77, 33): '(5, 5)', (77, 35, 77, 41): 'np.nan'}, {}), '((5, 5), np.nan)', True, 'import numpy as np\n'), ((80, 8, 80, 54), 'numpy.testing.assert_equal', 'npt.assert_equal', ({(80, 25, 80, 43): 'dm.distance_matrix', (80, 45, 80, 53): 'expected'}, {}), '(dm.distance_matrix, expected)', True, 'import numpy.testing as npt\n'), ((84, 8, 84, 70), 'numpy.testing.assert_equal', 'npt.assert_equal', ({(84, 25, 84, 43): 'dm.distance_matrix', (84, 45, 84, 69): 'self.dist_matrix[:5, :5]'}, {}), '(dm.distance_matrix, self.dist_matrix[:5, :5])', True, 'import numpy.testing as npt\n'), ((89, 18, 89, 41), 'numpy.full', 'np.full', ({(89, 26, 89, 32): '(5, 5)', (89, 34, 89, 40): 'np.nan'}, {}), '((5, 5), np.nan)', True, 'import numpy as np\n'), ((91, 8, 91, 53), 'numpy.testing.assert_equal', 'npt.assert_equal', ({(91, 25, 91, 43): 'dm.distance_matrix', (91, 45, 91, 52): 'correct'}, {}), '(dm.distance_matrix, correct)', True, 'import numpy.testing as npt\n'), ((95, 8, 95, 72), 'numpy.testing.assert_equal', 'npt.assert_equal', ({(95, 25, 95, 43): 'dm.distance_matrix', (95, 45, 95, 71): 'self.dist_matrix[1:6, 3:8]'}, {}), '(dm.distance_matrix, self.dist_matrix[1:6, 3:8])', True, 'import numpy.testing as npt\n'), ((101, 18, 101, 41), 'numpy.full', 'np.full', ({(101, 26, 101, 32): '(5, 5)', (101, 34, 101, 40): 'np.nan'}, {}), '((5, 5), np.nan)', True, 'import numpy as np\n'), ((104, 8, 104, 53), 'numpy.testing.assert_equal', 'npt.assert_equal', ({(104, 25, 104, 43): 'dm.distance_matrix', (104, 45, 104, 52): 'correct'}, {}), '(dm.distance_matrix, correct)', True, 'import numpy.testing as npt\n'), ((107, 13, 107, 29), 'distancematrix.consumer.distance_matrix.DistanceMatrix', 'DistanceMatrix', ({}, {}), '()', False, 'from distancematrix.consumer.distance_matrix import DistanceMatrix\n'), ((111, 19, 111, 63), 'distancematrix.util.diag_indices_of', 'diag_indices_of', ({(111, 35, 111, 59): 'self.dist_matrix[:3, :3]', (111, 61, 111, 62): '1'}, {}), '(self.dist_matrix[:3, :3], 1)', False, 'from distancematrix.util import diag_indices_of\n'), ((113, 19, 113, 42), 'numpy.full', 'np.full', ({(113, 27, 113, 33): '(5, 5)', (113, 35, 113, 41): 'np.nan'}, {}), '((5, 5), np.nan)', True, 'import numpy as np\n'), ((117, 8, 117, 54), 'numpy.testing.assert_equal', 'npt.assert_equal', ({(117, 25, 117, 43): 'dm.distance_matrix', (117, 45, 117, 53): 'expected'}, {}), '(dm.distance_matrix, expected)', True, 'import numpy.testing as npt\n'), ((123, 8, 123, 70), 'numpy.testing.assert_equal', 'npt.assert_equal', ({(123, 25, 123, 43): 'dm.distance_matrix', (123, 45, 123, 69): 'self.dist_matrix[:5, :5]'}, {}), '(dm.distance_matrix, self.dist_matrix[:5, :5])', True, 'import numpy.testing as npt\n'), ((130, 8, 130, 54), 'numpy.testing.assert_equal', 'npt.assert_equal', ({(130, 25, 130, 43): 'dm.distance_matrix', (130, 45, 130, 53): 'expected'}, {}), '(dm.distance_matrix, expected)', True, 'import numpy.testing as npt\n'), ((135, 8, 135, 70), 'numpy.testing.assert_equal', 'npt.assert_equal', ({(135, 25, 135, 43): 'dm.distance_matrix', (135, 45, 135, 69): 'self.dist_matrix[:5, :5]'}, {}), '(dm.distance_matrix, self.dist_matrix[:5, :5])', True, 'import numpy.testing as npt\n'), ((32, 23, 32, 62), 'distancematrix.util.diag_indices_of', 'diag_indices_of', ({(32, 39, 32, 55): 'self.dist_matrix', (32, 57, 32, 61): 'diag'}, {}), '(self.dist_matrix, diag)', False, 'from distancematrix.util import diag_indices_of\n'), ((44, 23, 44, 62), 'distancematrix.util.diag_indices_of', 'diag_indices_of', ({(44, 39, 44, 55): 'self.dist_matrix', (44, 57, 44, 61): 'diag'}, {}), '(self.dist_matrix, diag)', False, 'from distancematrix.util import diag_indices_of\n'), ((75, 29, 75, 66), 'numpy.atleast_2d', 'np.atleast_2d', ({(75, 43, 75, 65): 'self.dist_matrix[0, 0]'}, {}), '(self.dist_matrix[0, 0])', True, 'import numpy as np\n'), ((76, 29, 76, 67), 'numpy.atleast_2d', 'np.atleast_2d', ({(76, 43, 76, 66): 'self.dist_matrix[:2, (1)]'}, {}), '(self.dist_matrix[:2, (1)])', True, 'import numpy as np\n'), ((99, 29, 99, 68), 'numpy.atleast_2d', 'np.atleast_2d', ({(99, 43, 99, 67): 'self.dist_matrix[3:8, (8)]'}, {}), '(self.dist_matrix[3:8, (8)])', True, 'import numpy as np\n'), ((110, 31, 110, 68), 'numpy.atleast_2d', 'np.atleast_2d', ({(110, 45, 110, 67): 'self.dist_matrix[0, 0]'}, {}), '(self.dist_matrix[0, 0])', True, 'import numpy as np\n'), ((120, 23, 120, 70), 'distancematrix.util.diag_indices_of', 'diag_indices_of', ({(120, 39, 120, 63): 'self.dist_matrix[:5, :5]', (120, 65, 120, 69): 'diag'}, {}), '(self.dist_matrix[:5, :5], diag)', False, 'from distancematrix.util import diag_indices_of\n'), ((133, 23, 133, 70), 'distancematrix.util.diag_indices_of', 'diag_indices_of', ({(133, 39, 133, 63): 'self.dist_matrix[:5, :5]', (133, 65, 133, 69): 'diag'}, {}), '(self.dist_matrix[:5, :5], diag)', False, 'from distancematrix.util import diag_indices_of\n'), ((33, 38, 33, 79), 'numpy.atleast_2d', 'np.atleast_2d', ({(33, 52, 33, 78): 'self.dist_matrix[diag_ind]'}, {}), '(self.dist_matrix[diag_ind])', True, 'import numpy as np\n'), ((45, 38, 45, 79), 'numpy.atleast_2d', 'np.atleast_2d', ({(45, 52, 45, 78): 'self.dist_matrix[diag_ind]'}, {}), '(self.dist_matrix[diag_ind])', True, 'import numpy as np\n'), ((55, 38, 55, 80), 'numpy.atleast_2d', 'np.atleast_2d', ({(55, 52, 55, 79): 'self.dist_matrix[:, (column)]'}, {}), '(self.dist_matrix[:, (column)])', True, 'import numpy as np\n'), ((66, 38, 66, 80), 'numpy.atleast_2d', 'np.atleast_2d', ({(66, 52, 66, 79): 'self.dist_matrix[:, (column)]'}, {}), '(self.dist_matrix[:, (column)])', True, 'import numpy as np\n'), ((83, 38, 83, 88), 'numpy.atleast_2d', 'np.atleast_2d', ({(83, 52, 83, 87): 'self.dist_matrix[:5, :5][:, (column)]'}, {}), '(self.dist_matrix[:5, :5][:, (column)])', True, 'import numpy as np\n'), ((94, 38, 94, 90), 'numpy.atleast_2d', 'np.atleast_2d', ({(94, 52, 94, 89): 'self.dist_matrix[1:6, 3:8][:, (column)]'}, {}), '(self.dist_matrix[1:6, 3:8][:, (column)])', True, 'import numpy as np\n'), ((112, 45, 112, 86), 'numpy.atleast_2d', 'np.atleast_2d', ({(112, 59, 112, 85): 'self.dist_matrix[diag_ind]'}, {}), '(self.dist_matrix[diag_ind])', True, 'import numpy as np\n'), ((121, 38, 121, 79), 'numpy.atleast_2d', 'np.atleast_2d', ({(121, 52, 121, 78): 'self.dist_matrix[diag_ind]'}, {}), '(self.dist_matrix[diag_ind])', True, 'import numpy as np\n'), ((134, 38, 134, 79), 'numpy.atleast_2d', 'np.atleast_2d', ({(134, 52, 134, 78): 'self.dist_matrix[diag_ind]'}, {}), '(self.dist_matrix[diag_ind])', True, 'import numpy as np\n')]