import backend as F
import numpy as np
import scipy as sp
import dgl
from dgl import utils
import unittest
from numpy.testing import assert_array_equal

np.random.seed(42)

def generate_rand_graph(n):
    arr = (sp.sparse.random(n, n, density=0.1, format='coo') != 0).astype(np.int64)
    return dgl.DGLGraph(arr, readonly=True)

def test_create_full():
    g = generate_rand_graph(100)
    full_nf = dgl.contrib.sampling.sampler.create_full_nodeflow(g, 5)
    assert full_nf.number_of_nodes() == g.number_of_nodes() * 6
    assert full_nf.number_of_edges() == g.number_of_edges() * 5

def test_1neighbor_sampler_all():
    g = generate_rand_graph(100)
    # In this case, NeighborSampling simply gets the neighborhood of a single vertex.
    for i, subg in enumerate(dgl.contrib.sampling.NeighborSampler(
            g, 1, g.number_of_nodes(), neighbor_type='in', num_workers=4)):
        seed_ids = subg.layer_parent_nid(-1)
        assert len(seed_ids) == 1
        src, dst, eid = g.in_edges(seed_ids, form='all')
        assert subg.number_of_nodes() == len(src) + 1
        assert subg.number_of_edges() == len(src)

        assert seed_ids == subg.layer_parent_nid(-1)
        child_src, child_dst, child_eid = subg.in_edges(subg.layer_nid(-1), form='all')
        assert F.array_equal(child_src, subg.layer_nid(0))

        src1 = subg.map_to_parent_nid(child_src)
        assert F.array_equal(src1, src)

def is_sorted(arr):
    return np.sum(np.sort(arr) == arr, 0) == len(arr)

def verify_subgraph(g, subg, seed_id):
    seed_id = F.asnumpy(seed_id)
    seeds = F.asnumpy(subg.map_to_parent_nid(subg.layer_nid(-1)))
    assert seed_id in seeds
    child_seed = F.asnumpy(subg.layer_nid(-1))[seeds == seed_id]
    src, dst, eid = g.in_edges(seed_id, form='all')
    child_src, child_dst, child_eid = subg.in_edges(child_seed, form='all')

    child_src = F.asnumpy(child_src)
    # We don't allow duplicate elements in the neighbor list.
    assert(len(np.unique(child_src)) == len(child_src))
    # The neighbor list also needs to be sorted.
    assert(is_sorted(child_src))

    # a neighbor in the subgraph must also exist in parent graph.
    src = F.asnumpy(src)
    for i in subg.map_to_parent_nid(child_src):
        assert F.asnumpy(i) in src

def test_1neighbor_sampler():
    g = generate_rand_graph(100)
    # In this case, NeighborSampling simply gets the neighborhood of a single vertex.
    for subg in dgl.contrib.sampling.NeighborSampler(g, 1, 5, neighbor_type='in',
                                                     num_workers=4):
        seed_ids = subg.layer_parent_nid(-1)
        assert len(seed_ids) == 1
        assert subg.number_of_nodes() <= 6
        assert subg.number_of_edges() <= 5
        verify_subgraph(g, subg, seed_ids)

def test_prefetch_neighbor_sampler():
    g = generate_rand_graph(100)
    # In this case, NeighborSampling simply gets the neighborhood of a single vertex.
    for subg in dgl.contrib.sampling.NeighborSampler(g, 1, 5, neighbor_type='in',
                                                     num_workers=4, prefetch=True):
        seed_ids = subg.layer_parent_nid(-1)
        assert len(seed_ids) == 1
        assert subg.number_of_nodes() <= 6
        assert subg.number_of_edges() <= 5
        verify_subgraph(g, subg, seed_ids)

def test_10neighbor_sampler_all():
    g = generate_rand_graph(100)
    # In this case, NeighborSampling simply gets the neighborhood of a single vertex.
    for subg in dgl.contrib.sampling.NeighborSampler(g, 10, g.number_of_nodes(),
                                                     neighbor_type='in', num_workers=4):
        seed_ids = subg.layer_parent_nid(-1)
        assert F.array_equal(seed_ids, subg.map_to_parent_nid(subg.layer_nid(-1)))

        src, dst, eid = g.in_edges(seed_ids, form='all')
        child_src, child_dst, child_eid = subg.in_edges(subg.layer_nid(-1), form='all')
        src1 = subg.map_to_parent_nid(child_src)
        assert F.array_equal(src1, src)

def check_10neighbor_sampler(g, seeds):
    # In this case, NeighborSampling simply gets the neighborhood of a single vertex.
    for subg in dgl.contrib.sampling.NeighborSampler(g, 10, 5, neighbor_type='in',
                                                     num_workers=4, seed_nodes=seeds):
        seed_ids = subg.layer_parent_nid(-1)
        assert subg.number_of_nodes() <= 6 * len(seed_ids)
        assert subg.number_of_edges() <= 5 * len(seed_ids)
        for seed_id in seed_ids:
            verify_subgraph(g, subg, seed_id)

def test_10neighbor_sampler():
    g = generate_rand_graph(100)
    check_10neighbor_sampler(g, None)
    check_10neighbor_sampler(g, seeds=np.unique(np.random.randint(0, g.number_of_nodes(),
                                                                  size=int(g.number_of_nodes() / 10))))

