#!/usr/bin/env python

import numpy as np
import numpy.random as nr
from numpy.testing import(assert_array_equal, assert_array_almost_equal,
                          assert_almost_equal)
from nose.tools import assert_true, assert_equal

from ..graph import (WeightedGraph, complete_graph, mst, knn, eps_nn, 
                     wgraph_from_adjacency, wgraph_from_coo_matrix, 
                     concatenate_graphs, wgraph_from_3d_grid)



def basicdata():
    x = np.array( [[- 1.998, - 2.024], [- 0.117, - 1.010], [1.099, - 0.057],
                   [ 1.729, - 0.252], [1.003, - 0.021], [1.703, - 0.739],
                   [- 0.557, 1.382],[- 1.200, - 0.446],[- 0.331, - 0.256],
                   [- 0.800, - 1.584]])
    return x


def basic_graph():
    l = np.linspace(0, 2 * np.pi, 20, endpoint=False)
    x = np.column_stack((np.cos(l), np.sin(l)))
    G = knn(x, 2)
    return G


def basic_graph_2():
    l = np.linspace(0, 2 * np.pi, 20, endpoint=False)
    x = np.column_stack((np.cos(l), np.sin(l)))
    G = knn(x, 2)
    return G, x


def test_complete():
    v = 10
    G = complete_graph(v)
    a = G.get_edges()[:, 0]
    b = G.get_edges()[:, 1]
    inds = np.indices((v, v)).reshape( (2, v * v) )
    assert_array_equal(inds, (a, b))


def test_knn_1():
    x = basicdata()
    G = knn(x, 1)
    A = G.get_edges()[:, 0]
    assert_equal(np.shape(A)[0], 14)

    
def test_set_euclidian():
    G, x = basic_graph_2()
    d = G.weights
    G.set_euclidian(x / 10)
    D = G.weights
    assert_true(np.allclose(D, d / 10, 1e-7))


def test_set_gaussian():
    G, x = basic_graph_2()
    d = G.weights
    G.set_gaussian(x, 1.0)
    D = G.weights
    assert_true(np.allclose(D, np.exp(- d * d / 2), 1e-7))


def test_set_gaussian_2():
    G, x = basic_graph_2()
    d = G.weights
    G.set_gaussian(x)
    D = G.weights
    sigma = np.sum(d * d) / len(d)
    assert_true(np.allclose(D, np.exp(-d * d / (2 * sigma)), 1e-7))
    

def test_eps_1():
    x = basicdata()
    G = eps_nn(x, 1.)
    D = G.weights
    assert_equal(np.size(D), 16)
    assert_true((D < 1).all())

    
def test_mst_1():
    x = basicdata()
    G = mst(x)
    D = G.weights
    assert_equal(np.size(D), 18)


def test_3d_grid():
    """test the 6nn graph
    """
    x0 = np.array([0, 0, 0])
    x1 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1], [-1, 0, 0], [0, -1, 0],
                   [0, 0, -1]])
    x2 = np.array([[1, 1, 0], [0, 1, 1], [1, 0, 1], [1, -1, 0], [0, 1, -1],
                   [1, 0, -1], [-1, 1, 0], [0, -1, 1], [-1, 0, 1], 
                   [-1, -1, 0], [-1, 0, -1], [0, -1, -1]])
    x3 = np.array([[1, 1, 1], [1, 1, -1], [1, -1, 1], [1, -1, -1],
                   [-1, 1, 1], [-1, 1, -1], [-1, -1, 1], [-1, -1, -1]])
    for x in x1:
        xyz = np.vstack((x0, x))
        assert_equal(wgraph_from_3d_grid(xyz, 6).E, 2)
        assert_equal(wgraph_from_3d_grid(xyz, 18).E, 2)
        assert_equal(wgraph_from_3d_grid(xyz, 26).E, 2)
        for x in x2:
            xyz = np.vstack((x0, x))
            assert_equal(wgraph_from_3d_grid(xyz, 6).E, 0)
            assert_equal(wgraph_from_3d_grid(xyz, 18).E, 2)
            assert_equal(wgraph_from_3d_grid(xyz, 26).E, 2)
        for x in x3:
            xyz = np.vstack((x0, x))
            assert_equal(wgraph_from_3d_grid(xyz, 6).E, 0)
            assert_equal(wgraph_from_3d_grid(xyz, 18).E, 0)
            assert_equal(wgraph_from_3d_grid(xyz, 26).E, 2)
            

