# -*- coding: utf-8 -*-
#    Copyright (C) 2006-2011 by
#    Aric Hagberg <hagberg@lanl.gov>
#    Dan Schult <dschult@colgate.edu>
#    Pieter Swart <swart@lanl.gov>
#    All rights reserved.
#    BSD license.
#
# Authors: Aric Hagberg <aric.hagberg@gmail.com>
#                        Pieter Swart <swart@lanl.gov>
#                        Dan Schult <dschult@colgate.edu>
"""
Generators and functions for bipartite graphs.
"""
import math
import numbers
import random
from functools import reduce
import networkx as nx
from networkx.utils import nodes_or_number, py_random_state

__all__ = ['configuration_model',
           'havel_hakimi_graph',
           'reverse_havel_hakimi_graph',
           'alternating_havel_hakimi_graph',
           'preferential_attachment_graph',
           'random_graph',
           'gnmk_random_graph',
           'complete_bipartite_graph',
           ]


@nodes_or_number([0, 1])
def complete_bipartite_graph(n1, n2, create_using=None):
    """Return the complete bipartite graph `K_{n_1,n_2}`.

    Composed of two partitions with `n_1` nodes in the first
    and `n_2` nodes in the second. Each node in the first is
    connected to each node in the second.

    Parameters
    ----------
    n1 : integer
       Number of nodes for node set A.
    n2 : integer
       Number of nodes for node set B.
    create_using : NetworkX graph instance, optional
       Return graph of this type.

    Notes
    -----
    Node labels are the integers 0 to `n_1 + n_2 - 1`.

    The nodes are assigned the attribute 'bipartite' with the value 0 or 1
    to indicate which bipartite set the node belongs to.
    """
    G = nx.empty_graph(0, create_using)
    if G.is_directed():
        raise nx.NetworkXError("Directed Graph not supported")

    n1, top = n1
    n2, bottom = n2
    if isinstance(n2, numbers.Integral):
        bottom = [n1 + i for i in bottom]
    G.add_nodes_from(top, bipartite=0)
    G.add_nodes_from(bottom, bipartite=1)
    G.add_edges_from((u, v) for u in top for v in bottom)
    G.graph['name'] = "complete_bipartite_graph(%s,%s)" % (n1, n2)
    return G


@py_random_state(3)
def configuration_model(aseq, bseq, create_using=None, seed=None):
    """Return a random bipartite graph from two given degree sequences.

    Parameters
    ----------
    aseq : list
       Degree sequence for node set A.
    bseq : list
       Degree sequence for node set B.
    create_using : NetworkX graph instance, optional
       Return graph of this type.
    seed : integer, random_state, or None (default)
        Indicator of random number generation state.
        See :ref:`Randomness<randomness>`.

    Nodes from the set A are connected to nodes in the set B by
    choosing randomly from the possible free stubs, one in A and
    one in B.

    Notes
    -----
    The sum of the two sequences must be equal: sum(aseq)=sum(bseq)
    If no graph type is specified use MultiGraph with parallel edges.
    If you want a graph with no parallel edges use create_using=Graph()
    but then the resulting degree sequences might not be exact.

    The nodes are assigned the attribute 'bipartite' with the value 0 or 1
    to indicate which bipartite set the node belongs to.

    This function is not imported in the main namespace.
    To use it you have to explicitly import the bipartite package.
    """
    G = nx.empty_graph(0, create_using, default=nx.MultiGraph)
    if G.is_directed():
        raise nx.NetworkXError("Directed Graph not supported")

    # length and sum of each sequence
    lena = len(aseq)
    lenb = len(bseq)
    suma = sum(aseq)
    sumb = sum(bseq)

    if not suma == sumb:
        raise nx.NetworkXError(
            'invalid degree sequences, sum(aseq)!=sum(bseq),%s,%s'
            % (suma, sumb))

    G = _add_nodes_with_bipartite_label(G, lena, lenb)

    if len(aseq) == 0 or max(aseq) == 0:
        return G  # done if no edges

    # build lists of degree-repeated vertex numbers
    stubs = []
    stubs.extend([[v] * aseq[v] for v in range(0, lena)])
    astubs = []
    astubs = [x for subseq in stubs for x in subseq]

    stubs = []
    stubs.extend([[v] * bseq[v - lena] for v in range(lena, lena + lenb)])
    bstubs = []
    bstubs = [x for subseq in stubs for x in subseq]

    # shuffle lists
    seed.shuffle(astubs)
    seed.shuffle(bstubs)

    G.add_edges_from([[astubs[i], bstubs[i]] for i in range(suma)])

    G.name = "bipartite_configuration_model"
    return G


