import pandas as pd
import networkx as nx
import matplotlib.pyplot as plt

df = pd.read_csv('data\\temp.csv')
df['succ_equipment_code'] = df['equipment_code'].str.rsplit('.', n=1)
df['succ_equipment_code'] = df['succ_equipment_code'].apply(lambda lst: lst[0] if len(lst)>1 else None)

assert len(df['equipment_code']) == len((df['equipment_code'].unique())),\
    "column equipment_code contains duplicate value"
df.set_index('equipment_code', inplace=True)

# graph bom
g_bom = nx.DiGraph()
g_bom.add_nodes_from(df.index)

# edge
for index, row in df.iterrows():
    if row['succ_equipment_code']:
        g_bom.add_edge(index, row['succ_equipment_code'])

# attr
df_attr = df.loc[:, ['equipment_name']]
dct_attr = {}
for code in df_attr.index:
    dct_attr[code] = df_attr.loc[code].to_dict()
nx.set_node_attributes(g_bom, dct_attr)

# show
plt.rcParams['font.sans-serif'] = 'SimHei'
plt.figure(figsize=(12, 12), dpi=300)
pos = nx.nx_agraph.graphviz_layout(g_bom,
                                   prog='twopi',
                                   args='')
node_label = nx.get_node_attributes(g_bom, 'equipment_name')
nx.draw(g_bom,
        pos,
        labels=node_label,
        font_size=6)
plt.savefig('g_bom_20230712.png')
plt.close()

# graph firm
s_firm = pd.Series(df['firm_name'].unique())
g_firm = nx.MultiDiGraph()
g_firm.add_nodes_from(s_firm.index)

# attr
dct_attr = {}
for firm_code in s_firm.index:
    firm_name = s_firm[firm_code]
    lst_prod_code = df.index[df['firm_name'] == firm_name].to_list()
    dct_attr[firm_code] = {'firm_name': firm_name,
                           'lst_prod_code': lst_prod_code}
nx.set_node_attributes(g_firm, dct_attr)

print(g_firm.nodes)
print(g_firm.nodes[0])
# edge
for prod_code, row in df.iterrows():
    firm_name = row['firm_name']
    firm_code = s_firm.index[s_firm == firm_name].to_list()[0]
    for succ_prod_code in g_bom.successors(prod_code):
        # 目前df中一个prod code只对应一个firm
        succ_firm_name = df.loc[df.index == succ_prod_code, 'firm_name'].to_list()[0]
        succ_firm_code = s_firm.index[s_firm == succ_firm_name].to_list()[0]
        g_firm.add_edges_from([(firm_code, succ_firm_code, {'product': prod_code})])
        print((firm_name, succ_firm_name, prod_code))

# show
plt.rcParams['font.sans-serif'] = 'SimHei'
plt.figure(figsize=(12, 12), dpi=300)
pos = nx.nx_agraph.graphviz_layout(g_firm,
                                   prog='twopi',
                                   args='')
node_label = nx.get_node_attributes(g_firm, 'firm_name')
nx.draw(g_firm,
        pos,
        labels=node_label,
        font_size=6)
edge_label = nx.get_edge_attributes(g_firm, 'product')
edge_label = {(n1, n2): label for (n1, n2, _), label in edge_label.items()}
nx.draw_networkx_edge_labels(g_firm,
                             pos,
                             edge_label,
                             font_size=4)
plt.savefig('g_firm_20230712.png')
plt.close()

# pos = nx.random_layout(g_firm)
# nx.draw_networkx_nodes(g_firm, pos, node_color = 'r', node_size = 100, alpha = 1)
# ax = plt.gca()
# for e in g_firm.edges:
#     ax.annotate("",
#                 xy=pos[e[0]], xycoords='data',
#                 xytext=pos[e[1]], textcoords='data',
#                 arrowprops=dict(arrowstyle="->", color="0.5",
#                                 shrinkA=5, shrinkB=5,
#                                 patchA=None, patchB=None,
#                                 connectionstyle="arc3,rad=rrr".replace('rrr',str(0.3*e[2])
#                                 ),
#                                 ),
#                 )
# plt.axis('off')
# plt.show()