def load_train_test_graphs(dataset, recache_input):
    raw_mat_path = 'data/{}.npz'.format(dataset)
    train_graph_path = 'data/{}/train_graph.pkl'.format(dataset)
    test_graph_path = 'data/{}/test_graph.pkl'.format(dataset)

    if recache_input:
        print('loading sparse matrix from {}'.format(raw_mat_path))
        m = load_sparse_csr(raw_mat_path)

        print('splitting train and test...')
        train_m, test_m = split_train_test(
            m,
            weights=[0.9, 0.1])

        print('converting to nx.DiGraph')
        train_g = nx.from_scipy_sparse_matrix(train_m, create_using=nx.DiGraph(), edge_attribute='sign')
        test_g = nx.from_scipy_sparse_matrix(test_m, create_using=nx.DiGraph(), edge_attribute='sign')
                
        print('saving train and test graphs...')
        nx.write_gpickle(train_g, train_graph_path)
        nx.write_gpickle(test_g, test_graph_path)
    else:
        print('loading train and test graphs...')
        train_g = nx.read_gpickle(train_graph_path)
        test_g = nx.read_gpickle(test_graph_path)
    return train_g, test_g

    def __init__(self, post_fcn = None):
        self.parse_args()
        options = self.options

        input = options.input
        # Input graph
        input_path = os.path.join(input, input + '.gpk')
        self.g = nx.read_gpickle(input_path)
        if not self.g:
            raise Exception("null input file for input path %s" % input_path)
        # Output (reduced) graph.
        # Nodes: (lat, long) tuples w/ list of associated users & location strings
        # Edges: weight: number of links in this direction
        self.r = nx.DiGraph()

        conn = Connection()
        self.input_db = conn[self.options.input_db]
        self.input_coll = self.input_db[self.options.input_coll]

        print "now processing"
        self.reduce()
        print geo_stats(self.r)

        if options.write:
            geo_path = os.path.join(input, input + '.grg')
            nx.write_gpickle(self.r, geo_path)

def mainmain():
    redirects = process_redirects(sys.argv[1])
    print "redirects", len(redirects)
    sys.stderr.write(sys.argv[1] + " processed\n")

    links = process_links(sys.argv[2], redirects)
    links_t = len(links)
    print "links", links_t
    sys.stderr.write(sys.argv[2] + " processed\n")

    G = networkx.Graph()

    articles_processed = 0

    for article in links:
        articles_processed = articles_processed + 1
        if (articles_processed % 100000) == 0:
            sys.stdout.write(
                "links processed="
                + str(articles_processed)
                + "/"
                + str(links_t)
                + " (%"
                + str(articles_processed * 100 / links_t)
                + ")\n"
            )
        if len(links[article]) < 6:
            continue
        G.add_node(article)
        for l in links[article]:
            if (l in links) and (article in links[l]):  # back link is also present
                G.add_node(l)
                G.add_edge(article, l)

    networkx.write_gpickle(G, sys.argv[3])

def reduceGraph(read_g, write_g, minEdgeWeight, minNodeDegree, Lp, Sp):
    """
    Simplify the undirected graph and then update the 3 undirected weight properties.
    :param read_g: is the graph pickle to read
    :param write_g: is the updated graph pickle to write
    :param minEdgeWeight: the original weight of each edge should be >= minEdgeWeight
    :param minNodeDegree: the degree of each node should be >= minNodeDegree. the degree here is G.degree(node), NOT G.degree(node,weight='weight)
    :return: None
    """
    G=nx.read_gpickle(read_g)
    print 'number of original nodes: ', nx.number_of_nodes(G)
    print 'number of original edges: ', nx.number_of_edges(G)

    for (u,v,w) in G.edges(data='weight'):
        if w < minEdgeWeight:
            G.remove_edge(u,v)

    for n in G.nodes():
        if G.degree(n)<minNodeDegree:
            G.remove_node(n)

    print 'number of new nodes: ', nx.number_of_nodes(G)
    print 'number of new edges: ', nx.number_of_edges(G)

    for (a, b, w) in G.edges_iter(data='weight'):
        unweight_allocation(G, a, b, w,Lp,Sp)

    print 'update weight ok'
    nx.write_gpickle(G, write_g)

    return

