"""
  Test cases for functions/classes in nn_comparator.py
  -- kandasamy@cs.cmu.edu
"""

# pylint: disable=no-member
# pylint: disable=invalid-name

import numpy as np
# Local imports
from . import otmann
from .unittest_neural_network import generate_cnn_architectures, \
                                     generate_mlp_architectures
from ..utils.ancillary_utils import get_list_of_floats_as_str
from ..utils.base_test_class import BaseTestClass, execute_tests

_TOL = 1e-5


class TransportNNDistanceComputerTestCase(BaseTestClass):
  """ Contains unit tests for the TransportNNDistanceComputer class. """

  def __init__(self, *args, **kwargs):
    """ Constructor. """
    super(TransportNNDistanceComputerTestCase, self).__init__(*args, **kwargs)
    self.non_assignment_penalty = 1
    cnn_layer_labels, label_mismatch_penalty = \
      otmann.get_cnn_layer_label_mismatch_penalties(self.non_assignment_penalty)
    self.tp_comp = otmann.OTMANNDistanceComputer(cnn_layer_labels,
                     label_mismatch_penalty, self.non_assignment_penalty,
                     otmann.CNN_STRUCTURAL_PENALTY_GROUPS,
                     otmann.PATH_LENGTH_TYPES,
                     dflt_mislabel_coeffs=1.0, dflt_struct_coeffs=1.0,
                     dflt_dist_type='lp-emd')
    self.cnns = generate_cnn_architectures()

  def test_cnn_label_mismatch_penalties(self):
    """ Unit test for the label mismatch penalty of a CNN. """
    self.report('Testing generation of label mismatches for a CNN. ')
    cnn_layer_labels, label_mismatch_penalty = \
      otmann.get_cnn_layer_label_mismatch_penalties(self.non_assignment_penalty,
                       max_conv_size=9)
    self.report('cnn_layer_labels: %s'%(str(cnn_layer_labels)), 'test_result')
    self.report('cnn mismatch penalties: \n%s'%(str(np.round(label_mismatch_penalty, 3))),
                                            'test_result')
    assert np.all(label_mismatch_penalty == label_mismatch_penalty.T)
    assert np.all(np.diag(label_mismatch_penalty) == 0)

  def test_mlp_label_mismatch_penalties(self):
    """ Unit test for the label mismatch penalty of an MLP. """
    self.report('Testing generation of label mismatches for a MLP. ')
    mlp_layer_labels, label_mismatch_penalty = \
      otmann.get_mlp_layer_label_mismatch_penalties(self.non_assignment_penalty,
                                                            'reg')
    self.report('mlp_layer_labels: %s'%((str(mlp_layer_labels))), 'test_result')
    self.report('mlp mismatch penalties: \n%s'%(str(np.round(label_mismatch_penalty, 3))),
                'test_result')
    assert np.all(label_mismatch_penalty == label_mismatch_penalty.T)
    assert np.all(np.diag(label_mismatch_penalty) == 0)

  @classmethod
  def _is_cost_matrix_for_same_networks(cls, cost_matrix):
    """ Returns true if it is the cost matrix for the same network. """
    return np.all(np.diag(cost_matrix) == 0) and np.all(cost_matrix == cost_matrix.T)

  @classmethod
  def _has_corresponding_layers(cls, cost_matrix):
    """ Returns true if one network has a corresponding layer in the other and vice
        versa. """
    ret = True
    for row_idx in range(cost_matrix.shape[0]):
      ret = ret and np.any(cost_matrix[row_idx, :] == 0)
    for col_idx in range(cost_matrix.shape[1]):
      ret = ret and np.any(cost_matrix[:, col_idx] == 0)
    return ret

  def test_mislabel_cost_matrix(self):
    """ Tests the mislabel cost matrix for specific pairs of neural networks. """
    self.report('Testing generation of label cost matrices for specific cnns. ')
    num_cnns = len(self.cnns)
    for i in range(num_cnns):
      for j in range(i+1, num_cnns):
        cnn_i = self.cnns[i]
        cnn_j = self.cnns[j]
        mislabel_cost_matrix = self.tp_comp.get_mislabel_cost_matrix(cnn_i, cnn_j)
        assert mislabel_cost_matrix.shape[0] == cnn_i.num_layers
        assert mislabel_cost_matrix.shape[1] == cnn_j.num_layers
        if i == j:
          assert self._is_cost_matrix_for_same_networks(mislabel_cost_matrix)
        if i == 0 and j == 1:
          # These two matrices were designed so that there has to be a zero for each
          # column on each row and vice versa.
          assert self._has_corresponding_layers(mislabel_cost_matrix)
        if (i == 2 and j == 3) or (i == 0 and j == 1) or (i == 1 and j == 6):
          self.report('Mislabel cost matrix for cnn-%d and cnn-%d:\n%s'%(i, j,
                      str(np.round(mislabel_cost_matrix, 3))), 'test_result')

  def test_connectivity_cost_matrix(self):
    """ Tests the connectivity cost matrix for specific pairs of neural networks.  """
    self.report('Testing generation of connectivity cost matrices for specific cnns.')
    num_cnns = len(self.cnns)
    for i in range(num_cnns):
      for j in range(i, num_cnns):
        cnn_i = self.cnns[i]
        cnn_j = self.cnns[j]
        struct_cost_matrix = self.tp_comp.get_struct_cost_matrix(cnn_i, cnn_j)
        assert struct_cost_matrix.shape[0] == cnn_i.num_layers
        assert struct_cost_matrix.shape[1] == cnn_j.num_layers
        if i == j:
          assert self._is_cost_matrix_for_same_networks(struct_cost_matrix)
        if i == 0 and j == 1:
          # These two matrices were designed so that there has to be a zero for each
          # column on each row and vice versa.
          assert self._has_corresponding_layers(struct_cost_matrix)
        if (i == 2 and j == 3) or (i == 0 and j == 1) or (i == 1 and j == 6):
          self.report('Structural cost matrix for cnn-%d and cnn-%d:\n%s'%(i, j,
                      str(np.round(struct_cost_matrix, 3))), 'test_result')

  def test_ot_cost_matrix(self):
    """ Tests the OT cost matrix for specific pairs of neural networks.  """
    self.report('Testing generation of OT cost matrices for specific cnns.')
    nns = self.cnns
    num_nns = len(nns)
    for i in range(num_nns):
      for j in range(i, num_nns):
        nn_i = nns[i]
        nn_j = nns[j]
        mislabel_cost_matrix = self.tp_comp.get_mislabel_cost_matrix(nn_i, nn_j)
        struct_cost_matrix = self.tp_comp.get_struct_cost_matrix(nn_i, nn_j)
        ot_cost_matrix = self.tp_comp.get_ot_cost_matrix(mislabel_cost_matrix,
                           struct_cost_matrix, 1, 0.1, 1, None)
        assert ot_cost_matrix.shape[0] == nn_i.num_layers + 1
        assert ot_cost_matrix.shape[1] == nn_j.num_layers + 1
        if i == j:
          assert self._is_cost_matrix_for_same_networks(ot_cost_matrix)
        if i == 0 and j == 1:
          # These two matrices were designed so that there has to be a zero for each
          # column on each row and vice versa.
          assert self._has_corresponding_layers(ot_cost_matrix)
        if (i == 2 and j == 3) or (i == 0 and j == 1) or (i == 1 and j == 6):
          self.report('OT cost matrix for cnn-%d and cnn-%d:\n%s'%(i, j,
                      str(np.round(ot_cost_matrix, 3))), 'test_result')

  @classmethod
  def _get_dist_type_abbr(cls, dist_type):
    """ Shortens distance type. """
    if dist_type == 'lp-norm-by-max':
      return 'lnbm'
    else:
      return dist_type

  def _test_dist_comp_for_single_conn_coeff(self, nns, dist_types, tp_comp):
    """ Tests distance computation for a single connectivity coefficient. """
    num_nns = len(nns)
    for i in range(num_nns):
      for j in range(i, num_nns):
        nn_i = nns[i]
        nn_j = nns[j]
        dists = {}
        for dt in dist_types:
          dists[dt] = tp_comp(nn_i, nn_j, dist_type=dt)
        res_str = ' '.join(['%s=%s'%(self._get_dist_type_abbr(key),
                                     get_list_of_floats_as_str(val))
                            for key, val in dists.items()])
        self.report('(i,j)=(%d,%d) %s'%(i, j, res_str), 'test_result')

  def _test_dist_comp_for_multiple_coeffs(self, nns, dist_types, mislabel_coeffs,
                                          struct_coeffs, tp_comp):
    """ Tests distance computation for a single connectivity coefficient. """
    num_nns = len(nns)
    for i in range(num_nns):
      for j in range(i, num_nns):
        nn_i = nns[i]
        nn_j = nns[j]
        dists = {}
        for dt in dist_types:
          dists[dt] = tp_comp(nn_i, nn_j, dist_type=dt, mislabel_coeffs=mislabel_coeffs,
                              struct_coeffs=struct_coeffs)
          num_dists = len(dt.split('-')) * len(mislabel_coeffs)
          assert len(dists[dt]) == num_dists
        res_str = ' '.join(['%s=%s'%(self._get_dist_type_abbr(key),
                                     get_list_of_floats_as_str(val))
                            for key, val in dists.items()])
        self.report('(i,j)=(%d,%d) %s'%(i, j, res_str), 'test_result')

  def test_cnn_distance_computation(self):
    """ Tests the computation of the distance for CNNs. """
    struct_coeffs = [0.1, 0.4]
    mislabel_coeffs = [1.0] * len(struct_coeffs)
    dist_types = ['lp-emd', 'lp', 'emd']
    self.report('Testing distance computation for specific cnns with default coeffs.')
    self._test_dist_comp_for_single_conn_coeff(self.cnns, dist_types, self.tp_comp)
    # With multiple conn_coeffs
    self.report('Testing distance computation for specific cnns with conn_coeffs=%s.'%(
                struct_coeffs))
    self._test_dist_comp_for_multiple_coeffs(self.cnns, dist_types, mislabel_coeffs,
                                             struct_coeffs, self.tp_comp)

  def test_mlp_distance_computation(self):
    """ Tests the computation of the distance for CNNs. """
    self.report('Testing distance computation for specific mlps.')
    # Create the transport computation object
    mlp_layer_labels, label_mismatch_penalty = \
      otmann.get_mlp_layer_label_mismatch_penalties(self.non_assignment_penalty,
                                                            'reg')
    mlp_tp_comp = otmann.OTMANNDistanceComputer(mlp_layer_labels,
                     label_mismatch_penalty,
                     self.non_assignment_penalty,
                     otmann.MLP_STRUCTURAL_PENALTY_GROUPS,
                     otmann.PATH_LENGTH_TYPES,
                     dflt_mislabel_coeffs=1.0, dflt_struct_coeffs=1.0)
    # Create the mlp architectures
    mlps = generate_mlp_architectures()
    struct_coeffs = [0.1, 0.2]
    dist_types = ['lp-emd', 'lp', 'emd']
    mislabel_coeffs = [1.0] * len(struct_coeffs)
    self.report('Testing distance computation for specific mlps with default coeffs.')
    self._test_dist_comp_for_single_conn_coeff(mlps, dist_types, mlp_tp_comp)
    self.report('Testing distance computation for specific mlps with conn_coeffs=%s.'%(
                struct_coeffs))
    self._test_dist_comp_for_multiple_coeffs(mlps, dist_types, mislabel_coeffs,
                                             struct_coeffs, mlp_tp_comp)