def test_grid_3d_1():
    """ Test the 6 nn graphs on 3d grid
    """
    nx, ny, nz = 9, 6, 1
    xyz = np.mgrid[0:nx, 0:ny, 0:nz]
    xyz = np.reshape(xyz, (3, nx * ny * nz)).T
    G = wgraph_from_3d_grid(xyz, 6)
    assert_equal(G.E, 186)

    
def test_grid_3d_2():
    """ Test the 18-nn graph on a 3d grid
    """
    nx, ny, nz = 9, 6, 1
    xyz = np.mgrid[0:nx, 0:ny, 0:nz]
    xyz = np.reshape(xyz,(3, nx * ny * nz)).T
    G = wgraph_from_3d_grid(xyz, 18)
    assert_equal(G.E, 346)
        

def test_grid_3d_3():
    """ Test the 26-nn graph on a 3d grid
    """
    nx, ny, nz = 9, 6, 1
    xyz = np.mgrid[0:nx, 0:ny, 0:nz]
    xyz = np.reshape(xyz,(3, nx * ny * nz)).T
    G = wgraph_from_3d_grid(xyz, 26)
    assert_equal(G.E, 346)


def test_grid_3d_4():
    nx, ny, nz = 10, 10, 10
    xyz = np.reshape(np.indices((nx, ny, nz)), (3, nx * ny * nz)).T
    G = wgraph_from_3d_grid(xyz, 26)
    D = G.weights
    # 6 * 9 * 10 * 10
    assert_equal(sum(D == 1), 5400 )
    # 26 * 8 ** 3 + 6 * 8 ** 2 * 17 + 12 * 8 * 11 + 8 * 7 
    assert_equal(np.size(D), 20952 )
    # 18 * 8 ** 3 + 6 * 8 ** 2 * 13 + 12 * 8 * 9 + 8 * 6
    assert_equal(sum(D < 1.5), 15120)


def test_grid_3d_5():
    nx, ny, nz = 5, 5, 5
    xyz = np.reshape(np.indices((nx, ny, nz)), (3, nx * ny * nz)).T
    G = wgraph_from_3d_grid(xyz, 26)
    D = G.weights.copy()
    G.set_euclidian(xyz)
    assert_array_almost_equal(G.weights, D)


def test_grid_3d_6():
    nx, ny, nz = 5, 5, 5
    xyz = np.reshape(np.indices((nx, ny, nz)), (3, nx * ny * nz)).T
    adj = wgraph_from_3d_grid(xyz, 26).to_coo_matrix().tolil()
    assert_equal(len(adj.rows[63]), 26)
    for i in [62, 64, 58, 68, 38, 88, 57, 67, 37, 87, 59, 69, 39, 89, 33, 
              83, 43, 93, 32, 82, 42, 92, 34, 84, 44, 94]:
        assert_true(i in adj.rows[63])


def test_grid_3d_7():
    """ Check that the grid graph is symmetric
    """
    xyz = np.array(np.where(np.random.rand(5, 5, 5) > 0.5)).T
    adj = wgraph_from_3d_grid(xyz, 6).to_coo_matrix()
    assert_equal((adj - adj.T).nnz, 0)
    adj = wgraph_from_3d_grid(xyz, 18).to_coo_matrix()
    assert_equal((adj - adj.T).nnz, 0)
    adj = wgraph_from_3d_grid(xyz, 26).to_coo_matrix()
    assert_equal((adj - adj.T).nnz, 0)
        
    
def test_cut_redundancies():
    G = basic_graph()
    e = G.E
    edges = G.get_edges()
    weights = G.weights
    G.E = 2 * G.E
    G.edges = np.concatenate((edges, edges))
    G.weights = np.concatenate((weights, weights))
    K = G.cut_redundancies()