def createBridge(numOfNodes, edgeProb, bridgeNodes):
	'''
	numOfNodes: Number of nodes in the clustered part of the Bridge Graph
	edgeProb: Probability of existance of an edge between any two vertices.
	bridgeNodes: Number of nodes in the bridge
	This function creates a Bridge Graph with 2 main clusters connected by a bridge.
	'''	
	print "Generating and Saving Bridge Network..."	
	G1 = nx.erdos_renyi_graph(2*numOfNodes + bridgeNodes, edgeProb) #Create an ER graph with number of vertices equal to twice the number of vertices in the clusters plus the number of bridge nodes.
	G = nx.Graph() #Create an empty graph so that it can be filled with the required components from G1
	G.add_edges_from(G1.subgraph(range(numOfNodes)).edges()) #Generate an induced subgraph of the nodes, ranging from 0 to numOfNodes, from G1 and add it to G
	G.add_edges_from(G1.subgraph(range(numOfNodes + bridgeNodes,2*numOfNodes + bridgeNodes)).edges()) #Generate an induced subgraph of the nodes, ranging from (numOfNodes + bridgeNodes) to (2*numOfNodes + bridgeNodes)

	A = random.randrange(numOfNodes) #Choose a random vertex from the first component
	B = random.randrange(numOfNodes + bridgeNodes,2*numOfNodes + bridgeNodes) #Choose a random vertex from the second component

	prev = A #creating a connection from A to B via the bridge nodes
	for i in range(numOfNodes, numOfNodes + bridgeNodes):
		G.add_edge(prev, i)
		prev = i
	G.add_edge(i, B)
	
	StrMap = {}
	for node in G.nodes():
		StrMap[node] = str(node)
	G = nx.convert.relabel_nodes(G,StrMap)
	filename = "BG_" + str(numOfNodes) + "_" + str(edgeProb) + "_" + str(bridgeNodes) + ".gpickle"
	nx.write_gpickle(G,filename)#generate a gpickle file of the learnt graph.
	print "Successfully written into " + filename

def create_graph_df(vtask_paths, graphs_dir_out):
    """
    Creates a frame that maps sourcefiles to networkx digraphs in terms of DOT files
    :param source_path_list:
    :param dest_dir_path:
    :param relabel:
    :return:
    """
    if not isdir(graphs_dir_out):
        raise ValueError('Invalid destination directory.')
    data = []
    graphgen_times = []

    print('Writing graph representations of verification tasks to {}'.format(graphs_dir_out), flush=True)

    common_prefix = commonprefix(vtask_paths)
    for vtask in tqdm(vtask_paths):
        short_prefix = dirname(common_prefix)
        path = join(graphs_dir_out, vtask[len(short_prefix):][1:])

        if not os.path.exists(dirname(path)):
            os.makedirs(dirname(path))

        ret_path = path + '.pickle'

        # DEBUG
        if isfile(ret_path):
            data.append(ret_path)
            continue

        start_time = time.time()

        graph_path, node_labels_path, edge_types_path, edge_truth_path, node_depths_path \
            = _run_cpachecker(abspath(vtask))
        nx_digraph = nx.read_graphml(graph_path)

        node_labels = _read_node_labeling(node_labels_path)
        nx.set_node_attributes(nx_digraph, 'label', node_labels)

        edge_types = _read_edge_labeling(edge_types_path)
        parsed_edge_types = _parse_edge(edge_types)
        nx.set_edge_attributes(nx_digraph, 'type', parsed_edge_types)

        edge_truth = _read_edge_labeling(edge_truth_path)
        parsed_edge_truth = _parse_edge(edge_truth)
        nx.set_edge_attributes(nx_digraph, 'truth', parsed_edge_truth)

        node_depths = _read_node_labeling(node_depths_path)
        parsed_node_depths = _parse_node_depth(node_depths)
        nx.set_node_attributes(nx_digraph, 'depth', parsed_node_depths)

        assert not isfile(ret_path)
        assert node_labels and parsed_edge_types and parsed_edge_truth and parsed_node_depths
        nx.write_gpickle(nx_digraph, ret_path)
        data.append(ret_path)

        gg_time = time.time() - start_time
        graphgen_times.append(gg_time)

    return pd.DataFrame({'graph_representation': data}, index=vtask_paths), graphgen_times