def _test_layer_sampler(prefetch=False):
    g = generate_rand_graph(100)
    nid = g.nodes()
    src, dst, eid = g.all_edges(form='all', order='eid')
    n_batches = 5
    batch_size = 50
    seed_batches = [np.sort(np.random.choice(F.asnumpy(nid), batch_size, replace=False))
                    for i in range(n_batches)]
    seed_nodes = np.hstack(seed_batches)
    layer_sizes = [50] * 3
    LayerSampler = getattr(dgl.contrib.sampling, 'LayerSampler')
    sampler = LayerSampler(g, batch_size, layer_sizes, 'in',
                           seed_nodes=seed_nodes, num_workers=4, prefetch=prefetch)
    for sub_g in sampler:
        assert all(sub_g.layer_size(i) < size for i, size in enumerate(layer_sizes))
        sub_nid = F.arange(0, sub_g.number_of_nodes())
        assert all(np.all(np.isin(F.asnumpy(sub_g.layer_nid(i)), F.asnumpy(sub_nid)))
                   for i in range(sub_g.num_layers))
        assert np.all(np.isin(F.asnumpy(sub_g.map_to_parent_nid(sub_nid)),
                              F.asnumpy(nid)))
        sub_eid = F.arange(0, sub_g.number_of_edges())
        assert np.all(np.isin(F.asnumpy(sub_g.map_to_parent_eid(sub_eid)),
                              F.asnumpy(eid)))
        assert any(np.all(np.sort(F.asnumpy(sub_g.layer_parent_nid(-1))) == seed_batch)
                   for seed_batch in seed_batches)

        sub_src, sub_dst = sub_g.all_edges(order='eid')
        for i in range(sub_g.num_blocks):
            block_eid = sub_g.block_eid(i)
            block_src = sub_g.map_to_parent_nid(F.gather_row(sub_src, block_eid))
            block_dst = sub_g.map_to_parent_nid(F.gather_row(sub_dst, block_eid))

            block_parent_eid = sub_g.block_parent_eid(i)
            block_parent_src = F.gather_row(src, block_parent_eid)
            block_parent_dst = F.gather_row(dst, block_parent_eid)

            assert np.all(F.asnumpy(block_src == block_parent_src))

        n_layers = sub_g.num_layers
        sub_n = sub_g.number_of_nodes()
        assert sum(F.shape(sub_g.layer_nid(i))[0] for i in range(n_layers)) == sub_n
        n_blocks = sub_g.num_blocks
        sub_m = sub_g.number_of_edges()
        assert sum(F.shape(sub_g.block_eid(i))[0] for i in range(n_blocks)) == sub_m

def test_layer_sampler():
    _test_layer_sampler()
    _test_layer_sampler(prefetch=True)

@unittest.skipIf(dgl.backend.backend_name == "tensorflow", reason="Error occured when multiprocessing")
def test_nonuniform_neighbor_sampler():
    # Construct a graph with
    # (1) A path (0, 1, ..., 99) with weight 1
    # (2) A bunch of random edges with weight 0.
    edges = []
    for i in range(99):
        edges.append((i, i + 1))
    for i in range(1000):
        edge = (np.random.randint(100), np.random.randint(100))
        if edge not in edges:
            edges.append(edge)
    src, dst = zip(*edges)
    g = dgl.DGLGraph()
    g.add_nodes(100)
    g.add_edges(src, dst)
    g.readonly()

    g.edata['w'] = F.cat([
        F.ones((99,), F.float64, F.cpu()),
        F.zeros((len(edges) - 99,), F.float64, F.cpu())], 0)

    # Test 1-neighbor NodeFlow with 99 as target node.
    # The generated NodeFlow should only contain node i on layer i.
    sampler = dgl.contrib.sampling.NeighborSampler(
        g, 1, 1, 99, 'in', transition_prob='w', seed_nodes=[99])
    nf = next(iter(sampler))

    assert nf.num_layers == 100
    for i in range(nf.num_layers):
        assert nf.layer_size(i) == 1
        assert F.asnumpy(nf.layer_parent_nid(i)[0]) == i

    # Test the reverse direction
    sampler = dgl.contrib.sampling.NeighborSampler(
        g, 1, 1, 99, 'out', transition_prob='w', seed_nodes=[0])
    nf = next(iter(sampler))

    assert nf.num_layers == 100
    for i in range(nf.num_layers):
        assert nf.layer_size(i) == 1
        assert F.asnumpy(nf.layer_parent_nid(i)[0]) == 99 - i

def test_setseed():
    g = generate_rand_graph(100)

    nids = []

    dgl.random.seed(42)
    for subg in dgl.contrib.sampling.NeighborSampler(
            g, 5, 3, num_hops=2, neighbor_type='in', num_workers=1):
        nids.append(
            tuple(tuple(F.asnumpy(subg.layer_parent_nid(i))) for i in range(3)))

    # reinitialize
    dgl.random.seed(42)
    for i, subg in enumerate(dgl.contrib.sampling.NeighborSampler(
            g, 5, 3, num_hops=2, neighbor_type='in', num_workers=1)):
        item = tuple(tuple(F.asnumpy(subg.layer_parent_nid(i))) for i in range(3))
        assert item == nids[i]

    for i, subg in enumerate(dgl.contrib.sampling.NeighborSampler(
            g, 5, 3, num_hops=2, neighbor_type='in', num_workers=4)):
        pass