def havel_hakimi_graph(aseq, bseq, create_using=None):
    """Return a bipartite graph from two given degree sequences using a
    Havel-Hakimi style construction.

    Nodes from the set A are connected to nodes in the set B by
    connecting the highest degree nodes in set A to the highest degree
    nodes in set B until all stubs are connected.

    Parameters
    ----------
    aseq : list
       Degree sequence for node set A.
    bseq : list
       Degree sequence for node set B.
    create_using : NetworkX graph instance, optional
       Return graph of this type.

    Notes
    -----
    This function is not imported in the main namespace.
    To use it you have to explicitly import the bipartite package.

    The sum of the two sequences must be equal: sum(aseq)=sum(bseq)
    If no graph type is specified use MultiGraph with parallel edges.
    If you want a graph with no parallel edges use create_using=Graph()
    but then the resulting degree sequences might not be exact.

    The nodes are assigned the attribute 'bipartite' with the value 0 or 1
    to indicate which bipartite set the node belongs to.
    """
    G = nx.empty_graph(0, create_using, default=nx.MultiGraph)
    if G.is_directed():
        raise nx.NetworkXError("Directed Graph not supported")

    # length of the each sequence
    naseq = len(aseq)
    nbseq = len(bseq)

    suma = sum(aseq)
    sumb = sum(bseq)

    if not suma == sumb:
        raise nx.NetworkXError(
            'invalid degree sequences, sum(aseq)!=sum(bseq),%s,%s'
            % (suma, sumb))

    G = _add_nodes_with_bipartite_label(G, naseq, nbseq)

    if len(aseq) == 0 or max(aseq) == 0:
        return G  # done if no edges

    # build list of degree-repeated vertex numbers
    astubs = [[aseq[v], v] for v in range(0, naseq)]
    bstubs = [[bseq[v - naseq], v] for v in range(naseq, naseq + nbseq)]
    astubs.sort()
    while astubs:
        (degree, u) = astubs.pop()  # take of largest degree node in the a set
        if degree == 0:
            break  # done, all are zero
        # connect the source to largest degree nodes in the b set
        bstubs.sort()
        for target in bstubs[-degree:]:
            v = target[1]
            G.add_edge(u, v)
            target[0] -= 1  # note this updates bstubs too.
            if target[0] == 0:
                bstubs.remove(target)

    G.name = "bipartite_havel_hakimi_graph"
    return G


def reverse_havel_hakimi_graph(aseq, bseq, create_using=None):
    """Return a bipartite graph from two given degree sequences using a
    Havel-Hakimi style construction.

    Nodes from set A are connected to nodes in the set B by connecting
    the highest degree nodes in set A to the lowest degree nodes in
    set B until all stubs are connected.

    Parameters
    ----------
    aseq : list
       Degree sequence for node set A.
    bseq : list
       Degree sequence for node set B.
    create_using : NetworkX graph instance, optional
       Return graph of this type.

    Notes
    -----
    This function is not imported in the main namespace.
    To use it you have to explicitly import the bipartite package.

    The sum of the two sequences must be equal: sum(aseq)=sum(bseq)
    If no graph type is specified use MultiGraph with parallel edges.
    If you want a graph with no parallel edges use create_using=Graph()
    but then the resulting degree sequences might not be exact.

    The nodes are assigned the attribute 'bipartite' with the value 0 or 1
    to indicate which bipartite set the node belongs to.
    """
    G = nx.empty_graph(0, create_using, default=nx.MultiGraph)
    if G.is_directed():
        raise nx.NetworkXError("Directed Graph not supported")

    # length of the each sequence
    lena = len(aseq)
    lenb = len(bseq)
    suma = sum(aseq)
    sumb = sum(bseq)

    if not suma == sumb:
        raise nx.NetworkXError(
            'invalid degree sequences, sum(aseq)!=sum(bseq),%s,%s'
            % (suma, sumb))

    G = _add_nodes_with_bipartite_label(G, lena, lenb)

    if len(aseq) == 0 or max(aseq) == 0:
        return G  # done if no edges

    # build list of degree-repeated vertex numbers
    astubs = [[aseq[v], v] for v in range(0, lena)]
    bstubs = [[bseq[v - lena], v] for v in range(lena, lena + lenb)]
    astubs.sort()
    bstubs.sort()
    while astubs:
        (degree, u) = astubs.pop()  # take of largest degree node in the a set
        if degree == 0:
            break  # done, all are zero
        # connect the source to the smallest degree nodes in the b set
        for target in bstubs[0:degree]:
            v = target[1]
            G.add_edge(u, v)
            target[0] -= 1  # note this updates bstubs too.
            if target[0] == 0:
                bstubs.remove(target)

    G.name = "bipartite_reverse_havel_hakimi_graph"
    return G


