#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Author: <NAME>
# Contact: <EMAIL>
# Date: 18/12/2018

# This code generates train/test splits of edges from input graphs for evaluating graph embeddings   
# on link prediction. It also provides false train and test edge sets of the required sizes.

# The train/test sets are efficiently generated by: i) obtaining a spanning tree of the input graph
# selected uniformly at random. ii) adding more edges to the spanning tree until the required amount
# of train edges is reached.

from __future__ import division
from __future__ import print_function

import os
import random
import warnings

import networkx as nx
import numpy as np
import scipy as sp
from scipy.sparse import triu
from scipy.sparse import tril
from scipy.sparse.csgraph import depth_first_tree
from sklearn.externals.joblib import Parallel, delayed


def _sanity_check(G):
    r"""
    Helper function that checks if the input graphs contains a single connected component. Raises an error if not.

    Parameters
    ----------
    G : graph
       A NetworkX graph

    Raises
    ------
    ValueError
        If the graph has more than one (weakly) connected component.
    """
    # Compute the number of connected components
    if G.is_directed():
        num_ccs = nx.number_weakly_connected_components(G)
    else:
        num_ccs = nx.number_connected_components(G)

    # Rise an error if more than one CC exists
    if num_ccs != 1:
        raise ValueError("Input graph should contain one (weakly) connected component. "
                         "This graph contains: " + str(num_ccs))


def broder_alg(G, E):
    r"""
    Runs Andrei Broder's algorithm to select uniformly at random a spanning tree of the input
    graph.The direction of the edges included in train_E is taken from E which respects the
    edge directions in the original graph, thus, the results are still valid for directed graphs.
    For pairs of nodes in the original digraphs which have edges in both directions, we randomly
    select the direction of the edge included in the ST.

    Parameters
    ----------
    G : graph
       A NetworkX graph
    E : set
       A set of directed or undirected edges constituting the graph G.

    Returns
    -------
    train_E : set
       A set of edges of G describing the random spanning tree

     References
    ----------
    .. [1] <NAME>, "Generating Random Spanning Trees", Proc. of the 30th Annual Symposium
           on Foundations of Computer Science, pp. 442--447, 1989.
    """
    # Create two partitions, S and T. Initially store all nodes in S.
    S = set(G.nodes)
    T = set()

    # Pick a random node as the "current node" and mark it as visited.
    current_node = random.sample(S, 1).pop()
    S.remove(current_node)
    T.add(current_node)

    # Perform random walk on the graph
    train_E = set()
    while S:
        if G.is_directed():
            neighbour_node = random.sample(list(G.successors(current_node)) + list(G.predecessors(current_node)), 1).pop()
        else:
            neighbour_node = random.sample(list(G.neighbors(current_node)), 1).pop()
        if neighbour_node not in T:
            S.remove(neighbour_node)
            T.add(neighbour_node)
            if random.random() < 0.5:
                if (current_node, neighbour_node) in E:
                    train_E.add((current_node, neighbour_node))
                else:
                    train_E.add((neighbour_node, current_node))
            else:
                if (neighbour_node, current_node) in E:
                    train_E.add((neighbour_node, current_node))
                else:
                    train_E.add((current_node, neighbour_node))
        current_node = neighbour_node

    # Return the set of edges constituting the spanning tree
    return train_E


def wilson_alg(G, E):
    r"""
    Runs Willson's algorithm also known as loop erasing random walk to select uniformly at random
    a spanning tree of the input graph. A set E contains the original direction of edges in graph G,
    and train_E will only include edges which exist in E, thus, the results are still valid for
    digraphs. For pairs of nodes in the original digraphs, which have edges in both directions,
    we select the direction of the edge in the ST at random.

    Parameters
    ----------
    G : graph
       A NetworkX graph
    E : set
       A set of directed or undirected edges constituting the graph G.

    Returns
    -------
    train_E : set
       A set of edges of G describing the random spanning tree

    References
    ----------
    .. [1] <NAME>, "Generating Random Spanning Trees More Quickly than the Cover Time",
           In Proceedings of STOC, pp. 296--303, 1996.
    .. [2] <NAME> and <NAME>, "How to Get a Perfectly Random Sample from a Generic
           Markov Chain and Generate a Random Spanning Tree of a Directed Graph",
           Journal of Algorithms 27, pp. 170--217, 1998.
    """
    # Stores the nodes which are part of the trees created by the LERW.
    intree = set()

    # A dictionary which works as a linked list and stores the spanning tree
    tree = dict()

    # Pick a random node as the root of the spanning tree and add it to intree
    # For undirected graphs this is the correct approach
    r = random.sample(G.nodes, 1).pop()
    intree.add(r)

    for node in G.nodes:
        i = node
        while i not in intree:
            # This random successor works for weighted and unweighted graphs because we just
            # want to select a bunch of edges from the graph, no matter what the weights are.
            if G.is_directed():
                tree[i] = random.sample(list(G.successors(i)) + list(G.predecessors(i)), 1).pop()
            else:
                tree[i] = random.sample(list(G.neighbors(i)), 1).pop()
            i = tree[i]
        i = node
        while i not in intree:
            intree.add(i)
            i = tree[i]

    # Create a set to store the train edges
    train_E = set()

    # This is only relevant for directed graphs to make the selection of edge direction equiprobable
    for e in set(zip(tree.keys(), tree.values())):
        if random.random() < 0.5:
            if e in E:
                train_E.add(e)
            else:
                train_E.add(e[::-1])
        else:
            if e[::-1] in E:
                train_E.add(e[::-1])
            else:
                train_E.add(e)

    # Return the edges of the random spanning tree
    return train_E