def check_head_tail(g):
    lsrc, ldst, leid = g.all_edges(form='all', order='eid')

    lsrc = np.unique(F.asnumpy(lsrc))
    head_nid = np.unique(F.asnumpy(g.head_nid))
    assert len(head_nid) == len(g.head_nid)
    np.testing.assert_equal(lsrc, head_nid)

    ldst = np.unique(F.asnumpy(ldst))
    tail_nid = np.unique(F.asnumpy(g.tail_nid))
    assert len(tail_nid) == len(g.tail_nid)
    np.testing.assert_equal(tail_nid, ldst)


def check_negative_sampler(mode, exclude_positive, neg_size):
    g = generate_rand_graph(100)
    num_edges = g.number_of_edges()
    etype = np.random.randint(0, 10, size=g.number_of_edges(), dtype=np.int64)
    g.edata['etype'] = F.copy_to(F.tensor(etype), F.cpu())

    pos_gsrc, pos_gdst, pos_geid = g.all_edges(form='all', order='eid')
    pos_map = {}
    for i in range(len(pos_geid)):
        pos_d = int(F.asnumpy(pos_gdst[i]))
        pos_e = int(F.asnumpy(pos_geid[i]))
        pos_map[(pos_d, pos_e)] = int(F.asnumpy(pos_gsrc[i]))

    EdgeSampler = getattr(dgl.contrib.sampling, 'EdgeSampler')
    # Test the homogeneous graph.

    batch_size = 50
    total_samples = 0
    for pos_edges, neg_edges in EdgeSampler(g, batch_size,
                                            negative_mode=mode,
                                            reset=False,
                                            neg_sample_size=neg_size,
                                            exclude_positive=exclude_positive,
                                            return_false_neg=True):
        pos_lsrc, pos_ldst, pos_leid = pos_edges.all_edges(form='all', order='eid')
        assert_array_equal(F.asnumpy(F.gather_row(pos_edges.parent_eid, pos_leid)),
                           F.asnumpy(g.edge_ids(F.gather_row(pos_edges.parent_nid, pos_lsrc),
                                                F.gather_row(pos_edges.parent_nid, pos_ldst))))

        neg_lsrc, neg_ldst, neg_leid = neg_edges.all_edges(form='all', order='eid')

        neg_src = F.gather_row(neg_edges.parent_nid, neg_lsrc)
        neg_dst = F.gather_row(neg_edges.parent_nid, neg_ldst)
        neg_eid = F.gather_row(neg_edges.parent_eid, neg_leid)
        for i in range(len(neg_eid)):
            neg_d = int(F.asnumpy(neg_dst)[i])
            neg_e = int(F.asnumpy(neg_eid)[i])
            assert (neg_d, neg_e) in pos_map
            if exclude_positive:
                assert int(F.asnumpy(neg_src[i])) != pos_map[(neg_d, neg_e)]

        check_head_tail(neg_edges)
        pos_tails = F.gather_row(pos_edges.parent_nid, pos_edges.tail_nid)
        neg_tails = F.gather_row(neg_edges.parent_nid, neg_edges.tail_nid)
        pos_tails = np.sort(F.asnumpy(pos_tails))
        neg_tails = np.sort(F.asnumpy(neg_tails))
        np.testing.assert_equal(pos_tails, neg_tails)

        exist = neg_edges.edata['false_neg']
        if exclude_positive:
            assert np.sum(F.asnumpy(exist) == 0) == len(exist)
        else:
            assert F.array_equal(g.has_edges_between(neg_src, neg_dst), exist)
        total_samples += batch_size
    assert total_samples <= num_edges

    # check replacement = True
    # with reset = False (default setting)
    total_samples = 0
    for pos_edges, neg_edges in EdgeSampler(g, batch_size,
                                            replacement=True,
                                            reset=False,
                                            negative_mode=mode,
                                            neg_sample_size=neg_size,
                                            exclude_positive=exclude_positive,
                                            return_false_neg=True):
        _, _, pos_leid = pos_edges.all_edges(form='all', order='eid')
        assert len(pos_leid) == batch_size
        total_samples += len(pos_leid)
    assert total_samples == num_edges

    # check replacement = False
    # with reset = False (default setting)
    total_samples = 0
    for pos_edges, neg_edges in EdgeSampler(g, batch_size,
                                            replacement=False,
                                            reset=False,
                                            negative_mode=mode,
                                            neg_sample_size=neg_size,
                                            exclude_positive=exclude_positive,
                                            return_false_neg=True):
        _, _, pos_leid = pos_edges.all_edges(form='all', order='eid')
        assert len(pos_leid) == batch_size
        total_samples += len(pos_leid)
    assert total_samples == num_edges

    # check replacement = True
    # with reset = True
    total_samples = 0
    max_samples = 2 * num_edges
    for pos_edges, neg_edges in EdgeSampler(g, batch_size,
                                            replacement=True,
                                            reset=True,
                                            negative_mode=mode,
                                            neg_sample_size=neg_size,
                                            exclude_positive=exclude_positive,
                                            return_false_neg=True):
        _, _, pos_leid = pos_edges.all_edges(form='all', order='eid')
        assert len(pos_leid) <= batch_size
        total_samples += len(pos_leid)
        if (total_samples >= max_samples):
            break
    assert total_samples >= max_samples

    # check replacement = False
    # with reset = True
    total_samples = 0
    max_samples = 2 * num_edges
    for pos_edges, neg_edges in EdgeSampler(g, batch_size,
                                            replacement=False,
                                            reset=True,
                                            negative_mode=mode,
                                            neg_sample_size=neg_size,
                                            exclude_positive=exclude_positive,
                                            return_false_neg=True):
        _, _, pos_leid = pos_edges.all_edges(form='all', order='eid')
        assert len(pos_leid) <= batch_size
        total_samples += len(pos_leid)
        if (total_samples >= max_samples):
            break
    assert total_samples >= max_samples

    # Test the knowledge graph.
    total_samples = 0
    for _, neg_edges in EdgeSampler(g, batch_size,
                                    negative_mode=mode,
                                    reset=False,
                                    neg_sample_size=neg_size,
                                    exclude_positive=exclude_positive,
                                    relations=g.edata['etype'],
                                    return_false_neg=True):
        neg_lsrc, neg_ldst, neg_leid = neg_edges.all_edges(form='all', order='eid')
        neg_src = F.gather_row(neg_edges.parent_nid, neg_lsrc)
        neg_dst = F.gather_row(neg_edges.parent_nid, neg_ldst)
        neg_eid = F.gather_row(neg_edges.parent_eid, neg_leid)
        exists = neg_edges.edata['false_neg']
        neg_edges.edata['etype'] = F.gather_row(g.edata['etype'], neg_eid)
        for i in range(len(neg_eid)):
            u, v = F.asnumpy(neg_src[i]), F.asnumpy(neg_dst[i])
            if g.has_edge_between(u, v):
                eid = g.edge_id(u, v)
                etype = g.edata['etype'][eid]
                exist = neg_edges.edata['etype'][i] == etype
                assert F.asnumpy(exists[i]) == F.asnumpy(exist)
        total_samples += batch_size
    assert total_samples <= num_edges

