import unittest
import numpy as np
from numpy.testing import assert_array_equal
from scipy.sparse import coo_matrix
from sklearn.metrics.cluster import adjusted_rand_score
from graphs import Graph
from graphs.construction import neighbor_graph


class TestLabel(unittest.TestCase):

  def _make_blob_graphs(self, k=11):
    pts = np.random.random(size=(20, 2))
    pts[10:] += 2
    labels = np.zeros(20)
    labels[10:] = 1
    G_sparse = neighbor_graph(pts, k=k).symmetrize()
    G_dense = Graph.from_adj_matrix(G_sparse.matrix('dense'))
    return (G_sparse, G_dense), labels

  def test_greedy_coloring(self):
    pairs = np.array([[0,1],[0,2],[1,0],[1,2],[2,0],[2,1],[3,4],[4,3]])
    adj = [[0,1,1,0,0],
           [1,0,1,0,0],
           [1,1,0,0,0],
           [0,0,0,0,1],
           [0,0,0,1,0]]
    test_cases = [
        Graph.from_edge_pairs(pairs),
        Graph.from_adj_matrix(adj),
        Graph.from_adj_matrix(coo_matrix(adj)),
    ]
    for G in test_cases:
      assert_array_equal([1,2,3,1,2], G.color_greedy())

  def test_bicolor_spectral(self):
    pairs = np.array([[0,1],[0,2],[1,0],[1,2],[2,0],[2,1],[2,3],[3,2]])
    adj = [[0,1,1,0],
           [1,0,1,0],
           [1,1,0,1],
           [0,0,1,0]]
    test_cases = [
        Graph.from_edge_pairs(pairs),
        Graph.from_adj_matrix(adj),
        Graph.from_adj_matrix(coo_matrix(adj)),
    ]
    expected = np.array([1,1,0,1], dtype=bool)
    for G in test_cases:
      assert_array_equal(expected, G.bicolor_spectral())

  def test_spectral_clustering(self):
    blob_graphs, expected = self._make_blob_graphs(k=11)

    for g in blob_graphs:
      labels = g.cluster_spectral(2, kernel='rbf')
      self.assertGreater(adjusted_rand_score(expected, labels), 0.95)

  def test_nn_classifier(self):
    blob_graphs, expected = self._make_blob_graphs(k=4)
    partial = expected.copy()
    partial[1:-1] = -1

    for g in blob_graphs:
      labels = g.classify_nearest(partial)
      self.assertGreater(adjusted_rand_score(expected, labels), 0.95)

  def test_lgc_classifier(self):
    blob_graphs, expected = self._make_blob_graphs(k=11)
    partial = expected.copy()
    partial[1:-1] = -1

    for g in blob_graphs:
      labels = g.classify_lgc(partial, kernel='rbf', alpha=0.2, tol=1e-3,
                              max_iter=30)
      self.assertGreater(adjusted_rand_score(expected, labels), 0.95)

  def test_harmonic_classifier(self):
    blob_graphs, expected = self._make_blob_graphs(k=4)
    partial = expected.copy()
    partial[1:-1] = -1

    for g in blob_graphs:
      labels = g.classify_harmonic(partial, use_CMN=True)
      self.assertGreater(adjusted_rand_score(expected, labels), 0.95)

  def test_regression(self):
    t = np.linspace(0, 1, 31)
    pts = np.column_stack((np.sin(t), np.cos(t)))
    G = neighbor_graph(pts, k=3).symmetrize()
    y_mask = slice(None, None, 2)

    # test the interpolated case
    x = G.regression(t[y_mask], y_mask)
    assert_array_equal(t, np.linspace(0, 1, 31))  # ensure t hasn't changed
    self.assertLess(np.linalg.norm(t - x), 0.15)

    # test the boolean mask case
    y_mask = np.zeros_like(t, dtype=bool)
    y_mask[::2] = True
    x = G.regression(t[y_mask], y_mask)
    self.assertLess(np.linalg.norm(t - x), 0.15)

    # test the penalized case
    x = G.regression(t[y_mask], y_mask, smoothness_penalty=1e-4)
    self.assertLess(np.linalg.norm(t - x), 0.15)

    # test no kernel + dense laplacian case
    dG = Graph.from_adj_matrix(G.matrix('dense'))
    x = dG.regression(t[y_mask], y_mask, kernel='none')
    self.assertLess(np.linalg.norm(t - x), 0.25)
    x = dG.regression(t[y_mask], y_mask, smoothness_penalty=1e-4, kernel='none')
    self.assertLess(np.linalg.norm(t - x), 0.25)

    # test the multidimensional regression case
    tt = np.column_stack((t, t[::-1]))
    x = G.regression(tt[y_mask], y_mask)
    self.assertLess(np.linalg.norm(tt - x), 0.2)

    # check for bad inputs
    with self.assertRaisesRegexp(ValueError, r'^Invalid shape of y array'):
      G.regression([], y_mask)

if __name__ == '__main__':
  unittest.main()