def _compute_one_split(G, output_path, owa=True, train_frac=0.51, num_fe_train=None, num_fe_test=None, split_id=0):
    r"""
    Computes one split of train/test edges as well as non-edges from an input graph and writes the data to files.
    The train sets are always connected / weakly connected and span all nodes of the input graph.
    Input graphs (digraphs) cannot contain more than one (weakly) connected component.

    Parameters
    ----------
    G : graph
       A NetworkX graph
    output_path : string
       Indicates the path where data will be stored. Can include a name for all splits to share.
    owa : bool, optional
       Encodes the belief that the network respects or not the open world assumption. Default is True.
       If OWA=True, false train edges can be true test edges. False edges sampled from train graph.
       If OWA=False, closed world is assumed so false train edges are known to be false (not in G)
    train_frac : float, optional
       The relative size (in range (0.0, 1.0]) of the train set with respect to the total number of edges in the graph.
       Default is 0.51.
    num_fe_train : int, optional
       The number of train false edges to generate. Default is same number as true train edges.
    num_fe_test : int, optional
       The number of test false edges to generate. Default is same number as true test edges.
    split_id : int, optional
        The ID of train/test split. Default is 0.
    """
    # Generate train and test edge splits
    train_E, test_E = split_train_test(G, train_frac)

    # Generate the train/test false edges
    if owa:
        train_E_false, test_E_false = generate_false_edges_owa(G, train_E, test_E, num_fe_train, num_fe_test)
    else:
        train_E_false, test_E_false = generate_false_edges_cwa(G, train_E, test_E, num_fe_train, num_fe_test)

    #  Write the computed split to a file
    store_train_test_splits(output_path, train_E, train_E_false, test_E, test_E_false, split_id)


def compute_splits_parallel(G, output_path, owa=True, train_frac=0.51, num_fe_train=None, num_fe_test=None,
                            num_splits=10):
    r"""
    Computes in parallel the required number of train/test splits of edges and non-edges from an input graph
    and writes the data to files. The train sets are always connected / weakly connected and span all nodes
    of the input graph. Input graphs (digraphs) cannot contain more than one (weakly) connected component.
    
    Parameters
    ----------
    G : graph
       A NetworkX graph
    output_path : string
       Indicates the path where data will be stored. Can include a name for all splits to share.  
    owa : bool, optional
       Encodes the belief that the network respects or not the open world assumption. Default is True.
       If OWA=True, false train edges can be true test edges. False edges sampled from train graph. 
       If OWA=False, closed world is assumed so false train edges are known to be false (not in G)
    train_frac : float, optional
       The relative size (in range (0.0, 1.0]) of the train set with respect to the total number of edges in the graph.
       Default is 0.51.
    num_fe_train : int, optional
       The number of train false edges to generate. Default is same number as true train edges.
    num_fe_test : int, optional
       The number of test false edges to generate. Default is same number as true test edges.
    num_splits : int, optional
       The number of train/test splits to generate. Default is 10.
    """
    # Compute the splits sequentially or in parallel
    backend = 'multiprocessing'
    path_func = delayed(_compute_one_split)
    Parallel(n_jobs=num_splits, verbose=True, backend=backend)(
        path_func(G, output_path, owa, train_frac, num_fe_train, num_fe_test, split) for split in range(num_splits))


def split_train_test(G, train_frac=0.51, st_alg='wilson'):
    r"""
    Computes one train/test split of edges from an input graph and returns the results.
    The train set will be (weakly) connected and span all nodes of the input graph (digraph).
    Input graph (digraph) cannot contain more than one (weakly) connected component.
    
    Parameters
    ----------
    G : graph
        A NetworkX graph
    train_frac : float, optional
        The relative size (in range (0.0, 1.0]) of the train set with respect to the total number of edges in the graph.
        Default is 0.51.
    st_alg : basestring, optional
        The algorithm to use for generating the spanning tree constituting the backbone of the train set.
        Options are: 'wilson' and 'broder'. The first option, 'wilson', also known as LERW is much faster in most cases.
        Default is 'wilson'.

    Returns
    -------
    train_E : set
       The set of train edges
    test_E : set
       The set of test edges

    Raises
    ------
    ValueError
        If the train_frac parameter is not in range (0, 1].
        If the input graph G has more than one (weakly) connected component.
    """
    # Sanity check to make sure the input is correct
    _sanity_check(G)
    if train_frac <= 0.0 or train_frac > 1.0:
        raise ValueError('The train_frac parameter needs to be in range: (0.0, 1.0]')
    if train_frac == 1.0:
        return set(G.edges()), set()

    # Create a set of all edges in G
    E = set(G.edges)

    if st_alg == 'broder':
        # Compute a random spanning tree using broder's algorithm
        train_E = broder_alg(G, E)
    else:
        # Compute a random spanning tree using wilson's algorithm
        train_E = wilson_alg(G, E)

    # Fill test edge set as all edges not in the spanning tree
    test_E = E - train_E

    # Compute num train edges
    num_E = len(E)
    num_train_E = np.ceil(train_frac * num_E)

    # Check if the num edges in the spanning tree is already greater than the num train edges
    num_toadd = int(num_train_E - len(train_E))
    if num_toadd <= 0:
        print("WARNING: In order to return a connected train set the train_frac parameter needs to be higher!")
        print("In this case, the provided train set constitutes a random spanning tree of the input graph.")
        print("The train_frac value used is: {}".format(len(train_E) / num_E))
        print("Edges requested: train = {}, test = {}".format(num_train_E, num_E - num_train_E))
        print("Edges returned: train = {}, test = {}".format(len(train_E), num_E - len(train_E)))
    else:
        # Add more edges to train set from test set until it has desired size
        edges = set(random.sample(test_E, num_toadd))
        test_E = test_E - edges
        train_E = train_E | edges

    # Perform some simple checks
    assert E == (test_E | train_E)
    assert len(E) == len(test_E) + len(train_E)
    if num_toadd > 0:
        assert num_train_E == len(train_E)

    # Return the sets of edges
    return train_E, test_E


