import itertools


def vertex_color_qubo(G, colors):
    """Return the QUBO with ground states corresponding to a vertex coloring.

    If `V` is the set of nodes, `E` is the set of edges and `C` is the set of
    colors the resulting qubo will have:

    * :math:`|V|*|C|` variables/nodes
    * :math:`|V|*|C|*(|C| - 1) / 2 + |E|*|C|` interactions/edges

    The QUBO has ground energy :math:`-|V|` and an infeasible gap of 1.

    Parameters
    ----------
    G : NetworkX graph
        The graph on which to find a minimum vertex coloring.

    colors : int/sequence
        The colors. If an int, the colors are labelled `[0, n)`. The number of
        colors must be greater or equal to the chromatic number of the graph.

    Returns
    -------
    QUBO : dict
        The QUBO with ground states corresponding to valid colorings of the
        graph. The QUBO variables are labelled `(v, c)` where `v` is a node
        in `G` and `c` is a color. In the ground state of the QUBO, a variable
        `(v, c)` has value 1 if `v` should be colored `c` in a valid coloring.


    """
    _, colors = colors

    Q = {}

    # enforce that each variable in G has at most one color
    for v in G.nodes:
        # 1 in k constraint
        for c in colors:
            Q[(v, c), (v, c)] = -1

        for c0, c1 in itertools.combinations(colors, 2):
            Q[(v, c0), (v, c1)] = 2

    # enforce that adjacent nodes do not have the same color
    for u, v in G.edges:
        # NAND constraint
        for c in colors:
            Q[(u, c), (v, c)] = 1

    return Q