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

plt.rcParams['font.sans-serif'] = 'SimHei'

count_prod = pd.read_csv("output_result/risk/count_prod.csv")
print(count_prod)

# category
print(count_prod.describe())

# prod_networkx
BomNodes = pd.read_csv('input_data/input_product_data/BomNodes.csv', index_col=0)
BomNodes.set_index('Code', inplace=True)
BomCateNet = pd.read_csv('input_data/input_product_data/BomCateNet.csv', index_col=0)
BomCateNet.fillna(0, inplace=True)

G = nx.from_pandas_adjacency(BomCateNet.T, create_using=nx.MultiDiGraph())

labels_dict = {}
for code in G.nodes:
    node_attr = BomNodes.loc[code].to_dict()
    index_list = count_prod[count_prod['id_product'] == code].index.tolist()
    index = index_list[0] if len(index_list) == 1 else -1
    node_attr['count'] = count_prod['count'].get(index, 0)
    node_attr['node_size'] = count_prod['count'].get(index, 0)
    node_attr['node_color'] = count_prod['count'].get(index, 0)
    labels_dict[code] = node_attr
nx.set_node_attributes(G, labels_dict)
# print(labels_dict)

pos = nx.nx_agraph.graphviz_layout(G, prog="twopi", args="")
dict_node_name = nx.get_node_attributes(G, 'Name')
node_labels = {}
for node in nx.nodes(G):
    node_labels[node] = f"{node} {str(dict_node_name[node])}"
    # node_labels[node] = f"{str(dict_node_name[node])}"
colors = list(nx.get_node_attributes(G, 'node_color').values())
vmin = min(colors)
vmax = max(colors)
cmap = plt.cm.Blues
fig = plt.figure(figsize=(10, 10), dpi=300)
nx.draw(G,
        pos,
        node_size=list(nx.get_node_attributes(G, 'node_size').values()),
        labels=node_labels,
        font_size=6,
        node_color=colors,
        cmap=cmap,
        vmin=vmin,
        vmax=vmax,
        edge_color='grey')
sm = plt.cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(vmin=vmin, vmax=vmax))
sm._A = []
position = fig.add_axes([0.01, 0.05, 0.01, 0.3])
cb = plt.colorbar(sm, fraction=0.01, cax=position)
cb.ax.tick_params(labelsize=8)
cb.outline.set_visible(False)
plt.savefig("output_result\\risk\\count_prod_network")
plt.close()

# dcp_prod
count_dcp = pd.read_csv("output_result/risk/count_dcp.csv",
                        dtype={
                            'up_id_firm': str,
                            'down_id_firm': str
                        })
count_dcp_prod = count_dcp.groupby(
    ['up_id_product',
     'down_id_product'])['count'].sum()
count_dcp_prod = count_dcp_prod.reset_index()
count_dcp_prod.sort_values('count', inplace=True, ascending=False)
count_dcp_prod.to_csv('output_result\\risk\\count_dcp_prod.csv',
                      index=False,
                      encoding='utf-8-sig')
count_dcp_prod = count_dcp_prod[count_dcp_prod['count'] > 50]
# print(count_dcp_prod)

list_prod = count_dcp_prod['up_id_product'].tolist(
) + count_dcp['down_id_product'].tolist()
list_prod = list(set(list_prod))

# init graph bom

BomNodes = pd.read_csv('input_data/input_product_data/BomNodes.csv', index_col=0)
BomNodes.set_index('Code', inplace=True)

g_bom = nx.MultiDiGraph()
g_bom.add_nodes_from(list_prod)

bom_labels_dict = {}
for code in list_prod:
    dct_attr = BomNodes.loc[code].to_dict()
    bom_labels_dict[code] = dct_attr
nx.set_node_attributes(g_bom, bom_labels_dict)


count_max = count_dcp_prod['count'].max()
count_min = count_dcp_prod['count'].min()
k = 5 / (count_max - count_min)
for _, row in count_dcp_prod.iterrows():
    # print(row)
    lst_add_edge = [(
        row['up_id_product'],
        row['down_id_product'],
        {
            'count': row['count']
        })]
    g_bom.add_edges_from(lst_add_edge)

# dcp_networkx
pos = nx.nx_agraph.graphviz_layout(g_bom, prog="dot", args="")
node_labels = nx.get_node_attributes(g_bom, 'Name')
# rename node 1
# node_labels['1'] = '解决方案'
temp = {}
for key, value in node_labels.items():
    temp[key] = key + " " + value
node_labels = temp
colors = nx.get_edge_attributes(g_bom, "count")
colors = [w for (n1, n2, _), w in colors.items()]
vmin = min(colors)
vmax = max(colors)
cmap = plt.cm.Blues

pos_new = {}
for node, p in pos.items():
    pos_new[node] = (p[1], p[0])

fig = plt.figure(figsize=(6, 10), dpi=300)
# plt.subplots_adjust(right=0.7)
nx.draw(g_bom,
        pos_new,
        node_size=50,
        labels=node_labels,
        font_size=5,
        width=1.5,
        edge_color=colors,
        edge_cmap=cmap,
        edge_vmin=vmin,
        edge_vmax=vmax)
plt.axis('off')
axis = plt.gca()
axis.set_xlim([1.2*x for x in axis.get_xlim()])
axis.set_ylim([1.2*y for y in axis.get_ylim()])

sm = plt.cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(vmin=vmin, vmax=vmax))
sm._A = []
position = fig.add_axes([0.75, 0.1, 0.01, 0.2])
cb = plt.colorbar(sm, fraction=0.01, cax=position)
cb.ax.tick_params(labelsize=8)
cb.outline.set_visible(False)
plt.savefig("output_result\\risk\\count_dcp_prod_network")
plt.close()