def rand_split_train_test(G, train_frac=0.51):
    r"""
    Computes one train/test split of edges from an input graph and returns the results.
    The train/test split is computed by randomly removing 1-train_frac edges from the graph.
    From the remaining edges, those in the mainCC constitute the train edges. From the set
    of removed edges, those whose nodes are in the train set, are considered part or the
    test set. The proportion of train/test edges returned might not be the required one.
    The train set will be (weakly) connected and span all nodes of the input graph.
    Input graph (digraph) can contain one or many (weakly) connected components.

    Parameters
    ----------
    G : graph
        A NetworkX graph
    train_frac : float, optional
        The relative size (in range (0.0, 1.0]) of the train set with respect to the total number of edges in the graph.
        Default is 0.51.

    Returns
    -------
    train_E : set
        The set of train edges
    test_E : set
        The set of test edges

    Raises
    ------
    ValueError
        If the train_frac parameter is not in range (0, 1].
    """
    if train_frac <= 0.0 or train_frac > 1.0:
        raise ValueError('The train_frac parameter needs to be in range: (0.0, 1.0]')
    if train_frac == 1.0:
        return set(G.edges()), set()

    # Create a set of all edges in G
    E = set(G.edges)
    num_E = len(E)

    # Compute the potential number of train and test edges which corresponds to the fraction given
    num_train_E = int(np.ceil(train_frac * num_E))
    num_test_E = int(num_E - num_train_E)

    # Randomly remove 1-train_frac edges from the graph and store them as potential test edges
    pte_edges = set(random.sample(E, num_test_E))

    # The remaining edges are potential train edges
    ptr_edges = E - pte_edges

    # Create a graph containing all ptr_edges and compute the mainCC
    if G.is_directed():
        H = nx.DiGraph()
        H.add_edges_from(ptr_edges)
        maincc = max(nx.weakly_connected_component_subgraphs(H), key=len)
    else:
        H = nx.Graph()
        H.add_edges_from(ptr_edges)
        maincc = max(nx.connected_component_subgraphs(H), key=len)

    # The edges in the mainCC graph are the actual train edges
    train_E = set(maincc.edges)

    # Remove potential test edges for which the end nodes do not exist in the train_E
    test_E = set()
    for (src, dst) in pte_edges:
        if src in maincc.nodes and dst in maincc.nodes:
            test_E.add((src, dst))

    # Return the sets of edges
    return train_E, test_E


def naive_split_train_test(G, train_frac=0.51):
    r"""
    Computes one train/test split of edges from an input graph and returns the results.
    The sets are computed using the naive approach that checks connectivity of the graph
    for each removed edge. If graph gets disconnected, that edges is not removed.
    The train set will be (weakly) connected and span all nodes of the input graph.
    Input graph (digraph) cannot contain more than one (weakly) connected component.

    Parameters
    ----------
    G : graph
      A NetworkX graph
    train_frac : float, optional
        The relative size (in range (0.0, 1.0]) of the train set with respect to the total number of edges in the graph.
        Default is 0.51.

    Returns
    -------
    train_E : set
       The set of train edges
    test_E : set
        The set of test edges

    Raises
    ------
    ValueError
        If the train_frac parameter is not in range (0, 1].
        If the input graph G has more than one (weakly) connected component.
    """
    # Sanity check to make sure the input is correct
    _sanity_check(G)
    if train_frac <= 0.0 or train_frac > 1.0:
        raise ValueError('The train_frac parameter needs to be in range: (0.0, 1.0]')
    if train_frac == 1.0:
        return set(G.edges()), set()

    # Is directed
    directed = G.is_directed()

    G = G.copy()

    # Create a set of all edges in G
    aux = np.array(G.edges)
    np.random.shuffle(aux)
    E = set([tuple(edge) for edge in aux])

    # Compute num train edges
    num_E = len(E)
    num_train_E = np.ceil(train_frac * num_E)
    num_test_E = num_E - num_train_E

    # Initialize train edges to an empty set
    train_E = set(G.edges())

    # Initialize test edges to an empty set
    test_E = set()

    # Iterate over shuffled edges, add to train/val sets

    for i, edge in enumerate(E):
        # if i % 500 == 0:
        #    print('{}/{}'.format(i, num_test_E))
        node1 = edge[0]
        node2 = edge[1]

        # If removing edge would disconnect a connected component, backtrack and move on
        G.remove_edge(node1, node2)
        if directed:
            if nx.number_weakly_connected_components(G) > 1:
                G.add_edge(node1, node2)
                continue
        else:
            if nx.number_connected_components(G) > 1:
                G.add_edge(node1, node2)
                continue

        # Fill test_edges
        if len(test_E) < num_test_E:
            test_E.add(edge)
            train_E.remove(edge)
        else:
            break

    # Perform some simple checks
    assert E == (test_E | train_E)
    assert len(E) == len(train_E) + len(test_E)

    # Return the sets of edges
    return train_E, test_E


