import agentpy as ap
import pandas as pd
import numpy as np
import networkx as nx
from firm import FirmAgent

sample = 0
seed = 0
n_iter = 3
list_init_remove_firm_code = [0, 2]
dct_sample_para = {
    'sample': sample,
    'seed': seed,
    'n_iter': n_iter,
    'list_init_remove_firm_code': list_init_remove_firm_code
}


class Model(ap.Model):
    def setup(self):
        self.sample = self.p.sample
        self.nprandom = np.random.default_rng(self.p.seed)
        self.int_n_iter = int(self.p.n_iter)

        # init graph bom
        BomNodes = pd.read_csv('BomNodes.csv', index_col=0)
        BomNodes.set_index('Code', inplace=True)
        BomCateNet = pd.read_csv('BomCateNet.csv', index_col=0)
        BomCateNet.fillna(0, inplace=True)

        G_bom = nx.from_pandas_adjacency(BomCateNet,
                                         create_using=nx.MultiDiGraph())

        bom_labels_dict = {}
        for code in G_bom.nodes:
            bom_labels_dict[code] = BomNodes.loc[code].to_dict()
        nx.set_node_attributes(G_bom, bom_labels_dict)

        # init graph firm
        Firm = pd.read_csv("Firm_amended.csv")
        Firm.fillna(0, inplace=True)
        Firm_attr = Firm.loc[:, ["Code", "Name", "Type_Region", "Revenue_Log"]]
        firm_product = []
        for _, row in Firm.loc[:, '1':].iterrows():
            firm_product.append(row[row == 1].index.to_list())
        Firm_attr.loc[:, 'Product_Code'] = firm_product
        Firm_attr.set_index('Code')
        G_Firm = nx.MultiDiGraph()
        G_Firm.add_nodes_from(Firm["Code"])

        firm_labels_dict = {}
        for code in G_Firm.nodes:
            firm_labels_dict[code] = Firm_attr.loc[code].to_dict()
        nx.set_node_attributes(G_Firm, firm_labels_dict)

        # add edge to G_firm according to G_bom
        for node in nx.nodes(G_Firm):
            # print(node, '-'*20)
            list_pred_product_code = []
            for product_code in G_Firm.nodes[node]['Product_Code']:
                list_pred_product_code += list(
                    G_bom.predecessors(product_code))
            list_pred_product_code = list(set(list_pred_product_code))
            for pred_product_code in list_pred_product_code:
                # print(pred_product_code)
                list_pred_firms = Firm.index[Firm[pred_product_code] ==
                                             1].to_list()
                list_revenue_log = [
                    G_Firm.nodes[pred_firm]['Revenue_Log']
                    for pred_firm in list_pred_firms
                ]
                list_prob = [
                    (v - min(list_revenue_log) + 1) /
                    (max(list_revenue_log) - min(list_revenue_log) + 1)
                    for v in list_revenue_log
                ]
                list_flag = [
                    self.nprandom.choice([1, 0], p=[prob, 1 - prob])
                    for prob in list_prob
                ]
                # print(list(zip(list_pred_firms,list_flag, list_prob)))
                list_added_edges = [
                    (node, pred_firm)
                    for pred_firm, flag in zip(list_pred_firms, list_flag)
                    if flag == 1
                ]
                G_Firm.add_edges_from(list_added_edges)
                # print('-'*20)

        self.firm_network = ap.Network(self, G_Firm)
        # print([node.label for node in self.firm_network.nodes])
        # print([list(self.firm_network.graph.predecessors(node)) for node in self.firm_network.nodes])
        # print([self.firm_network.graph.nodes[node.label]['Name'] for node in self.firm_network.nodes])
        # print([v for v in self.firm_network.graph.nodes(data=True)])

        # init firm
        for ag_node, attr in self.firm_network.graph.nodes(data=True):
            firm_agent = FirmAgent(self,
                                   code=attr['Code'],
                                   name=attr['Name'],
                                   type_region=attr['Type_Region'],
                                   revenue_log=attr['Revenue_Log'])
            self.firm_network.add_agents([firm_agent], [ag_node])
        a_list_total_firms = ap.AgentList(self, self.firm_network.agents)
        # set the initial firm that is removed
        list_b_is_removed = list(
            map(lambda x: x in self.p.list_init_remove_firm_code,
                self.firm_network.agents.code))
        self.a_list_firms_removed = a_list_total_firms.select(
            list_b_is_removed)
        self.a_list_firms_removed.is_removed = True

    def update(self):
        # Update list of unhappy people
        # self.agents.update_happiness()
        # self.unhappy = self.agents.select(self.agents.happy == False)

        # Stop simulation if reached terminal number of iteration
        if self.t == self.int_n_iter:
            self.stop()

    def step(self):
        # Move unhappy people to new location
        self.unhappy.find_new_home()

    def end(self):
        pass

    def draw_network(self):
        import matplotlib.pyplot as plt
        plt.rcParams['font.sans-serif'] = 'SimHei'
        pos = nx.nx_agraph.graphviz_layout(self.firm_network.graph,
                                           prog="twopi",
                                           args="")
        node_label = nx.get_node_attributes(self.firm_network.graph, 'Name')
        node_size = list(
            nx.get_node_attributes(self.firm_network.graph,
                                   'Revenue_Log').values())
        node_size = list(map(lambda x: x**2, node_size))
        plt.figure(figsize=(12, 12), dpi=300)
        nx.draw(self.firm_network.graph,
                pos,
                node_size=node_size,
                labels=node_label,
                font_size=6)
        plt.savefig("network.png")


model = Model(dct_sample_para)
model.setup()
# model.draw_network()