def check_weighted_negative_sampler(mode, exclude_positive, neg_size):
    g = generate_rand_graph(100)
    num_edges = g.number_of_edges()
    num_nodes = g.number_of_nodes()
    edge_weight = F.copy_to(F.tensor(np.full((num_edges,), 1, dtype=np.float32)), F.cpu())
    node_weight = F.copy_to(F.tensor(np.full((num_nodes,), 1, dtype=np.float32)), F.cpu())
    etype = np.random.randint(0, 10, size=num_edges, dtype=np.int64)
    g.edata['etype'] = F.copy_to(F.tensor(etype), F.cpu())

    pos_gsrc, pos_gdst, pos_geid = g.all_edges(form='all', order='eid')
    pos_map = {}
    for i in range(len(pos_geid)):
        pos_d = int(F.asnumpy(pos_gdst[i]))
        pos_e = int(F.asnumpy(pos_geid[i]))
        pos_map[(pos_d, pos_e)] = int(F.asnumpy(pos_gsrc[i]))
    EdgeSampler = getattr(dgl.contrib.sampling, 'EdgeSampler')

    # Correctness check
    # Test the homogeneous graph.
    batch_size = 50
    # Test the knowledge graph with edge weight provied.
    total_samples = 0
    for pos_edges, neg_edges in EdgeSampler(g, batch_size,
                                            reset=False,
                                            edge_weight=edge_weight,
                                            negative_mode=mode,
                                            neg_sample_size=neg_size,
                                            exclude_positive=exclude_positive,
                                            return_false_neg=True):
        pos_lsrc, pos_ldst, pos_leid = pos_edges.all_edges(form='all', order='eid')
        assert_array_equal(F.asnumpy(F.gather_row(pos_edges.parent_eid, pos_leid)),
                           F.asnumpy(g.edge_ids(F.gather_row(pos_edges.parent_nid, pos_lsrc),
                                                F.gather_row(pos_edges.parent_nid, pos_ldst))))
        neg_lsrc, neg_ldst, neg_leid = neg_edges.all_edges(form='all', order='eid')

        neg_src = F.gather_row(neg_edges.parent_nid, neg_lsrc)
        neg_dst = F.gather_row(neg_edges.parent_nid, neg_ldst)
        neg_eid = F.gather_row(neg_edges.parent_eid, neg_leid)
        for i in range(len(neg_eid)):
            neg_d = int(F.asnumpy(neg_dst[i]))
            neg_e = int(F.asnumpy(neg_eid[i]))
            assert (neg_d, neg_e) in pos_map
            if exclude_positive:
                assert int(F.asnumpy(neg_src[i])) != pos_map[(neg_d, neg_e)]

        check_head_tail(neg_edges)
        pos_tails = F.gather_row(pos_edges.parent_nid, pos_edges.tail_nid)
        neg_tails = F.gather_row(neg_edges.parent_nid, neg_edges.tail_nid)
        pos_tails = np.sort(F.asnumpy(pos_tails))
        neg_tails = np.sort(F.asnumpy(neg_tails))
        np.testing.assert_equal(pos_tails, neg_tails)

        exist = neg_edges.edata['false_neg']
        if exclude_positive:
            assert np.sum(F.asnumpy(exist) == 0) == len(exist)
        else:
            assert F.array_equal(g.has_edges_between(neg_src, neg_dst), exist)
        total_samples += batch_size
    assert total_samples <= num_edges

    # Test the knowledge graph with edge weight provied.
    total_samples = 0
    for pos_edges, neg_edges in EdgeSampler(g, batch_size,
                                            reset=False,
                                            edge_weight=edge_weight,
                                            negative_mode=mode,
                                            neg_sample_size=neg_size,
                                            exclude_positive=exclude_positive,
                                            relations=g.edata['etype'],
                                            return_false_neg=True):
        neg_lsrc, neg_ldst, neg_leid = neg_edges.all_edges(form='all', order='eid')
        neg_src = F.gather_row(neg_edges.parent_nid, neg_lsrc)
        neg_dst = F.gather_row(neg_edges.parent_nid, neg_ldst)
        neg_eid = F.gather_row(neg_edges.parent_eid, neg_leid)
        exists = neg_edges.edata['false_neg']
        neg_edges.edata['etype'] = F.gather_row(g.edata['etype'], neg_eid)
        for i in range(len(neg_eid)):
            u, v = F.asnumpy(neg_src[i]), F.asnumpy(neg_dst[i])
            if g.has_edge_between(u, v):
                eid = g.edge_id(u, v)
                etype = g.edata['etype'][eid]
                exist = neg_edges.edata['etype'][i] == etype
                assert F.asnumpy(exists[i]) == F.asnumpy(exist)
        total_samples += batch_size
    assert total_samples <= num_edges

    # Test the knowledge graph with edge/node weight provied.
    total_samples = 0
    for pos_edges, neg_edges in EdgeSampler(g, batch_size,
                                            reset=False,
                                            edge_weight=edge_weight,
                                            node_weight=node_weight,
                                            negative_mode=mode,
                                            neg_sample_size=neg_size,
                                            exclude_positive=exclude_positive,
                                            relations=g.edata['etype'],
                                            return_false_neg=True):
        neg_lsrc, neg_ldst, neg_leid = neg_edges.all_edges(form='all', order='eid')
        neg_src = F.gather_row(neg_edges.parent_nid, neg_lsrc)
        neg_dst = F.gather_row(neg_edges.parent_nid, neg_ldst)
        neg_eid = F.gather_row(neg_edges.parent_eid, neg_leid)
        exists = neg_edges.edata['false_neg']
        neg_edges.edata['etype'] = F.gather_row(g.edata['etype'], neg_eid)
        for i in range(len(neg_eid)):
            u, v = F.asnumpy(neg_src[i]), F.asnumpy(neg_dst[i])
            if g.has_edge_between(u, v):
                eid = g.edge_id(u, v)
                etype = g.edata['etype'][eid]
                exist = neg_edges.edata['etype'][i] == etype
                assert F.asnumpy(exists[i]) == F.asnumpy(exist)
        total_samples += batch_size
    assert total_samples <= num_edges

    # check replacement = True with pos edges no-uniform sample
    # with reset = False
    total_samples = 0
    for pos_edges, neg_edges in EdgeSampler(g, batch_size,
                                            replacement=True,
                                            reset=False,
                                            edge_weight=edge_weight,
                                            negative_mode=mode,
                                            neg_sample_size=neg_size,
                                            exclude_positive=exclude_positive,
                                            return_false_neg=True):
        _, _, pos_leid = pos_edges.all_edges(form='all', order='eid')
        assert len(pos_leid) == batch_size
        total_samples += len(pos_leid)
    assert total_samples == num_edges

    # check replacement = True with pos edges no-uniform sample
    # with reset = True
    total_samples = 0
    max_samples = 4 * num_edges
    for pos_edges, neg_edges in EdgeSampler(g, batch_size,
                                            replacement=True,
                                            reset=True,
                                            edge_weight=edge_weight,
                                            negative_mode=mode,
                                            neg_sample_size=neg_size,
                                            exclude_positive=exclude_positive,
                                            return_false_neg=True):
        _, _, pos_leid = pos_edges.all_edges(form='all', order='eid')
        assert len(pos_leid) == batch_size
        total_samples += len(pos_leid)
        if total_samples >= max_samples:
            break
    assert total_samples == max_samples

    # check replacement = False with pos/neg edges no-uniform sample
    # reset = False
    total_samples = 0
    for pos_edges, neg_edges in EdgeSampler(g, batch_size,
                                            replacement=False,
                                            reset=False,
                                            edge_weight=edge_weight,
                                            node_weight=node_weight,
                                            negative_mode=mode,
                                            neg_sample_size=neg_size,
                                            exclude_positive=exclude_positive,
                                            relations=g.edata['etype'],
                                            return_false_neg=True):
        _, _, pos_leid = pos_edges.all_edges(form='all', order='eid')
        assert len(pos_leid) == batch_size
        total_samples += len(pos_leid)
    assert total_samples == num_edges

    # check replacement = False with pos/neg edges no-uniform sample
    # reset = True
    total_samples = 0
    for pos_edges, neg_edges in EdgeSampler(g, batch_size,
                                            replacement=False,
                                            reset=True,
                                            edge_weight=edge_weight,
                                            node_weight=node_weight,
                                            negative_mode=mode,
                                            neg_sample_size=neg_size,
                                            exclude_positive=exclude_positive,
                                            relations=g.edata['etype'],
                                            return_false_neg=True):
        _, _, pos_leid = pos_edges.all_edges(form='all', order='eid')
        assert len(pos_leid) == batch_size
        total_samples += len(pos_leid)
        if total_samples >= max_samples:
            break
    assert total_samples == max_samples

    # Check Rate
    dgl.random.seed(0)
    g = generate_rand_graph(1000)
    num_edges = g.number_of_edges()
    num_nodes = g.number_of_nodes()
    edge_weight = F.copy_to(F.tensor(np.full((num_edges,), 1, dtype=np.float32)), F.cpu())
    edge_weight[0] = F.sum(edge_weight, dim=0)
    node_weight = F.copy_to(F.tensor(np.full((num_nodes,), 1, dtype=np.float32)), F.cpu())
    node_weight[-1] = F.sum(node_weight, dim=0) / 200
    etype = np.random.randint(0, 20, size=num_edges, dtype=np.int64)
    g.edata['etype'] = F.copy_to(F.tensor(etype), F.cpu())

    # Test w/o node weight.
    max_samples = num_edges // 5
    total_samples = 0
    # Test the knowledge graph with edge weight provied.
    edge_sampled = np.full((num_edges,), 0, dtype=np.int32)
    node_sampled = np.full((num_nodes,), 0, dtype=np.int32)
    for pos_edges, neg_edges in EdgeSampler(g, batch_size,
                                            replacement=True,
                                            edge_weight=edge_weight,
                                            shuffle=True,
                                            negative_mode=mode,
                                            neg_sample_size=neg_size,
                                            exclude_positive=False,
                                            relations=g.edata['etype'],
                                            return_false_neg=True):
        _, _, pos_leid = pos_edges.all_edges(form='all', order='eid')
        neg_lsrc, neg_ldst, _ = neg_edges.all_edges(form='all', order='eid')
        if 'head' in mode:
            neg_src = neg_edges.parent_nid[neg_lsrc]
            np.add.at(node_sampled, F.asnumpy(neg_src), 1)
        else:
            neg_dst = neg_edges.parent_nid[neg_ldst]
            np.add.at(node_sampled, F.asnumpy(neg_dst), 1)
        np.add.at(edge_sampled, F.asnumpy(pos_edges.parent_eid[pos_leid]), 1)
        total_samples += batch_size
        if total_samples > max_samples:
            break

    # Check rate here
    edge_rate_0 = edge_sampled[0] / edge_sampled.sum()
    edge_tail_half_cnt = edge_sampled[edge_sampled.shape[0] // 2:-1].sum()
    edge_rate_tail_half = edge_tail_half_cnt / edge_sampled.sum()
    assert np.allclose(edge_rate_0, 0.5, atol=0.05)
    assert np.allclose(edge_rate_tail_half, 0.25, atol=0.05)

    node_rate_0 = node_sampled[0] / node_sampled.sum()
    node_tail_half_cnt = node_sampled[node_sampled.shape[0] // 2:-1].sum()
    node_rate_tail_half = node_tail_half_cnt / node_sampled.sum()
    assert node_rate_0 < 0.02
    assert np.allclose(node_rate_tail_half, 0.5, atol=0.02)

    # Test the knowledge graph with edge/node weight provied.
    edge_sampled = np.full((num_edges,), 0, dtype=np.int32)
    node_sampled = np.full((num_nodes,), 0, dtype=np.int32)
    total_samples = 0
    for pos_edges, neg_edges in EdgeSampler(g, batch_size,
                                            replacement=True,
                                            edge_weight=edge_weight,
                                            node_weight=node_weight,
                                            shuffle=True,
                                            negative_mode=mode,
                                            neg_sample_size=neg_size,
                                            exclude_positive=False,
                                            relations=g.edata['etype'],
                                            return_false_neg=True):
        _, _, pos_leid = pos_edges.all_edges(form='all', order='eid')
        neg_lsrc, neg_ldst, _ = neg_edges.all_edges(form='all', order='eid')
        if 'head' in mode:
            neg_src = F.gather_row(neg_edges.parent_nid, neg_lsrc)
            np.add.at(node_sampled, F.asnumpy(neg_src), 1)
        else:
            neg_dst = F.gather_row(neg_edges.parent_nid, neg_ldst)
            np.add.at(node_sampled, F.asnumpy(neg_dst), 1)
        np.add.at(edge_sampled, F.asnumpy(pos_edges.parent_eid[pos_leid]), 1)
        total_samples += batch_size
        if total_samples > max_samples:
            break

    # Check rate here
    edge_rate_0 = edge_sampled[0] / edge_sampled.sum()
    edge_tail_half_cnt = edge_sampled[edge_sampled.shape[0] // 2:-1].sum()
    edge_rate_tail_half = edge_tail_half_cnt / edge_sampled.sum()
    assert np.allclose(edge_rate_0, 0.5, atol=0.05)
    assert np.allclose(edge_rate_tail_half, 0.25, atol=0.05)

    node_rate = node_sampled[-1] / node_sampled.sum()
    node_rate_a = np.average(node_sampled[:50]) / node_sampled.sum()
    node_rate_b = np.average(node_sampled[50:100]) / node_sampled.sum()
    # As neg sampling does not contain duplicate nodes,
    # this test takes some acceptable variation on the sample rate.
    assert np.allclose(node_rate, node_rate_a * 5, atol=0.002)
    assert np.allclose(node_rate_a, node_rate_b, atol=0.0002)

def check_positive_edge_sampler():
    g = generate_rand_graph(1000)
    num_edges = g.number_of_edges()
    edge_weight = F.copy_to(F.tensor(np.full((num_edges,), 1, dtype=np.float32)), F.cpu())

    edge_weight[num_edges-1] = num_edges ** 3
    EdgeSampler = getattr(dgl.contrib.sampling, 'EdgeSampler')

    # Correctness check
    # Test the homogeneous graph.
    batch_size = 128
    edge_sampled = np.full((num_edges,), 0, dtype=np.int32)
    for pos_edges in EdgeSampler(g, batch_size,
                                    reset=False,
                                    edge_weight=edge_weight):
        _, _, pos_leid = pos_edges.all_edges(form='all', order='eid')
        np.add.at(edge_sampled, F.asnumpy(pos_edges.parent_eid[pos_leid]), 1)
    truth = np.full((num_edges,), 1, dtype=np.int32)
    edge_sampled = edge_sampled[:num_edges]
    assert np.array_equal(truth, edge_sampled)

    edge_sampled = np.full((num_edges,), 0, dtype=np.int32)
    for pos_edges in EdgeSampler(g, batch_size,
                                    reset=False,
                                    shuffle=True,
                                    edge_weight=edge_weight):
        _, _, pos_leid = pos_edges.all_edges(form='all', order='eid')
        np.add.at(edge_sampled, F.asnumpy(pos_edges.parent_eid[pos_leid]), 1)
    truth = np.full((num_edges,), 1, dtype=np.int32)
    edge_sampled = edge_sampled[:num_edges]
    assert np.array_equal(truth, edge_sampled)


@unittest.skipIf(dgl.backend.backend_name == "tensorflow", reason="TF doesn't support item assignment")
def test_negative_sampler():
    check_negative_sampler('chunk-head', False, 10)
    check_negative_sampler('head', True, 10)
    check_negative_sampler('head', False, 10)
    check_weighted_negative_sampler('chunk-head', False, 10)
    check_weighted_negative_sampler('head', True, 10)
    check_weighted_negative_sampler('head', False, 10)
    check_positive_edge_sampler()
    #disable this check for now. It might take too long time.
    #check_negative_sampler('head', False, 100)


if __name__ == '__main__':
    test_create_full()
    test_1neighbor_sampler_all()
    test_10neighbor_sampler_all()
    test_1neighbor_sampler()
    test_10neighbor_sampler()
    test_layer_sampler()
    test_nonuniform_neighbor_sampler()
    test_setseed()
    test_negative_sampler()

[(9, 'numpy.random.seed', 'np.random.seed', True, 'import numpy as np\n'), (161, 'unittest.skipIf', 'unittest.skipIf', False, 'import unittest\n'), (701, 'unittest.skipIf', 'unittest.skipIf', False, 'import unittest\n'), (13, 'dgl.DGLGraph', 'dgl.DGLGraph', False, 'import dgl\n'), (17, 'dgl.contrib.sampling.sampler.create_full_nodeflow', 'dgl.contrib.sampling.sampler.create_full_nodeflow', False, 'import dgl\n'), (43, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (50, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (57, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (64, 'dgl.contrib.sampling.NeighborSampler', 'dgl.contrib.sampling.NeighborSampler', False, 'import dgl\n'), (75, 'dgl.contrib.sampling.NeighborSampler', 'dgl.contrib.sampling.NeighborSampler', False, 'import dgl\n'), (98, 'dgl.contrib.sampling.NeighborSampler', 'dgl.contrib.sampling.NeighborSampler', False, 'import dgl\n'), (120, 'numpy.hstack', 'np.hstack', True, 'import numpy as np\n'), (174, 'dgl.DGLGraph', 'dgl.DGLGraph', False, 'import dgl\n'), (185, 'dgl.contrib.sampling.NeighborSampler', 'dgl.contrib.sampling.NeighborSampler', False, 'import dgl\n'), (195, 'dgl.contrib.sampling.NeighborSampler', 'dgl.contrib.sampling.NeighborSampler', False, 'import dgl\n'), (209, 'dgl.random.seed', 'dgl.random.seed', False, 'import dgl\n'), (210, 'dgl.contrib.sampling.NeighborSampler', 'dgl.contrib.sampling.NeighborSampler', False, 'import dgl\n'), (216, 'dgl.random.seed', 'dgl.random.seed', False, 'import dgl\n'), (232, 'numpy.testing.assert_equal', 'np.testing.assert_equal', True, 'import numpy as np\n'), (237, 'numpy.testing.assert_equal', 'np.testing.assert_equal', True, 'import numpy as np\n'), (393, 'numpy.random.randint', 'np.random.randint', True, 'import numpy as np\n'), (574, 'dgl.random.seed', 'dgl.random.seed', False, 'import dgl\n'), (579, 'backend.sum', 'F.sum', True, 'import backend as F\n'), (582, 'numpy.random.randint', 'np.random.randint', True, 'import numpy as np\n'), (589, 'numpy.full', 'np.full', True, 'import numpy as np\n'), (590, 'numpy.full', 'np.full', True, 'import numpy as np\n'), (617, 'numpy.allclose', 'np.allclose', True, 'import numpy as np\n'), (618, 'numpy.allclose', 'np.allclose', True, 'import numpy as np\n'), (624, 'numpy.allclose', 'np.allclose', True, 'import numpy as np\n'), (627, 'numpy.full', 'np.full', True, 'import numpy as np\n'), (628, 'numpy.full', 'np.full', True, 'import numpy as np\n'), (657, 'numpy.allclose', 'np.allclose', True, 'import numpy as np\n'), (658, 'numpy.allclose', 'np.allclose', True, 'import numpy as np\n'), (665, 'numpy.allclose', 'np.allclose', True, 'import numpy as np\n'), (666, 'numpy.allclose', 'np.allclose', True, 'import numpy as np\n'), (679, 'numpy.full', 'np.full', True, 'import numpy as np\n'), (685, 'numpy.full', 'np.full', True, 'import numpy as np\n'), (687, 'numpy.array_equal', 'np.array_equal', True, 'import numpy as np\n'), (689, 'numpy.full', 'np.full', True, 'import numpy as np\n'), (696, 'numpy.full', 'np.full', True, 'import numpy as np\n'), (698, 'numpy.array_equal', 'np.array_equal', True, 'import numpy as np\n'), (37, 'backend.array_equal', 'F.array_equal', True, 'import backend as F\n'), (94, 'backend.array_equal', 'F.array_equal', True, 'import backend as F\n'), (217, 'dgl.contrib.sampling.NeighborSampler', 'dgl.contrib.sampling.NeighborSampler', False, 'import dgl\n'), (222, 'dgl.contrib.sampling.NeighborSampler', 'dgl.contrib.sampling.NeighborSampler', False, 'import dgl\n'), (229, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (230, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (234, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (235, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (244, 'backend.tensor', 'F.tensor', True, 'import backend as F\n'), (244, 'backend.cpu', 'F.cpu', True, 'import backend as F\n'), (271, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (272, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (273, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (282, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (283, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (286, 'numpy.testing.assert_equal', 'np.testing.assert_equal', True, 'import numpy as np\n'), (372, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (373, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (374, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (376, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (391, 'backend.cpu', 'F.cpu', True, 'import backend as F\n'), (392, 'backend.cpu', 'F.cpu', True, 'import backend as F\n'), (394, 'backend.tensor', 'F.tensor', True, 'import backend as F\n'), (394, 'backend.cpu', 'F.cpu', True, 'import backend as F\n'), (422, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (423, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (424, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (433, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (434, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (437, 'numpy.testing.assert_equal', 'np.testing.assert_equal', True, 'import numpy as np\n'), (458, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (459, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (460, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (462, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (485, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (486, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (487, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (489, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (578, 'backend.cpu', 'F.cpu', True, 'import backend as F\n'), (580, 'backend.cpu', 'F.cpu', True, 'import backend as F\n'), (581, 'backend.sum', 'F.sum', True, 'import backend as F\n'), (583, 'backend.tensor', 'F.tensor', True, 'import backend as F\n'), (583, 'backend.cpu', 'F.cpu', True, 'import backend as F\n'), (661, 'numpy.average', 'np.average', True, 'import numpy as np\n'), (662, 'numpy.average', 'np.average', True, 'import numpy as np\n'), (671, 'backend.cpu', 'F.cpu', True, 'import backend as F\n'), (52, 'numpy.unique', 'np.unique', True, 'import numpy as np\n'), (59, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (145, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (146, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (170, 'numpy.random.randint', 'np.random.randint', True, 'import numpy as np\n'), (170, 'numpy.random.randint', 'np.random.randint', True, 'import numpy as np\n'), (249, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (250, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (251, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (284, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (285, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (391, 'numpy.full', 'np.full', True, 'import numpy as np\n'), (392, 'numpy.full', 'np.full', True, 'import numpy as np\n'), (399, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (400, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (401, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (435, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (436, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (578, 'numpy.full', 'np.full', True, 'import numpy as np\n'), (580, 'numpy.full', 'np.full', True, 'import numpy as np\n'), (608, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (643, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (646, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (648, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (671, 'numpy.full', 'np.full', True, 'import numpy as np\n'), (684, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (695, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (12, 'scipy.sparse.random', 'sp.sparse.random', True, 'import scipy as sp\n'), (40, 'numpy.sort', 'np.sort', True, 'import numpy as np\n'), (118, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (131, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (134, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (141, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (142, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (148, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (180, 'backend.cpu', 'F.cpu', True, 'import backend as F\n'), (181, 'backend.cpu', 'F.cpu', True, 'import backend as F\n'), (265, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (378, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (378, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (417, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (426, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (427, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (464, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (464, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (491, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (491, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (604, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (607, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (644, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (647, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (266, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (267, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (275, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (276, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (383, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (383, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (418, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (419, 'backend.gather_row', 'F.gather_row', True, 'import backend as F\n'), (469, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (469, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (496, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (496, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (128, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (279, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (290, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (430, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n'), (441, 'backend.asnumpy', 'F.asnumpy', True, 'import backend as F\n')]