"""
Multipartite layout
===================

This example from networkx shows how to use simple vertex styling.
"""

import itertools
import matplotlib.pyplot as plt
import networkx as nx
import iplotx as ipx

subset_sizes = [5, 5, 4, 3, 2, 4, 4, 3]
subset_color = [
    "gold",
    "violet",
    "violet",
    "violet",
    "violet",
    "limegreen",
    "limegreen",
    "darkorange",
]


def multilayered_graph(*subset_sizes):
    extents = nx.utils.pairwise(itertools.accumulate((0,) + subset_sizes))
    layers = [range(start, end) for start, end in extents]
    G = nx.Graph()
    for i, layer in enumerate(layers):
        G.add_nodes_from(layer, layer=i)
    for layer1, layer2 in nx.utils.pairwise(layers):
        G.add_edges_from(itertools.product(layer1, layer2))
    return G


G = multilayered_graph(*subset_sizes)
color = [subset_color[data["layer"]] for v, data in G.nodes(data=True)]
pos = nx.multipartite_layout(G, subset_key="layer")

fig, ax = plt.subplots(figsize=(6, 4))
ipx.plot(
    G,
    layout=pos,
    style={
        "vertex": {
            "facecolor": color,
        },
    },
    ax=ax,
)