def generate_false_edges_owa(G, train_E, test_E, num_fe_train=None, num_fe_test=None):
    r"""
    This method generates false train and test edges for both directed and undirected graphs.
    The train and test sets are non overlapping.
    Follows the open world assumption, so false train edges are generated only using the true train edges,
    so false train edges can be true test edges. This is the case for evolving graphs where edges can only appear.
    For undirected graphs the output is sorted (smallNodeID, bigNodeID)

    Parameters
    ----------
    G : graph
       A NetworkX graph
    train_E : set
       The set of train edges.
    test_E : set
       The set of test edges.
    num_fe_train : int, optional
       The number of train false edges to generate. Default is same number as true train edges.
    num_fe_test : int, optional
       The number of test false edges to generate. Default is same number as true test edges.

    Returns
    -------
    train_false_E : set
       The set of false train edges
    test_false_E : set
       The set of false test edges

    Raises
    ------
    ValueError
        If the input graph G has more than one (weakly) connected component.
        If more false edges than existing in the graph are required.
    """
    # Sanity check to make sure the input is correct
    _sanity_check(G)

    # Create a set of vertices
    V = set(G.nodes)

    # Initialize the sizes of the false edges
    if num_fe_train is None:
        num_fe_train = len(train_E)

    if num_fe_test is None:
        num_fe_test = len(test_E)

    # Make sure the required amount of false edges can be generated
    max_nonedges = len(V) * len(V) - len(train_E)
    if num_fe_train > max_nonedges:
        raise ValueError('Too many false train edges required! Max available for train+test is {}'.format(max_nonedges))
    else:
        if num_fe_train + num_fe_test > max_nonedges:
            warnings.warn('Too many false edges required in train+test! '
                          'Using maximum number of false test edges available: {}'.format(max_nonedges - num_fe_train))
            num_fe_test = max_nonedges - num_fe_train

    # Create sets to store the false edges
    train_E_false = set()
    test_E_false = set()

    # Generate negative train edges
    while len(train_E_false) < num_fe_train:
        edge = tuple(random.sample(V, 2))
        redge = tuple(reversed(edge))
        if edge not in train_E:
            if G.is_directed():
                train_E_false.add(edge)
            else:
                if redge not in train_E:
                    train_E_false.add(tuple(sorted(edge)))

    # Generate negative test edges
    while len(test_E_false) < num_fe_test:
        edge = tuple(random.sample(V, 2))
        redge = tuple(reversed(edge))
        if edge not in train_E and edge not in test_E and edge not in train_E_false:
            if G.is_directed():
                test_E_false.add(edge)
            else:
                if redge not in train_E and redge not in test_E and redge not in train_E_false:
                    test_E_false.add(tuple(sorted(edge)))

    # Perform some simple check before returning the result
    assert len(train_E_false) == num_fe_train
    assert len(test_E_false) == num_fe_test
    assert train_E_false.isdisjoint(test_E_false)
    assert train_E_false.isdisjoint(train_E)
    assert test_E_false.isdisjoint(train_E | test_E)

    # Return the sets of false edges
    return train_E_false, test_E_false


def generate_false_edges_cwa(G, train_E, test_E, num_fe_train=None, num_fe_test=None):
    r"""
    This method generates false train and test edges for both directed and undirected graphs.
    The train and test sets are non overlapping.
    Follows the closed world assumption, so false train edges are selected as known to be false.
    This is the case for some networks e.g. protein-protein interaction where information about
    both the positive class (existing edges) and the negative class (missing edges) exists.
    For undirected graphs the output is sorted (smallNodeID, bigNodeID)

    Parameters
    ----------
    G : graph
       A NetworkX graph
    train_E : set
       The set of train edges.
    test_E : set
       The set of test edges.
    num_fe_train : int, optional
       The number of train false edges to generate. Default is same number as true train edges.
    num_fe_test : int, optional
       The number of test false edges to generate. Default is same number as true test edges.

    Returns
    -------
    train_false_E : set
       The set of false train edges
    test_false_E : set
       The set of false test edges

    Raises
    ------
    ValueError
        If the input graph G has more than one (weakly) connected component.
        If more false edges than existing in the graph are required.
    """
    # Sanity check to make sure the input is correct
    _sanity_check(G)

    # Create a set of vertices
    V = set(G.nodes)

    # Initialize the sizes of the false edges
    if num_fe_train is None:
        num_fe_train = len(train_E)

    if num_fe_test is None:
        num_fe_test = len(test_E)

    # Make sure the required amount of false edges can be generated
    max_nonedges = len(V) * len(V) - len(G.edges)
    if num_fe_train > max_nonedges:
        raise ValueError(
            'Too many false train edges required! Max available for train+test is {}'.format(max_nonedges))
    else:
        if num_fe_train + num_fe_test > max_nonedges:
            warnings.warn('Too many false edges required in train+test! '
                          'Using maximum number of false test edges available: {}'.format(max_nonedges - num_fe_train))
            # num_fe_test = max_nonedges - num_fe_train
            return _getall_false_edges(G, (1.0*num_fe_train)/max_nonedges)

    # Create sets to store the false edges
    train_E_false = set()
    test_E_false = set()

    # Generate negative train edges
    while len(train_E_false) < num_fe_train:
        edge = tuple(random.sample(V, 2))
        redge = tuple(reversed(edge))
        if edge not in train_E and edge not in test_E:
            if G.is_directed():
                train_E_false.add(edge)
            else:
                if redge not in train_E and redge not in test_E:
                    train_E_false.add(tuple(sorted(edge)))

    # Generate negative test edges
    while len(test_E_false) < num_fe_test:
        edge = tuple(random.sample(V, 2))
        redge = tuple(reversed(edge))
        if edge not in train_E and edge not in test_E and edge not in train_E_false:
            if G.is_directed():
                test_E_false.add(edge)
            else:
                if redge not in train_E and redge not in test_E and redge not in train_E_false:
                    test_E_false.add(tuple(sorted(edge)))

    # Perform some simple check before returning the result
    assert len(train_E_false) == num_fe_train
    assert len(test_E_false) == num_fe_test
    assert train_E_false.isdisjoint(test_E_false)
    assert train_E_false.isdisjoint(train_E | test_E)
    assert test_E_false.isdisjoint(train_E | test_E)

    # Return the sets of false edges
    return train_E_false, test_E_false

