text: import numpy as np
  from unittest import TestCase
  import numpy.testing as npt
  
  from distancematrix.util import diag_indices_of
  from distancematrix.util import diag_indices
  from distancematrix.consumer.multidimensional_matrix_profile_lr import MultidimensionalMatrixProfileLR
  
  
  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]],
  
              [[5.89, 3.38, 2.54, 3.61, 6.17, 3.90, 3.41, 2.94, 5.02, 3.24, 0.65, 4.29, 7.93, 8.39, 0.34],
               [9.07, 6.32, 9.13, 0.22, 6.60, 9.55, 2.96, 6.94, 4.12, 4.06, 0.64, 0.55, 6.09, 0.47, 2.39],
               [6.53, 0.98, 2.08, 1.82, 3.91, 6.91, 0.90, 9.13, 4.49, 7.40, 1.67, 5.70, 3.04, 1.66, 8.26],
               [9.66, 5.53, 9.81, 8.96, 8.84, 2.23, 5.42, 6.41, 2.19, 8.69, 8.41, 9.63, 6.26, 2.77, 6.88],
               [5.68, 1.47, 3.30, 1.51, 0.97, 2.82, 6.85, 1.01, 0.86, 8.32, 0.83, 5.52, 4.61, 5.95, 9.56],
               [5.63, 5.17, 5.04, 2.77, 3.22, 8.44, 1.05, 2.06, 4.69, 0.16, 1.29, 2.02, 2.82, 1.16, 7.46],
               [4.39, 9.87, 2.19, 6.36, 6.88, 5.44, 7.18, 0.16, 6.62, 1.88, 6.46, 7.27, 7.59, 6.13, 5.34],
               [4.96, 8.46, 0.26, 7.01, 9.36, 0.61, 8.21, 1.82, 9.81, 1.28, 3.41, 3.20, 1.40, 8.20, 0.95],
               [6.36, 4.40, 7.43, 7.41, 8.58, 2.95, 5.04, 5.48, 1.33, 7.79, 1.53, 7.30, 2.02, 9.20, 3.19],
               [5.81, 5.85, 2.30, 5.09, 9.57, 8.08, 2.75, 2.99, 2.70, 5.14, 0.29, 7.95, 9.90, 9.71, 5.84]],
  
              [[7.93, 0.35, 1.02, 1.19, 6.96, 4.55, 1.85, 0.70, 6.80, 8.49, 0.54, 5.47, 7.57, 3.03, 5.93],
               [2.84, 9.46, 9.19, 4.49, 0.42, 2.26, 7.95, 4.36, 1.18, 9.94, 7.37, 4.88, 8.04, 6.01, 1.26],
               [2.73, 4.03, 2.50, 9.90, 3.51, 9.53, 1.03, 2.01, 0.71, 3.09, 1.24, 1.97, 4.78, 8.69, 4.2],
               [3.45, 4.46, 0.19, 0.54, 5.21, 9.66, 2.28, 5.84, 9.80, 9.72, 6.16, 5.17, 0.51, 8.07, 9.98],
               [5.82, 5.66, 0.02, 2.28, 8.50, 0.06, 5.98, 1.91, 3.69, 0.98, 7.10, 5.13, 6.22, 3.20, 1.05],
               [4.71, 2.63, 0.47, 6.70, 3.99, 5.36, 3.19, 3.50, 6.79, 7.85, 6.38, 2.26, 6.00, 5.31, 6.69],
               [9.67, 8.39, 7.69, 9.63, 1.80, 7.40, 0.53, 3.41, 4.18, 5.06, 6.33, 8.79, 4.18, 5.56, 4.18],
               [8.34, 8.47, 9.03, 9.66, 0.29, 6.72, 9.32, 0.56, 2.08, 6.74, 6.67, 4.30, 8.74, 8.59, 7.73],
               [3.83, 5.72, 4.78, 9.77, 8.38, 6.96, 0.90, 5.01, 9.55, 2.26, 3.10, 3.78, 2.83, 3.21, 1.7],
               [1.41, 7.79, 8.91, 0.13, 4.94, 9.35, 8.13, 4.78, 0.15, 4.76, 2.10, 7.92, 8.54, 2.19, 2.71]]
          ])
  
      def mock_initialise(self, mmp):
          mmp.initialise(*self.dist_matrix.shape)
  
      def bruteforce_mmp(self, dist_matrix):
          """
          Brute force calculation of multidimensional distance matrix.
          :param dist_matrix: 2D matrix
          :return: a tuple (profile, index, dimensions)
          """
  
          correct = np.full((dist_matrix.shape[0], dist_matrix.shape[2]), np.inf, dtype=float)
          correct_index = np.full((dist_matrix.shape[0], dist_matrix.shape[2]), -1, dtype=int)
          correct_dims = [np.full((i + 1, dist_matrix.shape[2]), -1, dtype=int) for i in range(dist_matrix.shape[0])]
  
          correct[0], correct_index[0], correct_dims[0] = self._bruteforce_mmp_for_dims(dist_matrix, ([0], [1], [2]))
          correct[1], correct_index[1], correct_dims[1] = self._bruteforce_mmp_for_dims(dist_matrix, ([0, 1], [0, 2], [1, 2]))
          correct[2], correct_index[2], correct_dims[2] = self._bruteforce_mmp_for_dims(dist_matrix, ([0, 1, 2],))
  
          return correct, correct_index, correct_dims
  
      def _bruteforce_mmp_for_dims(self, dist_matrix, dimensions):
          mp = np.full(dist_matrix.shape[2], np.inf, dtype=float)
          index = np.full(dist_matrix.shape[2], -1, dtype=int)
          dims = np.full((len(dimensions[0]), dist_matrix.shape[2]), -1, dtype=int)
  
          for selected_dims in dimensions:
              dim_view = np.sum(dist_matrix[selected_dims], axis=0) / len(selected_dims)
              dim_mp = np.min(dim_view, axis=0)
              dim_mp_better = dim_mp < mp
              mp[dim_mp_better] = dim_mp[dim_mp_better]
              index[dim_mp_better] = np.argmin(dim_view, axis=0)[dim_mp_better]
              dims[:, dim_mp_better] = np.asarray(selected_dims)[:, None]
  
          return mp, index, dims
  
      def test_process_diagonal(self):
          mmp = MultidimensionalMatrixProfileLR()
          self.mock_initialise(mmp)
  
          for diag in range(-self.dist_matrix.shape[1] + 1, self.dist_matrix.shape[2]):
              diag_ind = diag_indices_of(self.dist_matrix[0], diag)
              mmp.process_diagonal(diag, self.dist_matrix[:, diag_ind[0], diag_ind[1]])
  
          correct_mp, correct_index, correct_dims = self.bruteforce_mmp(self.dist_matrix)
  
          npt.assert_allclose(mmp.md_matrix_profile(), correct_mp)
          npt.assert_equal(mmp.md_profile_index(), correct_index)
          for i, dims in enumerate(correct_dims):
              sorted_mmp_dims = np.sort(mmp.md_profile_dimensions()[i], axis=0)
              npt.assert_equal(sorted_mmp_dims, dims)
  
      def test_process_diagonal_partial_calculation(self):
          part_dist_matrix = np.full_like(self.dist_matrix, np.inf)
  
          mmp = MultidimensionalMatrixProfileLR()
          self.mock_initialise(mmp)
  
          for diag in range(-8, self.dist_matrix.shape[1], 4):
              diag_ind = diag_indices_of(self.dist_matrix[0], diag)
              mmp.process_diagonal(diag, self.dist_matrix[:, diag_ind[0], diag_ind[1]])
              part_dist_matrix[:, diag_ind[0], diag_ind[1]] = self.dist_matrix[:, diag_ind[0], diag_ind[1]]
  
          correct_mp, correct_index, correct_dims = self.bruteforce_mmp(part_dist_matrix)
  
          npt.assert_allclose(mmp.md_matrix_profile(), correct_mp)
          npt.assert_equal(mmp.md_profile_index(), correct_index)
          for i, dims in enumerate(correct_dims):
              sorted_mmp_dims = np.sort(mmp.md_profile_dimensions()[i], axis=0)
              npt.assert_equal(sorted_mmp_dims, dims)
  
      def test_process_diagonal_lr(self):
          l_dist_matrix = self.dist_matrix.copy()
          r_dist_matrix = self.dist_matrix.copy()
  
          mmp = MultidimensionalMatrixProfileLR()
          self.mock_initialise(mmp)
  
          for diag in range(-self.dist_matrix.shape[1] + 1, self.dist_matrix.shape[2]):
              diag_ind = diag_indices_of(self.dist_matrix[0], diag)
              if diag < 0:
                  l_dist_matrix[:, diag_ind[0], diag_ind[1]] = np.inf
              else:
                  r_dist_matrix[:, diag_ind[0], diag_ind[1]] = np.inf
              mmp.process_diagonal(diag, self.dist_matrix[:, diag_ind[0], diag_ind[1]])
  
          correct_lmp, correct_lindex, correct_ldims = self.bruteforce_mmp(l_dist_matrix)
          correct_rmp, correct_rindex, correct_rdims = self.bruteforce_mmp(r_dist_matrix)
  
          npt.assert_allclose(mmp.md_matrix_profile_left, correct_lmp)
          npt.assert_allclose(mmp.md_matrix_profile_right, correct_rmp)
          npt.assert_equal(mmp.md_profile_index_left, correct_lindex)
          npt.assert_equal(mmp.md_profile_index_right, correct_rindex)
          for i, dims in enumerate(correct_ldims):
              sorted_mmp_ldims = np.sort(mmp.md_profile_dimension_left[i], axis=0)
              sorted_mmp_rdims = np.sort(mmp.md_profile_dimension_right[i], axis=0)
              npt.assert_equal(sorted_mmp_ldims, correct_ldims[i])
              npt.assert_equal(sorted_mmp_rdims, correct_rdims[i])
  
      def test_process_column(self):
          mmp = MultidimensionalMatrixProfileLR()
          self.mock_initialise(mmp)
  
          for column in range(0, self.dist_matrix.shape[2]):
              mmp.process_column(column, self.dist_matrix[:, :, column])
  
          correct_mp, correct_index, correct_dims = self.bruteforce_mmp(self.dist_matrix)
  
          npt.assert_allclose(mmp.md_matrix_profile(), correct_mp)
          npt.assert_equal(mmp.md_profile_index(), correct_index)
          for i, dims in enumerate(correct_dims):
              sorted_mmp_dims = np.sort(mmp.md_profile_dimensions()[i], axis=0)
              npt.assert_equal(sorted_mmp_dims, dims)
  
      def test_process_column_partial_calculation(self):
          part_dist_matrix = np.full_like(self.dist_matrix, np.inf)
  
          mmp = MultidimensionalMatrixProfileLR()
          self.mock_initialise(mmp)
  
          for column in [0, 3, 4, 6, 9, 12, 13]:
              mmp.process_column(column, self.dist_matrix[:, :, column])
              part_dist_matrix[:, :, column] = self.dist_matrix[:, :, column]
  
          correct_mp, correct_index, correct_dims = self.bruteforce_mmp(part_dist_matrix)
  
          npt.assert_allclose(mmp.md_matrix_profile(), correct_mp)
          npt.assert_equal(mmp.md_profile_index(), correct_index)
          for i, dims in enumerate(correct_dims):
              sorted_mmp_dims = np.sort(mmp.md_profile_dimensions()[i], axis=0)
              npt.assert_equal(sorted_mmp_dims, dims)
  
      def test_process_column_lr(self):
          l_dist_matrix = self.dist_matrix.copy()
          r_dist_matrix = self.dist_matrix.copy()
  
          mmp = MultidimensionalMatrixProfileLR()
          self.mock_initialise(mmp)
  
          for diag in range(-self.dist_matrix.shape[1] + 1, self.dist_matrix.shape[2]):
              diag_ind = diag_indices_of(self.dist_matrix[0], diag)
              if diag < 0:
                  l_dist_matrix[:, diag_ind[0], diag_ind[1]] = np.inf
              else:
                  r_dist_matrix[:, diag_ind[0], diag_ind[1]] = np.inf
  
          for column in range(self.dist_matrix.shape[2]):
              mmp.process_column(column, self.dist_matrix[:, :, column])
  
          correct_lmp, correct_lindex, correct_ldims = self.bruteforce_mmp(l_dist_matrix)
          correct_rmp, correct_rindex, correct_rdims = self.bruteforce_mmp(r_dist_matrix)
  
          npt.assert_allclose(mmp.md_matrix_profile_left, correct_lmp)
          npt.assert_allclose(mmp.md_matrix_profile_right, correct_rmp)
          npt.assert_equal(mmp.md_profile_index_left, correct_lindex)
          npt.assert_equal(mmp.md_profile_index_right, correct_rindex)
          for i, dims in enumerate(correct_ldims):
              sorted_mmp_ldims = np.sort(mmp.md_profile_dimension_left[i], axis=0)
              sorted_mmp_rdims = np.sort(mmp.md_profile_dimension_right[i], axis=0)
              npt.assert_equal(sorted_mmp_ldims, correct_ldims[i])
              npt.assert_equal(sorted_mmp_rdims, correct_rdims[i])
  
      def test_streaming_process_column(self):
          mmp = MultidimensionalMatrixProfileLR()
          mmp.initialise(3, 5, 7)
  
          available_data = np.full((3, 5, 7), np.inf)
          available_data[:, 0:5, 0] = self.dist_matrix[:, 0:5, 0]
          available_data[:, 0:5, 4] = self.dist_matrix[:, 0:5, 4]
          available_data[:, 0:5, 5] = self.dist_matrix[:, 0:5, 5]
          mmp.process_column(0, self.dist_matrix[:, 0:5, 0])
          mmp.process_column(4, self.dist_matrix[:, 0:5, 4])
          mmp.process_column(5, self.dist_matrix[:, 0:5, 5])
  
          correct_mp, correct_index, correct_dims = self.bruteforce_mmp(available_data)
          npt.assert_allclose(mmp.md_matrix_profile(), correct_mp)
          npt.assert_equal(mmp.md_profile_index(), correct_index)
          for i, dims in enumerate(correct_dims):
              sorted_mmp_dims = np.sort(mmp.md_profile_dimensions()[i], axis=0)
              npt.assert_equal(sorted_mmp_dims, dims)
  
          mmp.shift_series(2)
          mmp.shift_query(3)
          available_data = np.full((3, 8, 9), np.inf)
          available_data[:, 0:5, 4] = self.dist_matrix[:, 0:5, 4]
          available_data[:, 0:5, 5] = self.dist_matrix[:, 0:5, 5]
  
          correct_mp, correct_index, correct_dims = self.bruteforce_mmp(available_data)
          npt.assert_allclose(mmp.md_matrix_profile(), correct_mp[:, 2:])
          npt.assert_equal(mmp.md_profile_index(), correct_index[:, 2:])
          for i, dims in enumerate(correct_dims):
              sorted_mmp_dims = np.sort(mmp.md_profile_dimensions()[i], axis=0)
              npt.assert_equal(sorted_mmp_dims, dims[:, 2:])
  
          mmp.process_column(0, self.dist_matrix[:, 3:8, 2])
          mmp.process_column(2, self.dist_matrix[:, 3:8, 4])
          mmp.process_column(3, self.dist_matrix[:, 3:8, 5])
          mmp.process_column(5, self.dist_matrix[:, 3:8, 7])
          available_data[:, 3:8, 2] = self.dist_matrix[:, 3:8, 2]
          available_data[:, 3:8, 4] = self.dist_matrix[:, 3:8, 4]
          available_data[:, 3:8, 5] = self.dist_matrix[:, 3:8, 5]
          available_data[:, 3:8, 7] = self.dist_matrix[:, 3:8, 7]
  
          correct_mp, correct_index, correct_dims = self.bruteforce_mmp(available_data)
          npt.assert_allclose(mmp.md_matrix_profile(), correct_mp[:, 2:])
          npt.assert_equal(mmp.md_profile_index(), correct_index[:, 2:])
          for i, dims in enumerate(correct_dims):
              sorted_mmp_dims = np.sort(mmp.md_profile_dimensions()[i], axis=0)
              npt.assert_equal(sorted_mmp_dims, dims[:, 2:])
  
      def test_streaming_process_diagonal(self):
          mmp = MultidimensionalMatrixProfileLR()
          mmp.initialise(3, 5, 7)
  
          available_data = np.full((3, 5, 7), np.inf)
          a = self.dist_matrix[:, 0:5, 0:7]
          available_data[diag3(available_data, 0)] = a[diag3(a, 0)]
          available_data[diag3(available_data, 4)] = a[diag3(a, 4)]
          available_data[diag3(available_data, -3)] = a[diag3(a, -3)]
          mmp.process_diagonal(0, a[diag3(a, 0)])
          mmp.process_diagonal(4, a[diag3(a, 4)])
          mmp.process_diagonal(-3, a[diag3(a, -3)])
  
          correct_mp, correct_index, correct_dims = self.bruteforce_mmp(available_data)
          npt.assert_allclose(mmp.md_matrix_profile(), correct_mp)
          npt.assert_equal(mmp.md_profile_index(), correct_index)
          for i, dims in enumerate(correct_dims):
              sorted_mmp_dims = np.sort(mmp.md_profile_dimensions()[i], axis=0)
              npt.assert_equal(sorted_mmp_dims, dims)
  
          mmp.shift_series(3)
          mmp.shift_query(2)
          available_data = np.full((3, 7, 10), np.inf)
          a_win = available_data[:, 0:5, 0:7]
          a = self.dist_matrix[:, 0:5, 0:7]
          a_win[diag3(a_win, 0)] = a[diag3(a, 0)]
          a_win[diag3(a_win, 4)] = a[diag3(a, 4)]
          a_win[diag3(a_win, -3)] = a[diag3(a, -3)]
  
          correct_mp, correct_index, correct_dims = self.bruteforce_mmp(available_data)
          npt.assert_allclose(mmp.md_matrix_profile(), correct_mp[:, 3:])
          npt.assert_equal(mmp.md_profile_index(), correct_index[:, 3:])
          for i, dims in enumerate(correct_dims):
              sorted_mmp_dims = np.sort(mmp.md_profile_dimensions()[i], axis=0)
              npt.assert_equal(sorted_mmp_dims, dims[:, 3:])
  
          a_win = available_data[:, 2:7, 3:10]
          a = self.dist_matrix[:, 2:7, 3:10]
          a_win[diag3(a_win, -1)] = a[diag3(a, -1)]  # Old diag 0
          a_win[diag3(a_win, 3)] = a[diag3(a, 3)]  # Old diag 4
          a_win[diag3(a_win, -4)] = a[diag3(a, -4)]  # Old diag -3
          a_win[diag3(a_win, 2)] = a[diag3(a, 2)]
          mmp.process_diagonal(-1, a[diag3(a, -1)])
          mmp.process_diagonal(3, a[diag3(a, 3)])
          mmp.process_diagonal(-4, a[diag3(a, -4)])
          mmp.process_diagonal(2, a[diag3(a, 2)])
  
          correct_mp, correct_index, correct_dims = self.bruteforce_mmp(available_data)
          npt.assert_allclose(mmp.md_matrix_profile(), correct_mp[:, 3:])
          npt.assert_equal(mmp.md_profile_index(), correct_index[:, 3:])
          for i, dims in enumerate(correct_dims):
              sorted_mmp_dims = np.sort(mmp.md_profile_dimensions()[i], axis=0)
              npt.assert_equal(sorted_mmp_dims, dims[:, 3:])
  
  
  def diag3(array, diag):
      """
      Returns the diagonal indices for a 3D array, taking all elements from the first dimension.
      """
      return (slice(None), *diag_indices_of(array[0], diag))
  