"""
This code provides an example on how to plot a network graph with the igraph object generated by the `network_p2p` function.
Use `return_type = 'network'` to generate the igraph object, and use functions from matplotlib, igraph, and seaborn to generate a network visualization.
"""
import seaborn as sns
import vivainsights as vi
import matplotlib.pyplot as plt
import igraph as ig
import matplotlib.patches as mpatches

# Adjust placeholder text to your needs
hrvar_text = 'Organization'

# Generate graph
g = vi.network_p2p(
    data = vi.p2p_data_sim(), # Replace with your own P2P data
    hrvar = hrvar_text,
    return_type = "network"
    )

# Remove loops and edges
# g = g.simplify()

# Unique values in hrvar_text
unique_values = len(set(g.vs[hrvar_text]))
unique_nodes = set(g.vs[hrvar_text])

# Generate a list of colors from the 'hsv' palette, based on unique values
colors = sns.color_palette('hsv', unique_values)

# Create a dictionary that maps each unique org to a unique color
org_to_color = {org: color for org, color in zip(set(g.vs[hrvar_text]), colors)}

# Create a new vertex attribute that contains the corresponding color for each vertex
g.vs["node_color"] = [org_to_color[org] for org in g.vs[hrvar_text]]

# Create legend handles
legend_handles = [mpatches.Patch(color=org_to_color[org], label=org) for org in unique_nodes]

# Start plotting
fig, ax = plt.subplots(figsize=(8, 8))

g.vs["node_size"] = [x*120 for x in g.vs["node_size"]] # scale the size of the nodes

ig.plot(
    g,
    layout = g.layout("mds"),
    target=ax,
    vertex_label = None,
    vertex_size = g.vs["node_size"],
    vertex_color = g.vs["node_color"],
    edge_arrow_mode = "0",
    edge_arrow_size=0,
    edge_color = "#adadad",
    edge_width = 0.1
)

# Create the legend
plt.legend(
    frameon = False,
    markerscale = 1,
    fontsize= 5,
    handles = legend_handles,
    labelcolor = 'grey',
    ncols = 2
)

plt.show()