def generate_false_edges_cwa_close1(G, train_E, test_E, num_fe_train=None, num_fe_test=None, length=None):
    r"""
    This method generates false train and test edges for both directed and undirected graphs.
    The train and test sets are non overlapping.
    Follows the closed world assumption, so false train edges are selected as known to be false.
    This is the case for some networks e.g. protein-protein interaction where information about
    both the positive class (existing edges) and the negative class (missing edges) exists.
    For undirected graphs the output is sorted (smallNodeID, bigNodeID)

    Parameters
    ----------
    G : graph
       A NetworkX graph
    train_E : set
       The set of train edges.
    test_E : set
       The set of test edges.
    num_fe_train : int, optional
       The number of train false edges to generate. Default is same number as true train edges.
    num_fe_test : int, optional
       The number of test false edges to generate. Default is same number as true test edges.

    Returns
    -------
    train_false_E : set
       The set of false train edges
    test_false_E : set
       The set of false test edges

    Raises
    ------
    ValueError
        If the input graph G has more than one (weakly) connected component.
        If more false edges than existing in the graph are required.
    """
    # Sanity check to make sure the input is correct
    _sanity_check(G)

    # Create a set of vertices
    V = set(G.nodes)

    # Initialize the sizes of the false edges
    if num_fe_train is None:
        num_fe_train = len(train_E)

    if num_fe_test is None:
        num_fe_test = len(test_E)

    # Make sure the required amount of false edges can be generated
    max_nonedges = len(V) * len(V) - len(G.edges)
    if num_fe_train > max_nonedges:
        raise ValueError(
            'Too many false train edges required! Max available for train+test is {}'.format(max_nonedges))
    else:
        if num_fe_train + num_fe_test > max_nonedges:
            warnings.warn('Too many false edges required in train+test! '
                          'Using maximum number of false test edges available: {}'.format(max_nonedges - num_fe_train))
            # num_fe_test = max_nonedges - num_fe_train
            return _getall_false_edges(G, (1.0*num_fe_train)/max_nonedges)

    # Create sets to store the false edges
    train_E_false = set()
    test_E_false = set()

    # Generate negative train edges
    while len(train_E_false) < num_fe_train:
        edge = tuple(random.sample(V, 2))
        redge = tuple(reversed(edge))
        if edge not in train_E and edge not in test_E:
            if G.is_directed():
                if (not nx.has_path(G, source=edge[0], target=edge[1])) or (nx.shortest_path_length(G, source=edge[0], target=edge[1]) >= length):
                    train_E_false.add(edge)
            else:
                if redge not in train_E and redge not in test_E:
                    if (not nx.has_path(G, source=edge[0], target=edge[1])) or (nx.shortest_path_length(G, source=edge[0], target=edge[1]) >= length):
                        train_E_false.add(tuple(sorted(edge)))

    # Generate negative test edges
    while len(test_E_false) < num_fe_test:
        edge = tuple(random.sample(V, 2))
        redge = tuple(reversed(edge))
        if edge not in train_E and edge not in test_E and edge not in train_E_false:
            if G.is_directed():
                if (not nx.has_path(G, source=edge[0], target=edge[1])) or (nx.shortest_path_length(G, source=edge[0], target=edge[1]) >= length):
                    test_E_false.add(edge)
            else:
                if redge not in train_E and redge not in test_E and redge not in train_E_false:
                    if (not nx.has_path(G, source=edge[0], target=edge[1])) or (nx.shortest_path_length(G, source=edge[0], target=edge[1]) >= length):
                        test_E_false.add(tuple(sorted(edge)))

    # Perform some simple check before returning the result
    assert len(train_E_false) == num_fe_train
    assert len(test_E_false) == num_fe_test
    assert train_E_false.isdisjoint(test_E_false)
    assert train_E_false.isdisjoint(train_E | test_E)
    assert test_E_false.isdisjoint(train_E | test_E)

    # Return the sets of false edges
    return train_E_false, test_E_false



def generate_false_edges_cwa_close(G, train_E, test_E, num_fe_train=None, num_fe_test=None, length=None):
    r"""
    This method generates false train and test edges for both directed and undirected graphs.
    The train and test sets are non overlapping.
    Follows the closed world assumption, so false train edges are selected as known to be false.
    This is the case for some networks e.g. protein-protein interaction where information about
    both the positive class (existing edges) and the negative class (missing edges) exists.
    For undirected graphs the output is sorted (smallNodeID, bigNodeID)

    Parameters
    ----------
    G : graph
       A NetworkX graph
    train_E : set
       The set of train edges.
    test_E : set
       The set of test edges.
    num_fe_train : int, optional
       The number of train false edges to generate. Default is same number as true train edges.
    num_fe_test : int, optional
       The number of test false edges to generate. Default is same number as true test edges.
    length: int, optional
        false train and test edges whose shortest path length is larger than the value of length.

    Returns
    -------
    train_false_E : set
       The set of false train edges
    test_false_E : set
       The set of false test edges

    Raises
    ------
    ValueError
        If the input graph G has more than one (weakly) connected component.
        If more false edges than existing in the graph are required.
    """
    # Sanity check to make sure the input is correct
    _sanity_check(G)

    # Create a set of vertices
    V = set(G.nodes)

    # Initialize the sizes of the false edges
    if num_fe_train is None:
        num_fe_train = len(train_E)

    if num_fe_test is None:
        num_fe_test = len(test_E)

    # Make sure the required amount of false edges can be generated
    max_nonedges = len(V) * len(V) - len(G.edges)
    if num_fe_train > max_nonedges:
        raise ValueError(
            'Too many false train edges required! Max available for train+test is {}'.format(max_nonedges))
    else:
        if num_fe_train + num_fe_test > max_nonedges:
            warnings.warn('Too many false edges required in train+test! '
                          'Using maximum number of false test edges available: {}'.format(max_nonedges - num_fe_train))
            # num_fe_test = max_nonedges - num_fe_train
            return _getall_false_edges(G, (1.0*num_fe_train)/max_nonedges)

    # Create sets to store the false edges
    train_E_false = set()
    test_E_false = set()

    # Generate negative train edges
    while len(train_E_false) < num_fe_train:
        edge = tuple(random.sample(V, 2))
        redge = tuple(reversed(edge))
        if length is not None:
            if G.is_directed():
                if not nx.has_path(G, source=edge[0], target=edge[1]):
                    if not nx.has_path(G, source=edge[1], target=edge[0]):
                        if edge not in train_E and edge not in test_E:
                                train_E_false.add(edge)
                                train_E_false.add(redge)
                                continue
                        continue
                    continue
                if not nx.has_path(G, source=edge[1], target=edge[0]):
                    continue
            if G.is_directed() and (nx.shortest_path_length(G, source=edge[0], target=edge[1]) <= length
                                    or nx.shortest_path_length(G, source=edge[1], target=edge[0]) <= length):
                continue
            else:
                if nx.shortest_path_length(G, source=edge[0], target=edge[1]) <= length:
                    continue
        if edge not in train_E and edge not in test_E:
            if G.is_directed():
                train_E_false.add(edge)
            else:
                if redge not in train_E and redge not in test_E:
                    train_E_false.add(tuple(sorted(edge)))

    # Generate negative test edges
    while len(test_E_false) < num_fe_test:
        edge = tuple(random.sample(V, 2))
        redge = tuple(reversed(edge))
        if length is not None:
            if (edge in test_E) or (edge in train_E_false) or (edge in test_E_false):
                continue
            if G.is_directed() :
                if not nx.has_path(G, source=edge[0], target=edge[1]):
                    if not nx.has_path(G, source=edge[1], target=edge[0]):
                        if edge not in train_E and edge not in test_E:
                                train_E_false.add(edge)
                                train_E_false.add(redge)
                                continue
                        continue
                    continue
                if not nx.has_path(G, source=edge[1], target=edge[0]):
                    continue
            if G.is_directed() and (nx.shortest_path_length(G, source=edge[0], target=edge[1]) <= length
                                    or nx.shortest_path_length(G, source=edge[1], target=edge[0]) <= length):
                continue
            else:
                if nx.shortest_path_length(G, source=edge[0], target=edge[1]) <= length:
                    continue
        if (edge not in train_E) and (edge not in test_E) and (edge not in train_E_false):
            if G.is_directed():
                test_E_false.add(edge)
            else:
                if (redge not in train_E) and (redge not in test_E) and (redge not in train_E_false):
                    test_E_false.add(redge)

    # Perform some simple check before returning the result
    assert len(train_E_false) == num_fe_train
    assert len(test_E_false) == num_fe_test
    assert train_E_false.isdisjoint(test_E_false)
    assert train_E_false.isdisjoint(train_E | test_E)
    assert test_E_false.isdisjoint(train_E | test_E)

    # Return the sets of false edges
    return train_E_false, test_E_false