def alternating_havel_hakimi_graph(aseq, bseq, create_using=None):
    """Return a bipartite graph from two given degree sequences using
    an alternating Havel-Hakimi style construction.

    Nodes from the set A are connected to nodes in the set B by
    connecting the highest degree nodes in set A to alternatively the
    highest and the lowest degree nodes in set B until all stubs are
    connected.

    Parameters
    ----------
    aseq : list
       Degree sequence for node set A.
    bseq : list
       Degree sequence for node set B.
    create_using : NetworkX graph instance, optional
       Return graph of this type.

    Notes
    -----
    This function is not imported in the main namespace.
    To use it you have to explicitly import the bipartite package.

    The sum of the two sequences must be equal: sum(aseq)=sum(bseq)
    If no graph type is specified use MultiGraph with parallel edges.
    If you want a graph with no parallel edges use create_using=Graph()
    but then the resulting degree sequences might not be exact.

    The nodes are assigned the attribute 'bipartite' with the value 0 or 1
    to indicate which bipartite set the node belongs to.
    """
    G = nx.empty_graph(0, create_using, default=nx.MultiGraph)
    if G.is_directed():
        raise nx.NetworkXError("Directed Graph not supported")

    # length of the each sequence
    naseq = len(aseq)
    nbseq = len(bseq)
    suma = sum(aseq)
    sumb = sum(bseq)

    if not suma == sumb:
        raise nx.NetworkXError(
            'invalid degree sequences, sum(aseq)!=sum(bseq),%s,%s'
            % (suma, sumb))

    G = _add_nodes_with_bipartite_label(G, naseq, nbseq)

    if len(aseq) == 0 or max(aseq) == 0:
        return G  # done if no edges
    # build list of degree-repeated vertex numbers
    astubs = [[aseq[v], v] for v in range(0, naseq)]
    bstubs = [[bseq[v - naseq], v] for v in range(naseq, naseq + nbseq)]
    while astubs:
        astubs.sort()
        (degree, u) = astubs.pop()  # take of largest degree node in the a set
        if degree == 0:
            break  # done, all are zero
        bstubs.sort()
        small = bstubs[0:degree // 2]  # add these low degree targets
        large = bstubs[(-degree + degree // 2):]  # and these high degree targets
        stubs = [x for z in zip(large, small) for x in z]  # combine, sorry
        if len(stubs) < len(small) + len(large):  # check for zip truncation
            stubs.append(large.pop())
        for target in stubs:
            v = target[1]
            G.add_edge(u, v)
            target[0] -= 1  # note this updates bstubs too.
            if target[0] == 0:
                bstubs.remove(target)

    G.name = "bipartite_alternating_havel_hakimi_graph"
    return G


@py_random_state(3)
def preferential_attachment_graph(aseq, p, create_using=None, seed=None):
    """Create a bipartite graph with a preferential attachment model from
    a given single degree sequence.

    Parameters
    ----------
    aseq : list
       Degree sequence for node set A.
    p :  float
       Probability that a new bottom node is added.
    create_using : NetworkX graph instance, optional
       Return graph of this type.
    seed : integer, random_state, or None (default)
        Indicator of random number generation state.
        See :ref:`Randomness<randomness>`.

    References
    ----------
    .. [1] Jean-Loup Guillaume and Matthieu Latapy,
       Bipartite structure of all complex networks,
       Inf. Process. Lett. 90, 2004, pg. 215-221
       https://doi.org/10.1016/j.ipl.2004.03.007

    Notes
    -----

    This function is not imported in the main namespace.
    To use it you have to explicitly import the bipartite package.
    """
    G = nx.empty_graph(0, create_using, default=nx.MultiGraph)
    if G.is_directed():
        raise nx.NetworkXError("Directed Graph not supported")

    if p > 1:
        raise nx.NetworkXError("probability %s > 1" % (p))

    naseq = len(aseq)
    G = _add_nodes_with_bipartite_label(G, naseq, 0)
    vv = [[v] * aseq[v] for v in range(0, naseq)]
    while vv:
        while vv[0]:
            source = vv[0][0]
            vv[0].remove(source)
            if seed.random() < p or G.number_of_nodes() == naseq:
                target = G.number_of_nodes()
                G.add_node(target, bipartite=1)
                G.add_edge(source, target)
            else:
                bb = [[b] * G.degree(b) for b in range(naseq, G.number_of_nodes())]
                # flatten the list of lists into a list.
                bbstubs = reduce(lambda x, y: x + y, bb)
                # choose preferentially a bottom node.
                target = seed.choice(bbstubs)
                G.add_node(target, bipartite=1)
                G.add_edge(source, target)
        vv.remove(vv[0])
    G.name = "bipartite_preferential_attachment_model"
    return G


@py_random_state(3)
def random_graph(n, m, p, seed=None, directed=False):
    """Return a bipartite random graph.

    This is a bipartite version of the binomial (Erdős-Rényi) graph.

    Parameters
    ----------
    n : int
        The number of nodes in the first bipartite set.
    m : int
        The number of nodes in the second bipartite set.
    p : float
        Probability for edge creation.
    seed : integer, random_state, or None (default)
        Indicator of random number generation state.
        See :ref:`Randomness<randomness>`.
    directed : bool, optional (default=False)
        If True return a directed graph

    Notes
    -----
    This function is not imported in the main namespace.
    To use it you have to explicitly import the bipartite package.

    The bipartite random graph algorithm chooses each of the n*m (undirected)
    or 2*nm (directed) possible edges with probability p.

    This algorithm is $O(n+m)$ where $m$ is the expected number of edges.

    The nodes are assigned the attribute 'bipartite' with the value 0 or 1
    to indicate which bipartite set the node belongs to.

    See Also
    --------
    gnp_random_graph, configuration_model

    References
    ----------
    .. [1] Vladimir Batagelj and Ulrik Brandes,
       "Efficient generation of large random networks",
       Phys. Rev. E, 71, 036113, 2005.
    """
    G = nx.Graph()
    G = _add_nodes_with_bipartite_label(G, n, m)
    if directed:
        G = nx.DiGraph(G)
    G.name = "fast_gnp_random_graph(%s,%s,%s)" % (n, m, p)

    if p <= 0:
        return G
    if p >= 1:
        return nx.complete_bipartite_graph(n, m)

    lp = math.log(1.0 - p)

    v = 0
    w = -1
    while v < n:
        lr = math.log(1.0 - seed.random())
        w = w + 1 + int(lr / lp)
        while w >= m and v < n:
            w = w - m
            v = v + 1
        if v < n:
            G.add_edge(v, n + w)

    if directed:
        # use the same algorithm to
        # add edges from the "m" to "n" set
        v = 0
        w = -1
        while v < n:
            lr = math.log(1.0 - seed.random())
            w = w + 1 + int(lr / lp)
            while w >= m and v < n:
                w = w - m
                v = v + 1
            if v < n:
                G.add_edge(n + w, v)

    return G


@py_random_state(3)
def gnmk_random_graph(n, m, k, seed=None, directed=False):
    """Return a random bipartite graph G_{n,m,k}.

    Produces a bipartite graph chosen randomly out of the set of all graphs
    with n top nodes, m bottom nodes, and k edges.

    Parameters
    ----------
    n : int
        The number of nodes in the first bipartite set.
    m : int
        The number of nodes in the second bipartite set.
    k : int
        The number of edges
    seed : integer, random_state, or None (default)
        Indicator of random number generation state.
        See :ref:`Randomness<randomness>`.
    directed : bool, optional (default=False)
        If True return a directed graph

    Examples
    --------
    from nx.algorithms import bipartite
    G = bipartite.gnmk_random_graph(10,20,50)

    See Also
    --------
    gnm_random_graph

    Notes
    -----
    This function is not imported in the main namespace.
    To use it you have to explicitly import the bipartite package.

    If k > m * n then a complete bipartite graph is returned.

    This graph is a bipartite version of the `G_{nm}` random graph model.
    """
    G = nx.Graph()
    G = _add_nodes_with_bipartite_label(G, n, m)
    if directed:
        G = nx.DiGraph(G)
    G.name = "bipartite_gnm_random_graph(%s,%s,%s)" % (n, m, k)
    if n == 1 or m == 1:
        return G
    max_edges = n * m  # max_edges for bipartite networks
    if k >= max_edges:  # Maybe we should raise an exception here
        return nx.complete_bipartite_graph(n, m, create_using=G)

    top = [n for n, d in G.nodes(data=True) if d['bipartite'] == 0]
    bottom = list(set(G) - set(top))
    edge_count = 0
    while edge_count < k:
        # generate random edge,u,v
        u = seed.choice(top)
        v = seed.choice(bottom)
        if v in G[u]:
            continue
        else:
            G.add_edge(u, v)
            edge_count += 1
    return G


def _add_nodes_with_bipartite_label(G, lena, lenb):
    G.add_nodes_from(range(0, lena + lenb))
    b = dict(zip(range(0, lena), [0] * lena))
    b.update(dict(zip(range(lena, lena + lenb), [1] * lenb)))
    nx.set_node_attributes(G, b, 'bipartite')
    return G