def load_data():
    start = time.time()
    try:
        print("Loading data from /data pickles and hfd5 adj matrices")
        f = h5py.File('data/cosponsorship_data.hdf5', 'r')
        for chamber in ['house', 'senate']:
            for congress in SUPPORTED_CONGRESSES:
                adj_matrix_lookup[(chamber, congress)] = np.asarray(f[chamber + str(congress)])

                igraph_graph = igraph.load("data/" + chamber + str(congress) + "_igraph.pickle", format="pickle")
                igraph_graph_lookup[(chamber, congress, False)] = igraph_graph

                nx_graph = nx.read_gpickle("data/" + chamber + str(congress) + "_nx.pickle")
                nx_graph_lookup[(chamber, congress, False)] = nx_graph
    except IOError as e:
        print("Loading data from cosponsorship files")
        f = h5py.File("data/cosponsorship_data.hdf5", "w")
        for chamber in ['house', 'senate']:
            for congress in SUPPORTED_CONGRESSES:
                print("Starting %s %s" % (str(congress), chamber))
                adj_matrix = load_adjacency_matrices(congress, chamber)
                data = f.create_dataset(chamber + str(congress), adj_matrix.shape, dtype='f')
                data[0: len(data)] = adj_matrix

                # igraph
                get_cosponsorship_graph(congress, chamber, False).save("data/" + chamber + str(congress) + "_igraph.pickle", "pickle")
                # networkx
                nx.write_gpickle(get_cosponsorship_graph_nx(congress, chamber, False), "data/" + chamber + str(congress) + "_nx.pickle")

                print("Done with %s %s" % (str(congress), chamber))
    print("Data loaded in %d seconds" % (time.time() - start))

    def save_celltype_graph(self, filename="celltype_conn.gml", format="gml"):
        """
        Save the celltype-to-celltype connectivity information in a file.
        
        filename -- path of the file to be saved.

        format -- format to save in. Using GML as GraphML support is
        not complete in NetworkX.  

        """
        start = datetime.now()
        if format == "gml":
            nx.write_gml(self.__celltype_graph, filename)
        elif format == "yaml":
            nx.write_yaml(self.__celltype_graph, filename)
        elif format == "graphml":
            nx.write_graphml(self.__celltype_graph, filename)
        elif format == "edgelist":
            nx.write_edgelist(self.__celltype_graph, filename)
        elif format == "pickle":
            nx.write_gpickle(self.__celltype_graph, filename)
        else:
            raise Exception("Supported formats: gml, graphml, yaml. Received: %s" % (format))
        end = datetime.now()
        delta = end - start
        config.BENCHMARK_LOGGER.info(
            "Saved celltype_graph in file %s of format %s in %g s"
            % (filename, format, delta.seconds + delta.microseconds * 1e-6)
        )
        print "Saved celltype connectivity graph in", filename

def graph_preprocessing_with_counts(G_input=None, save_file=None):

    if not G_input:
        graph_file = os.path.join(work_dir, "adj_graph.p")
        G = nx.read_gpickle(graph_file)
    else:
        G = G_input.copy()

    print "Raw graph size:", G.size()
    print "Raw graph nodes", G.number_of_nodes()

    profile2prob = {l.split()[0]: float(l.split()[1]) for l in open(os.path.join(work_dir, 'profile_weight.txt'))}

    for edge in G.edges(data=True):
        nodes = edge[:2]
        _weight = edge[2]['weight']
        _count = edge[2]['count']
        
        if _count < 3:
            G.remove_edge(*nodes)

    print "Pre-processed graph size", G.size()
    print "Pre-processed graph nodes", G.number_of_nodes()

    G.remove_nodes_from(nx.isolates(G))

    print "Pre-processed graph size", G.size()
    print "Pre-processed graph nodes", G.number_of_nodes()
    
    if save_file:
        print "Saving to", save_file
        nx.write_gpickle(G,save_file)

    return G

def store_graph(graph,name=None):
	filename=datetime.datetime.now().strftime('%Y-%m-%d_%H:%M:%S')+"_Network.gpickle"
	if name != None:
		filename=name+".gpickle"
	nx.write_gpickle(graph,filename)
	print("Finish storing the graph"+" "+filename)
	return filename

def createScaleFreeNetwork(numOfNodes, degree):
	'''
	numOfNodes: The number of nodes that the scale free network should have
	degree: The degree of the Scale Free Network
	This function creates a Scale Free Network containing 'numOfNodes' nodes, each of degree 'degree'
	It generates the required graph and saves it in a file. It runs the Reinforcement Algorithm to create a weightMatrix and an ordering of the vertices based on their importance by Flagging.
	'''
	global reinforce_time
	G = nx.barabasi_albert_graph(numOfNodes, degree) #Create a Scale Free Network of the given number of nodes and degree
	StrMap = {}
	for node in G.nodes():
		StrMap[node] = str(node)
	G = nx.convert.relabel_nodes(G,StrMap)

	print "Undergoing Machine Learning..."

	start = time.time()
	H = reinforce(G) #Enforce Machine Learning to generate a gml file of the learnt graph.
	finish = time.time()
	reinforce_time = finish - start

	print "Machine Learning Completed..."
	filename = "SFN_" + str(numOfNodes) + "_" + str(degree) + '.gpickle' 
	nx.write_gpickle(H,filename)#generate a gpickle file of the learnt graph.
	print "Learnt graph Successfully written into " + filename