def _getall_false_edges(G, fe_train_frac):
    print("Generating all non-edges and splitting them in train and test...")
    train_E_false = list()
    test_E_false = list()
    for e in nx.non_edges(G):
        r = random.uniform(0, 1)
        if r <= fe_train_frac:
            train_E_false.append(e)
        else:
            test_E_false.append(e)

    return train_E_false, test_E_false


def redges_false(train_E, test_E, output_path=None):
    r"""
    For directed graphs computes all non-edges (a->b) such that the opposite edge (a<-b) exists in the graph.
    It does this for both the train and test edge sets. These non-edges can be used to asses the performance
    of the embedding methods on predicting non-reciprocated edges.

    Parameters
    ----------
    train_E : set
       The set of train edges.
    test_E : set
       The set of test edges.
    output_path : string, optional
        A path or file where to store the results. Default None.

    Returns
    -------
    train_redges_false : set
        A set of edges respecting the mentioned property regarding the train edges
    test_redges_false : set
        A set of edges respecting the mentioned property on the complete graph
    """
    # Reverse all train and test edges
    train_redges_false = set(tuple(reversed(edge_tuple)) for edge_tuple in train_E)
    test_redges_false = set(tuple(reversed(edge_tuple)) for edge_tuple in test_E)

    # Keep only the reversed edges which are not real train edges
    train_redges_false = train_redges_false - train_E

    # Keep only the test reversed edges which are not true edges in the graph
    test_redges_false = test_redges_false - train_E
    test_redges_false = test_redges_false - test_E

    if output_path is not None:
        # Store the reversed edges
        train_redges_false_np = np.array([list(edge_tuple) for edge_tuple in train_redges_false])
        test_redges_false_np = np.array([list(edge_tuple) for edge_tuple in test_redges_false])
        # Save the splits in different files
        np.savetxt(output_path, train_redges_false_np, delimiter=',', fmt='%d')
        np.savetxt(output_path, test_redges_false_np, delimiter=',', fmt='%d')

    # Return the computed sets
    return train_redges_false, test_redges_false


def store_train_test_splits(output_path, train_E, train_E_false, test_E, test_E_false, split_id=0):
    r"""
    Writes the sets of true and false edges to files in the provided path. All files will share
    the same split number as an identifier. If any folder in the path do not exist, it will be generated.
    
    Parameters
    ----------
    output_path : string
       Indicates the path where data will be stored. It can also include a name for all the splits to share.
    train_E : set
       Set of train edges
    train_E_false : set
       Set of train non-edges
    test_E : set
       Set of test edges
    test_E_false : set
       Set of test non-edges
    split_id : int, optional
       The ID of train/test split to be stored. Default is 0.

    Returns
    -------
    filenames : list
        A list of strings, the names given to the 4 files where the true and false train and test edge are stored.
    """
    # Create path if it does not exist
    if not os.path.exists(output_path):
        os.makedirs(output_path)

    # Convert edge-lists to numpy arrays
    train_E = np.array([list(edge_tuple) for edge_tuple in train_E])
    train_E_false = np.array([list(edge_tuple) for edge_tuple in train_E_false])
    test_E = np.array([list(edge_tuple) for edge_tuple in test_E])
    test_E_false = np.array([list(edge_tuple) for edge_tuple in test_E_false])

    filenames = (os.path.join(output_path, "trE_{}.csv".format(split_id)),
                 os.path.join(output_path, "negTrE_{}.csv".format(split_id)),
                 os.path.join(output_path, "teE_{}.csv".format(split_id)),
                 os.path.join(output_path, "negTeE_{}.csv".format(split_id)))

    # Save the splits in different files
    np.savetxt(fname=filenames[0], X=train_E, delimiter=',', fmt='%d')
    np.savetxt(fname=filenames[1], X=train_E_false, delimiter=',', fmt='%d')
    np.savetxt(fname=filenames[2], X=test_E, delimiter=',', fmt='%d')
    np.savetxt(fname=filenames[3], X=test_E_false, delimiter=',', fmt='%d')

    # Return the names given to the 4 files where data is stored
    return filenames


