import agentpy as ap
import math


class FirmAgent(ap.Agent):
    def setup(self, code, name, type_region, revenue_log, a_lst_product):
        self.firm_network = self.model.firm_network
        self.product_network = self.model.product_network

        # self para
        self.code = code
        self.name = name
        self.type_region = type_region
        self.ori_size = revenue_log
        self.size = revenue_log
        self.dct_prod_up_prod_stat = {}
        self.dct_prod_capacity = {}

        # para in trial
        self.dct_n_trial_up_prod_disrupted = {}
        self.dct_cand_alt_supp_up_prod_disrupted = {}
        self.dct_request_prod_from_firm = {}

        # external variable
        self.is_prf_size = self.model.is_prf_size
        self.is_prf_conn = bool(self.p.prf_conn)
        self.str_cap_limit_prob_type = str(self.p.cap_limit_prob_type)
        self.flt_cap_limit_level = float(self.p.cap_limit_level)
        self.flt_diff_new_conn = float(self.p.diff_new_conn)
        self.flt_crit_supplier = float(self.p.crit_supplier)

        # init dct_prod_up_prod_stat (self para)
        for prod in a_lst_product:
            self.dct_prod_up_prod_stat[prod] = {
                # (Normal / Affected / Disrupted / Removed, time step)
                'status': [('N', 0)],
                # have or have no supply
                'supply': dict.fromkeys(prod.a_predecessors(), True)
            }

        # init extra capacity (self para)
        for product in a_lst_product:
            # init extra capacity based on discrete uniform distribution
            assert self.str_cap_limit_prob_type in ['uniform', 'normal'], \
                "cap_limit_prob_type other than uniform, normal"
            if self.str_cap_limit_prob_type == 'uniform':
                extra_cap_mean = \
                    self.size / self.flt_cap_limit_level
                extra_cap = self.model.nprandom.integers(extra_cap_mean-2,
                                                         extra_cap_mean+2)
                extra_cap = 0 if round(extra_cap) < 0 else round(extra_cap)
                # print(firm_agent.name, extra_cap)
                self.dct_prod_capacity[product] = extra_cap
            elif self.str_cap_limit_prob_type == 'normal':
                extra_cap_mean = \
                    self.size / self.flt_cap_limit_level
                extra_cap = self.model.nprandom.normal(extra_cap_mean, 1)
                extra_cap = 0 if round(extra_cap) < 0 else round(extra_cap)
                # print(firm_agent.name, extra_cap)
                self.dct_prod_capacity[product] = extra_cap

    def remove_edge_to_cus_affect_cus_up_prod(self, disrupted_prod):
        # para remove_product is the product that self got disrupted
        lst_out_edge = list(
            self.firm_network.graph.out_edges(
                self.firm_network.positions[self], keys=True, data='Product'))
        for n1, n2, key, product_code in lst_out_edge:
            if product_code == disrupted_prod.code:
                # remove edge to customer
                self.firm_network.graph.remove_edge(n1, n2, key)

                # customer up product affected conditionally
                customer = ap.AgentIter(self.model, n2).to_list()[0]
                lst_in_edge = list(
                    self.firm_network.graph.in_edges(n2,
                                                     keys=True,
                                                     data='Product'))
                lst_select_in_edge = [
                    edge for edge in lst_in_edge
                    if edge[-1] == disrupted_prod.code
                ]
                prob_lost_supp = math.exp(-1 * self.flt_crit_supplier *
                                          len(lst_select_in_edge))
                if self.model.nprandom.choice([True, False],
                                              p=[prob_lost_supp,
                                                 1 - prob_lost_supp]):
                    customer.dct_n_trial_up_prod_disrupted[disrupted_prod] = 0
                    for prod in customer.dct_prod_up_prod_stat.keys():
                        if disrupted_prod in \
                                customer.dct_prod_up_prod_stat[
                                    prod]['supply'].keys():
                            customer.dct_prod_up_prod_stat[
                                prod]['supply'][disrupted_prod] = False
                            customer.dct_prod_up_prod_stat[
                                prod]['status'].append(('A', self.model.t))
                            print(self.name, disrupted_prod.code, 'affect',
                                  customer.name, prod.code)

    def seek_alt_supply(self, product):
        # para product is the product that self is seeking
        print(f"{self.name} seek alt supply for {product.code}")
        if self.dct_n_trial_up_prod_disrupted[
                product] <= self.model.int_n_max_trial:
            if self.dct_n_trial_up_prod_disrupted[product] == 0:
                # select a list of candidate firm that has the product
                self.dct_cand_alt_supp_up_prod_disrupted[product] = \
                    self.model.a_lst_total_firms.select([
                        firm.is_prod_in_current_normal(product)
                        for firm in self.model.a_lst_total_firms
                    ])
            if self.dct_cand_alt_supp_up_prod_disrupted[product]:
                # select based on connection
                lst_firm_connect = []
                if self.is_prf_conn:
                    for firm in \
                            self.dct_cand_alt_supp_up_prod_disrupted[product]:
                        out_edges = self.model.firm_network.graph.out_edges(
                            self.model.firm_network.positions[firm], keys=True)
                        in_edges = self.model.firm_network.graph.in_edges(
                            self.model.firm_network.positions[firm], keys=True)
                        lst_adj_firm = []
                        lst_adj_firm += \
                            [ap.AgentIter(self.model, edge[1]).to_list()[
                                0].code for edge in out_edges]
                        lst_adj_firm += \
                            [ap.AgentIter(self.model, edge[0]).to_list()[
                                0].code for edge in in_edges]
                        if self.code in lst_adj_firm:
                            lst_firm_connect.append(firm)
                if len(lst_firm_connect) == 0:
                    # select based on size or not
                    if self.is_prf_size:
                        lst_size = \
                            [size for size in
                                self.dct_cand_alt_supp_up_prod_disrupted[
                                    product].size]
                        lst_prob = [size / sum(lst_size)
                                    for size in lst_size]
                        select_alt_supply = self.model.nprandom.choice(
                            self.dct_cand_alt_supp_up_prod_disrupted[product],
                            p=lst_prob)
                    else:
                        select_alt_supply = self.model.nprandom.choice(
                            self.dct_cand_alt_supp_up_prod_disrupted[product])
                elif len(lst_firm_connect) > 0:
                    # select based on size of connected firm or not
                    if self.is_prf_size:
                        lst_firm_size = \
                            [firm.size for firm in lst_firm_connect]
                        lst_prob = \
                            [size / sum(lst_firm_size)
                                for size in lst_firm_size]
                        select_alt_supply = \
                            self.model.nprandom.choice(lst_firm_connect,
                                                       p=lst_prob)
                    else:
                        select_alt_supply = \
                            self.model.nprandom.choice(lst_firm_connect)
                print(
                    f"{self.name} selct alt supply for {product.code} "
                    f"from {select_alt_supply.name}"
                )
                assert select_alt_supply.is_prod_in_current_normal(product), \
                    f"{select_alt_supply} \
                        does not produce requested product {product}"

                if product in select_alt_supply.dct_request_prod_from_firm.\
                        keys():
                    select_alt_supply.dct_request_prod_from_firm[
                        product].append(self)
                else:
                    select_alt_supply.dct_request_prod_from_firm[product] = [
                        self
                    ]
                print(
                    select_alt_supply.name, 'dct_request_prod_from_firm', {
                        key.code: [v.name for v in value]
                        for key, value in
                        select_alt_supply.dct_request_prod_from_firm.items()
                    })

                self.dct_n_trial_up_prod_disrupted[product] += 1

    def handle_request(self):
        print(self.name, 'handle_request')
        for product, lst_firm in self.dct_request_prod_from_firm.items():
            if self.dct_prod_capacity[product] > 0:
                if len(lst_firm) == 0:
                    continue
                elif len(lst_firm) == 1:
                    self.accept_request(lst_firm[0], product)
                elif len(lst_firm) > 1:
                    # handling based on connection
                    lst_firm_connect = []
                    if self.is_prf_conn:
                        for firm in lst_firm:
                            out_edges = \
                                self.model.firm_network.graph.out_edges(
                                    self.model.firm_network.positions[firm],
                                    keys=True)
                            in_edges = \
                                self.model.firm_network.graph.in_edges(
                                    self.model.firm_network.positions[firm],
                                    keys=True)
                            lst_adj_firm = []
                            lst_adj_firm += \
                                [ap.AgentIter(self.model, edge[1]).to_list()[
                                    0].code for edge in out_edges]
                            lst_adj_firm += \
                                [ap.AgentIter(self.model, edge[0]).to_list()[
                                    0].code for edge in in_edges]
                            if self.code in lst_adj_firm:
                                lst_firm_connect.append(firm)
                    if len(lst_firm_connect) == 0:
                        # handling based on size or not
                        if self.is_prf_size:
                            lst_firm_size = \
                                [firm.size for firm in lst_firm]
                            lst_prob = \
                                [size / sum(lst_firm_size)
                                 for size in lst_firm_size]
                            select_customer = \
                                self.model.nprandom.choice(lst_firm,
                                                           p=lst_prob)
                        else:
                            select_customer = \
                                self.model.nprandom.choice(lst_firm)
                        self.accept_request(select_customer, product)
                    elif len(lst_firm_connect) > 0:
                        # handling based on size of connected firm or not
                        if self.is_prf_size:
                            lst_firm_size = \
                                [firm.size for firm in lst_firm_connect]
                            lst_prob = \
                                [size / sum(lst_firm_size)
                                 for size in lst_firm_size]
                            select_customer = \
                                self.model.nprandom.choice(lst_firm_connect,
                                                           p=lst_prob)
                        else:
                            select_customer = \
                                self.model.nprandom.choice(lst_firm_connect)
                        self.accept_request(select_customer, product)

    def accept_request(self, down_firm, product):
        # para product is the product that self is selling
        prod_accept = self.flt_diff_new_conn
        if self.model.nprandom.choice([True, False],
                                      p=[prod_accept, 1 - prod_accept]):
            self.firm_network.graph.add_edges_from([
                (self.firm_network.positions[self],
                 self.firm_network.positions[down_firm], {
                     'Product': product.code
                })
            ])
            self.dct_prod_capacity[product] -= 1
            self.dct_request_prod_from_firm[product].remove(down_firm)

            for prod in down_firm.dct_prod_up_prod_stat.keys():
                if product in down_firm.dct_prod_up_prod_stat[
                        prod]['supply'].keys():
                    down_firm.dct_prod_up_prod_stat[
                        prod]['supply'][product] = True
                    down_firm.dct_prod_up_prod_stat[
                        prod]['status'].append(('N', self.model.t))
            del down_firm.dct_n_trial_up_prod_disrupted[product]
            del down_firm.dct_cand_alt_supp_up_prod_disrupted[product]

            print(
                f"{self.name} accept {product.code} request "
                f"from {down_firm.name}"
            )
        else:
            down_firm.dct_cand_alt_supp_up_prod_disrupted[product].remove(self)

    def clean_before_trial(self):
        self.dct_request_prod_from_firm = {}

    def clean_before_time_step(self):
        self.dct_n_trial_up_prod_disrupted = \
            dict.fromkeys(self.dct_n_trial_up_prod_disrupted.keys(), 0)
        self.dct_cand_alt_supp_up_prod_disrupted = {}

    def get_firm_network_node(self):
        return self.firm_network.positions[self]

    def is_prod_in_current_normal(self, prod):
        if prod in self.dct_prod_up_prod_stat.keys():
            if self.dct_prod_up_prod_stat[prod]['status'][-1][0] == 'N':
                return True
            else:
                return False
        else:
            return False