def handle_by_file(out_dir, tweet_file, country, net_func=comprehend_network):
    try:
        if not os.path.exists(out_dir):
            os.makedirs(out_dir)

        if net_func == comprehend_network:
            net_type = "comprehend"
        elif net_func == user2user_network:
            net_type = "user2user"
        elif net_func == content_based_network:
            net_type = "content"
        elif net_func == entity_network:
            net_type = "entity"
        elif net_func == entity_corr_network:
            net_type = "entity_corr"
        else:
            net_type = "normal"

        net = net_func(tweet_file)
        net.graph["country"] = country
        g_date = re.search(r'\d{4}-\d{2}-\d{2}', tweet_file).group()
        net.graph["date"] = g_date
        out_file = os.path.join(out_dir,
                                "graph_%s_%s" % (net_type,
                                                 tweet_file.split(os.sep)[-1]))
        nx.write_gpickle(net, out_file + ".gpickle")
        nx.write_graphml(net, out_file + ".graphml")
    except Exception, e:
        print "Error Encoutered: %s, \n %s" \
            % (tweet_file, sys.exc_info()[0]), e

def fit_forestFire_mod(graphSize, graphID, dkPath, original2k, resultPath):
    """
    Runs synthetic graph tests for various 'p' values (burn rate).
    """
    outfile = open(resultPath + graphID + '_ff_dkDistances.txt', 'w')
    
    p = 0.01
    while p < 1.0:
        print 'Running modified Forest Fire with parameters: n = ', graphSize, ' p = ', p
        
        newFile = graphID + '_ff_' + str(p)

        # Create synthetic graph
        syntheticGraph = sm.forestFire_mod(graphSize, p)

        # Write pickle, edge list, and 2k distro to file
        print 'Writing pickle and calculating dK-2...\n'
        nx.write_gpickle(syntheticGraph, resultPath + newFile + '.pickle')
        getdk2(syntheticGraph, newFile, dkPath, resultPath)

        # Find distance between the dK-2 distributions
        dkDistance = tk.get_2k_distance(original2k, resultPath + newFile + '_target.2k')
        outfile.write(str(dkDistance) + '\tp = ' + str(p) + '\n')
        outfile.flush()

        p += 0.01

    outfile.close()

def main(**kwargs):
    cells = kwargs["cells"]
    frames = kwargs["frames"]
    if isinstance(cells, str):
        cells = np.load(cells)
    if isinstance(frames, str):
        frames = np.load(frames)

    logger.info("Creating correlation graph")
    N = int(cells.max())
    for i in range(1, N):
        logger.info("\tDone %d out of %d" % (i, N))
        indices = list(zip(*np.where(cells == i)))
        if len(indices) < 2:
            continue
        pixals = []
        for y, x in indices:
            pixals.append(frames[x, y, :])
        pixals = np.mean(pixals, axis=0)
        g_.add_node(i, timeseries=pixals, indices=indices)

    g_.graph["shape"] = frames[:, :, 0].shape
    create_correlate_graph(g_)
    outfile = kwargs.get("output", False) or "correlation_graph.pickle"
    logger.info("Writing pickle of graph to %s" % outfile)
    nx.write_gpickle(g_, outfile)
    logger.info("Graph pickle is saved to %s" % outfile)

def create_nodes(paths, args):
    """ creates nodes
    Parameters
    ----------
    paths.node_file       : file
    args.fasta_file       : file
    """

    # read in fasta to dictionary.
    seqs = io.load_fasta(args.contig_file)

    # create graph.
    G = nx.MultiGraph()

    # add nodes to graph.
    for name, seq in seqs.items():

        # skip split names.
        tmp = name.split(" ")
        name = tmp[0]

        # add node.
        G.add_node(name, {'seq':seq, 'width':len(seq), 'cov':0})

    # write to disk.
    nx.write_gpickle(G, paths.node_file)