def store_edgelists(train_path, test_path, train_edges, test_edges):
    r"""
    Writes the train and test edgelists to files with the specified names.

    Parameters
    ----------
    train_path : string
       Indicates the path where the train data will be stored.
    test_path : string
       Indicates the path where the test data will be stored.
    train_edges : array_like
       Set of train true and false edges
    test_edges : array_like
       Set of test true and false edges
    """
    # Convert edge-lists to numpy arrays
    train_edges = np.array([list(edge_tuple) for edge_tuple in train_edges])
    test_edges = np.array([list(edge_tuple) for edge_tuple in test_edges])

    # Save the splits in different files
    np.savetxt(fname=train_path, X=train_edges, delimiter=',', fmt='%d')
    np.savetxt(fname=test_path, X=test_edges, delimiter=',', fmt='%d')


def check_overlap(filename, num_sets):
    r"""
    Shows the amount of overlap (shared elements) between edge sets from different random splits.
    The path and name of the set (without split ID) for which to compute the overlap is required. 
    The method will read num_sets from the same path and compute the overlap between them. 

    Parameters
    ----------
    filename : string
       Indicates the path and name (without split ID) of the first set.
       The sets are assumed to have sequential split IDs starting at 0. 
    num_sets : int
       The number of sets for which to check the overlap.
    """
    # Load the first set and transform it into a list of tuples
    S = np.loadtxt(filename+"_0.csv", delimiter=',', dtype=int)
    S = set(map(tuple, S))

    # Initialize the intersection and union sets as all elements in first edge set
    intrs = S
    union = S

    # Sequentially add the rest of the sets and check overlap
    for i in range(num_sets-1):
        # Read a new edge set
        S = np.loadtxt(filename+"_{}.csv".format(i+1), delimiter=',', dtype=int)
        S = set(map(tuple, S))
        # Update intersection and union sets
        intrs = intrs & S
        union = union | S
        # Print the information on screen
        print("Intersection of {} sets is {}".format(i+2, len(intrs)))
        print("Union of {} sets is {}".format(i+2, len(union)))
        print("Jaccard coefficient: {}".format(len(intrs)/len(union)))
        print("")


def random_edge_sample(a, samp_frac=0.01, directed=False):
    r"""
    Returns a sample of positive and negative edges from the given graph represented by `a` selected uniformly at
    random without replacement. If the directed flag is set to False the samples are obtained only from the upper
    triangle.

    Parameters
    ----------
    a : sparse matrix
        A sparse adjacency matrix representing a graph.
    samp_frac : float, optional
        An float representing the fraction of elements to sample. Default is 1.0 (1%)
    directed : bool, optional
        A flag indicating if the adjacency matrix should be considered directed or undirected. If undirected
        indices are obtained only from the lower triangle. Default is False.

    Returns
    -------
    pos_e : ndarray
        Positive edges
    neg_e : ndarray
        Negative edges
    """
    n = a.shape[0]

    if directed:
        num_samp = int(n ** 2 * samp_frac / 100)
        lin_indx_a = np.ravel_multi_index(a.nonzero(), (n, n))
        # randomly generate linear indices
        lin_indx = np.random.randint(0, n ** 2, num_samp)
    else:
        # For undir graphs we only need to sample half the num nodes
        num_samp = int((n*(n-1))/2 * (samp_frac / 100))
        lin_indx_a = np.ravel_multi_index(triu(a, k=1).nonzero(), (n, n))
        ij = np.random.randint(0, n, size=(2, num_samp))
        ij.sort(axis=0)
        lin_indx = np.ravel_multi_index((ij[0], ij[1]), (n, n))

    pos_e = np.intersect1d(lin_indx, lin_indx_a)
    neg_e = np.setdiff1d(lin_indx, lin_indx_a)

    # Remove the self edges
    lin_diag_indxs = np.ravel_multi_index(np.diag_indices(n), (n, n))
    pos_e = np.setdiff1d(pos_e, lin_diag_indxs)
    neg_e = np.setdiff1d(neg_e, lin_diag_indxs)

    # Unravel the linear indices to obtain src, dst pairs
    pos_e = np.array(np.unravel_index(np.array(pos_e), (n, n))).T
    neg_e = np.array(np.unravel_index(np.array(neg_e), (n, n))).T

    return pos_e, neg_e


def random_edge_sample_other(a, samp_frac=0.01, directed=False):
    r"""
    Returns a sample of positive and negative edges from the given graph represented by `a` selected uniformly at
    random without replacement. If the directed flag is set to False the samples are obtained only from the upper
    triangle.

    A different take on the random sampling technique. Probably less efficient than the other one. For undir graphs
    generates lots of candidates also from the bottom triangle to reach the desired density, this is not as efficient
    as the other version.

    Parameters
    ----------
    a : sparse matrix
        A sparse adjacency matrix representing a graph.
    samp_frac : float, optional
        An float representing the fraction of elements to sample. Default is 0.01 (1%)
    directed : bool, optional
        A flag indicating if the adjacency matrix should be considered directed or undirected. If undirected
        indices are obtained only from the lower triangle. Default is False.

    Returns
    -------
    pos_e : ndarray
        Positive edges
    neg_e : ndarray
        Negative edges
    """
    n = a.shape[0]
    num_samp = int(n**2 * samp_frac)

    # Generate sparse random matrix representing mask of samples
    density = (num_samp + n) / n**2
    mask = sp.sparse.rand(n, n, density)

    if not directed:
        # For undir graphs we only look at the upper triangle
        mask = triu(mask, k=1)
    else:
        # Remove elements from diagonal
        mask.setdiag(0)
        mask.eliminate_zeros()

    mask.data[:] = 1
    lin_indx_samp = np.ravel_multi_index(mask.nonzero(), (n, n))

    # All positive edges sampled in mask will stay in aux
    aux = mask.multiply(a)
    pos_e = np.array(aux.nonzero()).T

    # The rest of the lin indx not positive are negative
    lin_indx_ne = np.setdiff1d(lin_indx_samp, np.ravel_multi_index(aux.nonzero(), (n, n)))
    neg_e = np.array(np.unravel_index(lin_indx_ne, (n, n)))

    return pos_e, neg_e


