import networkx as nx
import numpy as np
import os
from scipy import sparse
from sklearn import manifold
import unittest
import warnings

import csrgraph as cg
import nodevectors


class TestGraphEmbedding(unittest.TestCase):
    """
    higher level tests
    """
    def test_n2v(self):
        tt = nx.generators.complete_graph(50)
        wordsize = 32
        # Gensim triggers deprecation warnings...
        warnings.simplefilter("ignore", category=PendingDeprecationWarning)
        warnings.simplefilter("ignore", category=DeprecationWarning)
        g2v = nodevectors.Node2Vec(
            walklen=5, 
            epochs=5,
            threads=6,
            n_components=wordsize,
            keep_walks=True,
            verbose=False,
            w2vparams={"window":3, "negative":3, "iter":3,
                       "batch_words":32, "workers": 2})
        g2v2 = nodevectors.Node2Vec(
            walklen=5, 
            epochs=5,
            threads=6, 
            n_components=wordsize,
            keep_walks=False,
            verbose=False,
            w2vparams={"window":3, "negative":3, "iter":3,
                       "batch_words":32, "workers": 2})
        g2v.fit(tt)
        g2v2.fit(tt)
        res_v = g2v.predict(9)
        self.assertTrue(len(res_v) == wordsize)
        self.assertTrue(len(g2v2.predict(9)) == wordsize)
        self.assertTrue(hasattr(g2v, 'walks'))
        self.assertFalse(hasattr(g2v2, 'walks'))
        warnings.resetwarnings()
        # Test save/load
        fname = 'test_saving'
        try:
            g2v.save(fname)
            g2v_l = nodevectors.Node2Vec.load(fname + '.zip')
            res_l = g2v_l.predict(9)
            self.assertTrue(len(res_l) == wordsize)
            np.testing.assert_array_almost_equal(res_l, res_v)
        finally:
            os.remove(fname + '.zip')

    def test_skl(self):
        tt = nx.generators.complete_graph(25)
        ndim = 3
        skle = nodevectors.SKLearnEmbedder(
            manifold.Isomap, 
            n_components=ndim,
            n_neighbors=3)
        skle.fit(tt)
        res_v = skle.predict(9)
        self.assertTrue(len(res_v) == ndim)
        # Test save/load
        fname = 'test_saving'
        try:
            skle.save(fname)
            g2v_l = nodevectors.SKLearnEmbedder.load(fname + '.zip')
            res_l = g2v_l.predict(9)
            self.assertTrue(len(res_l) == ndim)
            np.testing.assert_array_almost_equal(res_l, res_v)
        finally:
            os.remove(fname + '.zip')

    def test_ggvec(self):
        tt = nx.generators.complete_graph(25)
        ndim = 3
        skle = nodevectors.GGVec(
            n_components=ndim,
            tol=0.15,
            max_loss=15,
            learning_rate=0.1,
            negative_ratio=0.3,
            order=1,
            verbose=False,
            max_epoch=2000)
        skle.fit(tt)
        res_v = skle.predict(9)
        self.assertTrue(len(res_v) == ndim)
        # Test save/load
        fname = 'test_saving'
        try:
            skle.save(fname)
            g2v_l = nodevectors.SKLearnEmbedder.load(fname + '.zip')
            res_l = g2v_l.predict(9)
            self.assertTrue(len(res_l) == ndim)
            np.testing.assert_array_almost_equal(res_l, res_v)
        finally:
            os.remove(fname + '.zip')

    def test_ggvec_order2(self):
        tt = nx.generators.complete_graph(25)
        ndim = 3
        skle = nodevectors.GGVec(
            n_components=ndim,
            tol=0.15,
            max_loss=15,
            learning_rate=0.1,
            negative_ratio=0.3,
            order=2,
            verbose=False,
            max_epoch=2000)
        skle.fit(tt)
        res_v = skle.predict(9)
        self.assertTrue(len(res_v) == ndim)
        # Test save/load
        fname = 'test_saving'
        try:
            skle.save(fname)
            g2v_l = nodevectors.SKLearnEmbedder.load(fname + '.zip')
            res_l = g2v_l.predict(9)
            self.assertTrue(len(res_l) == ndim)
            np.testing.assert_array_almost_equal(res_l, res_v)
        finally:
            os.remove(fname + '.zip')

    def test_glove(self):
        tt = nx.generators.complete_graph(25)
        ndim = 3
        skle = nodevectors.Glove(
            n_components=ndim,
            tol=0.005,
            max_loss=15,
            learning_rate=0.1,
            verbose=False,
            max_epoch=2000)
        skle.fit(tt)
        res_v = skle.predict(9)
        self.assertTrue(len(res_v) == ndim)
        # Test save/load
        fname = 'test_saving'
        try:
            skle.save(fname)
            g2v_l = nodevectors.SKLearnEmbedder.load(fname + '.zip')
            res_l = g2v_l.predict(9)
            self.assertTrue(len(res_l) == ndim)
            np.testing.assert_array_almost_equal(res_l, res_v)
        finally:
            os.remove(fname + '.zip')

    def test_prone(self):
        tt = nx.generators.complete_graph(25)
        ndim = 3
        skle = nodevectors.ProNE(n_components=ndim)
        skle.fit(tt)
        res_v = skle.predict(9)
        self.assertTrue(len(res_v) == ndim)
        # Test save/load
        fname = 'test_saving'
        try:
            skle.save(fname)
            g2v_l = nodevectors.SKLearnEmbedder.load(fname + '.zip')
            res_l = g2v_l.predict(9)
            self.assertTrue(len(res_l) == ndim)
            np.testing.assert_array_almost_equal(res_l, res_v)
        finally:
            os.remove(fname + '.zip')


    def test_grarep(self):
        tt = nx.generators.complete_graph(25)
        ndim = 3
        skle = nodevectors.GraRep(n_components=ndim)
        skle.fit(tt)
        res_v = skle.predict(9)
        self.assertTrue(len(res_v) == ndim)
        # Test save/load
        fname = 'test_saving'
        try:
            skle.save(fname)
            g2v_l = nodevectors.SKLearnEmbedder.load(fname + '.zip')
            res_l = g2v_l.predict(9)
            self.assertTrue(len(res_l) == ndim)
            np.testing.assert_array_almost_equal(res_l, res_v)
        finally:
            os.remove(fname + '.zip')


    def test_node2vec_factored_names(self):
        tt = cg.read_edgelist("./tests/unfactored_edgelist.csv", sep=",")
        ndim = 3
        skle = nodevectors.Node2Vec(
            walklen=5, 
            epochs=5,
            threads=1,
            n_components=ndim,
            keep_walks=True,
            verbose=False,
            w2vparams={"window":3, "negative":3, "iter":3,
                       "batch_words":32, "workers": 2})
        skle.fit(tt)
        res_v = skle.predict(9)
        self.assertTrue(len(res_v) == ndim)
        # Test save/load
        fname = 'test_saving'
        try:
            skle.save(fname)
            g2v_l = nodevectors.SKLearnEmbedder.load(fname + '.zip')
            res_l = g2v_l.predict(9)
            self.assertTrue(len(res_l) == ndim)
            np.testing.assert_array_almost_equal(res_l, res_v)
        finally:
            os.remove(fname + '.zip')


    def test_node2vec_fit_transform(self):
        tt = cg.read_edgelist("./tests/unfactored_edgelist.csv", sep=",")
        ndim = 3
        skle = nodevectors.Node2Vec(
            walklen=5, 
            epochs=5,
            threads=1,
            n_components=ndim,
            keep_walks=True,
            verbose=False,
            w2vparams={"window":3, "negative":3, "iter":3,
                       "batch_words":32, "workers": 2})
        skle.fit_transform(tt)