def buildGraph(db):

    DATA_DIR = os.environ['OPENSHIFT_DATA_DIR']
    file = os.path.join(DATA_DIR, "UserTagsGraph.gpickle")

    if not os.path.isfile(file):
        user_network = nx.Graph()
        users = DAL.UserTags.getAll(db)
        user_tags = {}
        all_tags = {}

        for user in users:
            tags = user.tags_to_list()
            user_tags[user.id] = tags
            for tag in tags:
                all_tags.setdefault(tag, set()).add(user.id)

        for tag in all_tags:
            user_network.add_edges_from([perm for perm in itertools.permutations(all_tags[tag], 2)])

        #save graph to file
        nx.write_gpickle(user_network, file)
    else:
        user_network = nx.read_gpickle(file)

    return user_network

def G_init(slim_dict,slim):
    def check_field(pos,G,rate):
        for node in G:
            node_pos = G.node[node]["pos"]
            if math.sqrt(np.sum((pos-node_pos)**2)) < rate:
                return True
        return False

    if len(slim_dict) < 25:
        rate = 0.42
    elif len(slim_dict) < 42:
        rate = 0.3
    else:
        rate = 0.25

    G = nx.Graph()
    for node in slim_dict:
        pos = np.array((0.0,0.0))
        while check_field(pos,G,rate):
            pos = np.array((2.8*random.random()-1.4, 2.8*random.random()-1.4))
        G.add_node(node,size=len(GOID2group[node]),color=0,pos=pos)
    for node in G:
        G.node[node]["sum"] = 0
        G.node[node]["1"] = 0
        G.node[node]["2"] = 0
    for GOID_1 in slim_dict:
        for GOID_2 in GO2interact[GOID_1]:
            if GOID_2 != GOID_1 and GOID_2 in slim_dict:
                G.add_edge(GOID_1,GOID_2,weight=GO2interact[GOID_1][GOID_2],percent=0.5,percent_1=0.5,percent_2=0.5)
    png = path+"/results/slim_"+slim+"/slim_"+slim+".png"
    G = nx.relabel_nodes(G,slim_dict)
    nx.write_gpickle(G,path+"/results/G_slim_"+slim)
    draw_G(G,png)

def better_display(G,slim,pic,fit=None,pngpath=None,bw=None):
    import engine.Genetic as Genetic
    _fit = "C" if fit == "C" else "L"
    R = Genetic.genetic(G,_fit,pngpath,bw)
    for node in G:
        G.node[node]["pos"] = R[node]
    nx.write_gpickle(G,pic)

def addNodeDe_EdgeDist():
    """
    Add node degree and edge distance on the filtered Graph
    :return: Graph
    """
    schema = 'total_v3_csvneo4j'
    Graph_type = 'undirected'
    alpha_thred = 0.65
    nodeDegree_thred = 1.0
    DisTypes = ['G','SP','R']

    G = nx.read_gpickle('../filteredG_{}_alpha{}_nodeD{}_{}.gpickle'.format(Graph_type,alpha_thred,nodeDegree_thred,schema))
    print 'after read'
    print 'edges: ', len(G.edges())
    print 'nodes: ', len(G.nodes())
    G = main.addNode_degree(G)
    print 'finish adding node degree'
    G = main.addEdge_distance(G,DisTypes)
    print 'finish adding edge degree'

    nx.write_gpickle(G,'../addNodeEdgeDegree_{}_{}_alpha{}_nodeD{}_{}.gpickle'.format('+'.join(DisTypes),Graph_type,alpha_thred,nodeDegree_thred,schema))
    print 'finishing write gpickle'
    print 'edges: ', len(G.edges())
    print 'nodes: ', len(G.nodes())

    return

def main():
    if not (len(sys.argv) == 2 and direction in ["forward, backward"]):
        print("usage: ./gen_graph.py [forward/backward]", file=sys.stderr)
        sys.exit(1)

    direction = sys.argv[1]
    if direction == "forward":
        f = roundf
    else:
        f = inv_roundf
    n = 65536

    g = nx.DiGraph()
    for x in range(n):
        for ns, w in f(convert_int(x)):
            y = convert_states(ns)
            g.add_edge(x, y, weight=w)
        print(x)
    nx.write_gpickle(g, "{}.gpickle".format(direction))

    print("Generated {}.gpickle.".format(direction))

    nx.reverse(g, copy=False)
    nx.write_gpickle(g, "rev_{}.gpickle".format(direction))

    print("Generated rev_{}.gpickle.".format(direction))