def quick_split(G, train_frac=0.51):
    r"""
    Computes one train/test split of edges from an input graph and returns the results.
    The train set will be (weakly) connected and span all nodes of the input graph (digraph).
    This implementation uses a depth first tree to obtain edges covering all nodes for the train graph.
    Input graph (digraph) cannot contain more than one (weakly) connected component.

    Parameters
    ----------
    G : graph
        A NetworkX graph
    train_frac : float, optional
        The relative size (in range (0.0, 1.0]) of the train set with respect to the total number of edges in the graph.
        Default is 0.51.

    Returns
    -------
    train_E : array
       Column array of train edges as pairs src, dst
    test_E : array
       Column array of test edges as pairs src, dst

    Raises
    ------
    ValueError
        If the train_frac parameter is not in range (0, 1].
        If the input graph G has more than one (weakly) connected component.
    """
    _sanity_check(G)
    if train_frac <= 0.0 or train_frac > 1.0:
        raise ValueError('The train_frac parameter needs to be in range: (0.0, 1.0]')
    if train_frac == 1.0:
        return set(G.edges()), set()

    # Restrict input graph to its main cc
    if nx.is_directed(G):
        a = nx.adj_matrix(G)
    else:
        a = triu(nx.adj_matrix(G), k=1)

    # Compute initial statistics and linear indx of nonzeros
    n = a.shape[0]
    num_tr_e = int(a.nnz * train_frac)
    nz_lin_ind = np.ravel_multi_index(a.nonzero(), (n, n))

    # Build a dft starting at a random node. If dir false returns only upper triang
    dft = depth_first_tree(a, np.random.randint(0, a.shape[0]), directed=nx.is_directed(G))
    if nx.is_directed(G):
        dft_lin_ind = np.ravel_multi_index(dft.nonzero(), (n, n))
    else:
        dft_lin_ind = np.ravel_multi_index(triu(tril(dft).T + dft, k=1).nonzero(), (n, n))

    # From all nonzero indx remove those in dft. From the rest take enough to fill train quota. Rest are test
    rest_lin_ind = np.setdiff1d(nz_lin_ind, dft_lin_ind)
    aux = np.random.choice(rest_lin_ind, num_tr_e-len(dft_lin_ind), replace=False)
    lin_tr_e = np.union1d(dft_lin_ind, aux)
    lin_te_e = np.setdiff1d(rest_lin_ind, aux)

    # Unravel the linear indices to obtain src, dst pairs
    tr_e = np.array(np.unravel_index(np.array(lin_tr_e), (n, n))).T
    te_e = np.array(np.unravel_index(np.array(lin_te_e), (n, n))).T
    return tr_e, te_e


def quick_nonedges(G, train_frac=0.51, fe_ratio=1.0):
    r"""
    Computes one train/test split of non-edges from an input graph and returns the results.
    The negative train and test edges will have no overlap. Also there will be no overlap between false train and test
    edges and real ones. No selfloop false edges will be generated.
    Input graph (digraph) cannot contain more than one (weakly) connected component.

    Parameters
    ----------
    G : graph
        A NetworkX graph
    train_frac : float, optional
        The relative size (in range (0.0, 1.0]) of the train false edge set w.r.t. total number of edges in graph.
        Default is 0.51.
    fe_ratio : float, optional
        The ratio of negative to positive edges to sample. For fr_ratio > 0 and < 1 less false than true edges will be
        generated. For fe_edges > 1 more false than true edges will be generated. Default 1, same amounts.

    Returns
    -------
    train_E : array
       Column array of train edges as pairs src, dst
    test_E : array
       Column array of test edges as pairs src, dst

    Raises
    ------
    ValueError
        If more false edges than existing in the graph are required.
    """
    # fe_ration can be any float or keyword 'prop'
    a = nx.adj_matrix(G)
    n = a.shape[0]
    density = a.nnz / n ** 2
    if fe_ratio == 'prop':
        fe_ratio = np.floor(1.0 / density)
    if not nx.is_directed(G):
        num_fe = int((a.nnz/2.0) * fe_ratio)
    else:
        num_fe = int(a.nnz * fe_ratio)
    num_fe_tr = int(train_frac * num_fe)

    # Make sure we have enough false edges
    if num_fe > (n**2 - (a.nnz + n)):
        raise ValueError('Too many false edges required!')
        # warnings.warn('Too many false edges required in train+test! '
        #              'Using maximum number of false test edges available: {}'.format(n**2-ut.nnz))
        # return _getall_false_edges(G, (1.0 * num_fe_train) / max_nonedges)

    # Get linear indexes of 1s in A
    lin_indexes = np.ravel_multi_index(a.nonzero(), (n, n))
    inv_indx = np.union1d(lin_indexes, np.ravel_multi_index(np.diag_indices(n), (n, n)))
    # we could generate more FE than we need to make sure we find enough 0s
    candidates = np.random.randint(0, n**2, size=int(num_fe/(1-density)))

    # make sure there is no overlap
    fe_lin_ind = np.setdiff1d(candidates, inv_indx)

    while len(fe_lin_ind) < num_fe:
        new_cands = np.random.randint(0, n ** 2, size=num_fe-len(fe_lin_ind))
        valid_cands = np.setdiff1d(new_cands, inv_indx)
        fe_lin_ind = np.union1d(fe_lin_ind, valid_cands)

    fe_lin_ind = fe_lin_ind[:num_fe]
    aux = np.array(np.unravel_index(fe_lin_ind, (n, n))).T
    fe_tr = aux[:num_fe_tr, :]
    fe_te = aux[num_fe_tr:, :]
    return fe_tr, fe_te