class DistProdNNKernelTestCase(BaseTestClass):
  """ Unit tests for the Transport NNKernels. """

  def __init__(self, *args, **kwargs):
    """ Constructor. """
    super(DistProdNNKernelTestCase, self).__init__(*args, **kwargs)
    self.non_assignment_penalty = 1
    cnn_layer_labels, label_mismatch_penalty = \
      otmann.get_cnn_layer_label_mismatch_penalties(self.non_assignment_penalty)
    self.all_layer_labels = cnn_layer_labels
    self.label_mismatch_penalty = label_mismatch_penalty
    self.tp_comp = otmann.OTMANNDistanceComputer(cnn_layer_labels,
                     label_mismatch_penalty,
                     self.non_assignment_penalty,
                     otmann.CNN_STRUCTURAL_PENALTY_GROUPS,
                     otmann.PATH_LENGTH_TYPES
                     )
    self.mislabel_coeffs = [2.0, 2.0, 1.0, 1.0, 1.0]
    self.struct_coeffs = [0.25, 0.5, 1.0, 2.0, 4.0]
    self.lp_betas = [1e-6] * len(self.struct_coeffs)
    self.emd_betas = [1] * len(self.struct_coeffs)
    self.scale = 1
    self.cnns = generate_cnn_architectures()

  def test_instantiation(self):
    """ Testing instantiation. """
    self.report('Testing instantiation of DistProdNNKernelTestCase and computation ' +
                'for specific networks.')
    dist_type_vals = ['lp', 'emd', 'lp-emd']
    all_kernels = []
    for dist_type in dist_type_vals:
      if dist_type == 'lp':
        betas = self.lp_betas
      elif dist_type == 'emd':
        betas = self.emd_betas
      else:
        betas = [j for i in zip(self.lp_betas, self.emd_betas) for j in i]
      tp_kernel = otmann.get_otmann_kernel_from_params('prod',
                    self.all_layer_labels, self.label_mismatch_penalty,
                    self.non_assignment_penalty,
                    otmann.CNN_STRUCTURAL_PENALTY_GROUPS,
                    otmann.PATH_LENGTH_TYPES,
                    self.mislabel_coeffs, self.struct_coeffs, dist_type,
                    betas, self.scale)
      cnn_K = tp_kernel(self.cnns)
      all_kernels.append(cnn_K)
      cnn_eig_vals, _ = np.linalg.eig(cnn_K)
      self.report('dist-type: %s, eigvals: %s.'%(dist_type,
                     get_list_of_floats_as_str(sorted(cnn_eig_vals))))
      self.report('%s transport kernel:\n%s'%(dist_type,
                  str(np.round(cnn_K, 3))), 'test_result')
      assert cnn_K.shape == (len(self.cnns), len(self.cnns))
      assert np.all(np.diag(cnn_K) == 1)
    # Check if it is in fact the product
    if 'lp' in dist_type_vals and 'emd' in dist_type_vals and 'lp-emd' in dist_type_vals:
      lp_kernel = all_kernels[dist_type_vals.index('lp')]
      emd_kernel = all_kernels[dist_type_vals.index('emd')]
      lpemd_kernel = all_kernels[dist_type_vals.index('lp-emd')]
      assert np.linalg.norm(lpemd_kernel - lp_kernel * emd_kernel) < _TOL

  def test_kernel_computation(self):
    """ Testing computation of the lp distance. """
    self.report('Testing computed kernel values for specific cnns.')
    betas = [0.0001]
    scale = 2.1
    struct_coeffs = 1.0
    mislabel_coeffs = 1.0
    dist_type = 'lp'
    tp_comp = otmann.OTMANNDistanceComputer(self.all_layer_labels,
                self.label_mismatch_penalty, self.non_assignment_penalty,
                otmann.CNN_STRUCTURAL_PENALTY_GROUPS,
                otmann.PATH_LENGTH_TYPES,
                dflt_mislabel_coeffs=mislabel_coeffs,
                dflt_struct_coeffs=struct_coeffs,
                dflt_dist_type=dist_type)
    tp_kernel = otmann.get_otmann_kernel_from_params('prod',
                self.all_layer_labels, self.label_mismatch_penalty,
                self.non_assignment_penalty, otmann.CNN_STRUCTURAL_PENALTY_GROUPS,
                otmann.PATH_LENGTH_TYPES,
                mislabel_coeffs, struct_coeffs, dist_type, betas, scale)
    cnn_dists = tp_comp(self.cnns, self.cnns)
    cnn_kernel = tp_kernel(self.cnns)
    diff = np.linalg.norm(scale * np.exp(-cnn_dists[0] * betas[0]) - cnn_kernel)
    assert diff < np.linalg.norm(cnn_kernel) * 1e-6


class DistSumNNKernelTestCase(BaseTestClass):
  """ Unit tests for the Transport NNKernels. """

  def __init__(self, *args, **kwargs):
    """ Constructor. """
    super(DistSumNNKernelTestCase, self).__init__(*args, **kwargs)
    self.non_assignment_penalty = 1
    mlp_layer_labels, label_mismatch_penalty = \
      otmann.get_mlp_layer_label_mismatch_penalties(self.non_assignment_penalty,
                                                            'reg')
    self.all_layer_labels = mlp_layer_labels
    self.label_mismatch_penalty = label_mismatch_penalty
    self.tp_comp = otmann.OTMANNDistanceComputer(mlp_layer_labels,
                     label_mismatch_penalty,
                     self.non_assignment_penalty,
                     otmann.MLP_STRUCTURAL_PENALTY_GROUPS,
                     otmann.PATH_LENGTH_TYPES
                     )
    self.mislabel_coeffs = [2.0, 2.0, 1.0, 1.0, 1.0]
    self.struct_coeffs = [0.25, 0.5, 1.0, 2.0, 4.0]
    self.lp_betas = [1e-6] * len(self.struct_coeffs)
    self.emd_betas = [1] * len(self.struct_coeffs)
    self.mlps = generate_mlp_architectures()

  def test_instantiation_and_computation(self):
    """ Testing instantiation. """
    self.report('Testing instantiation of DistSumNNKernelTestCase and computation ' +
                'for specific networks.')
    dist_type_vals = ['lp', 'emd', 'lp-emd']
    all_kernels = []
    for dist_type in dist_type_vals:
      if dist_type == 'lp':
        betas = self.lp_betas
        scales = [1]
      elif dist_type == 'emd':
        betas = self.emd_betas
        scales = [1]
      else:
        betas = [j for i in zip(self.lp_betas, self.emd_betas) for j in i]
        scales = [1, 1]
      tp_kernel = otmann.get_otmann_kernel_from_params('sum',
                    self.all_layer_labels, self.label_mismatch_penalty,
                    self.non_assignment_penalty,
                    otmann.MLP_STRUCTURAL_PENALTY_GROUPS,
                    otmann.PATH_LENGTH_TYPES,
                    self.mislabel_coeffs, self.struct_coeffs, dist_type, betas, scales)
      nn_K = tp_kernel(self.mlps)
      nn_eig_vals, _ = np.linalg.eig(nn_K)
      self.report('dist-type: %s, eigvals: %s.'%(dist_type,
                     get_list_of_floats_as_str(sorted(nn_eig_vals))))
      assert nn_K.shape == (len(self.mlps), len(self.mlps))
      self.report('%s transport kernel:\n%s'%(dist_type,
                  str(np.round(nn_K, 3))), 'test_result')
      assert np.all(np.diag(nn_K) == sum(scales))
      all_kernels.append(nn_K)
    # Check if it is in fact the sum
    if 'lp' in dist_type_vals and 'emd' in dist_type_vals and 'lp-emd' in dist_type_vals:
      lp_kernel = all_kernels[dist_type_vals.index('lp')]
      emd_kernel = all_kernels[dist_type_vals.index('emd')]
      lpemd_kernel = all_kernels[dist_type_vals.index('lp-emd')]
      assert np.linalg.norm(lpemd_kernel - lp_kernel - emd_kernel) < _TOL

  def test_sum_product_equivalence(self):
    """ Unit-test for testing that both kernels compute the same thing in certain cases.
    """
    dist_type_vals = ['lp', 'emd']
    for dist_type in dist_type_vals:
      if dist_type == 'lp':
        betas = self.lp_betas
        scales = [1]
      elif dist_type == 'emd':
        betas = self.emd_betas
        scales = [1]
      sum_kernel = otmann.get_otmann_kernel_from_params('sum',
                     self.all_layer_labels, self.label_mismatch_penalty,
                     self.non_assignment_penalty,
                     otmann.MLP_STRUCTURAL_PENALTY_GROUPS,
                     otmann.PATH_LENGTH_TYPES,
                     self.mislabel_coeffs, self.struct_coeffs, dist_type, betas, scales)
      prod_kernel = otmann.get_otmann_kernel_from_params('prod',
                      self.all_layer_labels, self.label_mismatch_penalty,
                      self.non_assignment_penalty,
                      otmann.MLP_STRUCTURAL_PENALTY_GROUPS,
                      otmann.PATH_LENGTH_TYPES,
                      self.mislabel_coeffs, self.struct_coeffs, dist_type, betas, scales)
      sum_nn_K = sum_kernel(self.mlps)
      prod_nn_K = prod_kernel(self.mlps)
      assert np.linalg.norm(sum_nn_K - prod_nn_K) < _TOL


if __name__ == '__main__':
  execute_tests()

