import agentpy as ap


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.size_stat = []
        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)

        # init size_stat (self para)
        # (size, time step)
        self.size_stat.append((revenue_log, 0))

        # init dct_prod_up_prod_stat (self para)
        for prod in a_lst_product:
            self.dct_prod_up_prod_stat[prod] = {
                # (Normal / Disrupted / Removed, time step)
                'p_stat': [('N', 0)],
                # supply for each component and respective disrupted supplier
                # set_disrupt_firm is refreshed to empty at each update
                's_stat': {up_prod: {'stat': True,
                                     'set_disrupt_firm': set()}
                           for up_prod in prod.a_predecessors()}
                # Note: do not use fromkeys as it's a shallow copy
            }

        # 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_stat[0][0] / 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)
                self.dct_prod_capacity[product] = extra_cap
            elif self.str_cap_limit_prob_type == 'normal':
                extra_cap_mean = \
                    self.size_stat[0][0] / 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)
                self.dct_prod_capacity[product] = extra_cap

    def remove_edge_to_cus(self, disrupted_prod):
        # para disrupted_prod 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:
                # update customer up product supplier status
                customer = ap.AgentIter(self.model, n2).to_list()[0]
                for prod in customer.dct_prod_up_prod_stat.keys():
                    if disrupted_prod in \
                            customer.dct_prod_up_prod_stat[
                                prod]['s_stat'].keys():
                        customer.dct_prod_up_prod_stat[
                            prod]['s_stat'][disrupted_prod][
                            'set_disrupt_firm'].add(self)
                        print(f"{self.name} disrupt {customer.name}'s "
                              f"{prod.code} due to {disrupted_prod.code}")
                # remove edge to customer
                self.firm_network.graph.remove_edge(n1, n2, key)

    def disrupt_cus_prod(self, prod, disrupted_up_prod):
        # para prod is the product that has disrupted_up_prod
        # para disrupted_up_prod is the product that
        # self's component exists disrupted supplier
        num_lost = \
            len(self.dct_prod_up_prod_stat[prod]['s_stat']
                [disrupted_up_prod]['set_disrupt_firm'])
        num_remain = \
            len([u for u, _, _, d in
                 self.firm_network.graph.in_edges(self.get_firm_network_node(),
                                                  keys=True,
                                                  data='Product')
                 if d == disrupted_up_prod.code])
        lost_percent = num_lost / (num_lost + num_remain)
        lst_size = \
            [firm.size_stat[-1][0] for firm in self.model.a_lst_total_firms]
        std_size = (self.size_stat[-1][0] - min(lst_size) + 1) \
            / (max(lst_size) - min(lst_size) + 1)
        prob_disrupt = 1 - std_size * (1 - lost_percent)
        if self.model.nprandom.choice([True, False],
                                      p=[prob_disrupt,
                                         1 - prob_disrupt]):
            self.dct_n_trial_up_prod_disrupted[disrupted_up_prod] = 0
            self.dct_prod_up_prod_stat[
                prod]['s_stat'][disrupted_up_prod]['stat'] = False
            status, _ = self.dct_prod_up_prod_stat[
                prod]['p_stat'][-1]
            if status != 'D':
                self.dct_prod_up_prod_stat[
                    prod]['p_stat'].append(('D', self.model.t))
                print(f"{self.name}'s {prod.code} turn to D status due to "
                      f"disrupted supplier of {disrupted_up_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]:
                        node_self = self.get_firm_network_node()
                        node_firm = firm.get_firm_network_node()
                        if self.model.firm_network.graph.\
                                has_edge(node_self, node_firm) or \
                                self.model.firm_network.graph.\
                                has_edge(node_firm, node_self):
                            lst_firm_connect.append(firm)
                if len(lst_firm_connect) == 0:
                    # select based on size or not
                    if self.is_prf_size:
                        lst_size = \
                            [firm.size_stat[-1][0] for firm in
                                self.dct_cand_alt_supp_up_prod_disrupted[
                                    product]]
                        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_stat[-1][0]
                             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:
                            node_self = self.get_firm_network_node()
                            node_firm = firm.get_firm_network_node()
                            if self.model.firm_network.graph.\
                                    has_edge(node_self, node_firm) or \
                                    self.model.firm_network.graph.\
                                    has_edge(node_firm, node_self):
                                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_stat[-1][0] 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_stat[-1][0]
                                 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)
            else:
                for down_firm in lst_firm:
                    down_firm.dct_cand_alt_supp_up_prod_disrupted[
                        product].remove(self)

                    print(
                        f"{self.name} denied {product.code} request "
                        f"from {down_firm.name} for lack of capacity"
                    )

    def accept_request(self, down_firm, product):
        # para product is the product that self is selling
        # connected firm has no probability
        node_self = self.get_firm_network_node()
        node_d_firm = down_firm.get_firm_network_node()
        if self.model.firm_network.graph.has_edge(node_self, node_d_firm) or \
                self.model.firm_network.graph.has_edge(node_d_firm, node_self):
            prod_accept = 1.0
        else:
            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]['s_stat'].keys():
                    down_firm.dct_prod_up_prod_stat[
                        prod]['s_stat'][product]['stat'] = True
                    down_firm.dct_prod_up_prod_stat[
                        prod]['p_stat'].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)

            print(
                f"{self.name} denied {product.code} request "
                f"from {down_firm.name}"
            )

    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 = {}

        # update the status of firm
        for prod in self.dct_prod_up_prod_stat.keys():
            status, ts = self.dct_prod_up_prod_stat[prod]['p_stat'][-1]
            if ts != self.model.t:
                self.dct_prod_up_prod_stat[prod]['p_stat'].append(
                    (status, self.model.t))
            # refresh set_disrupt_firm
            for up_prod in self.dct_prod_up_prod_stat[prod]['s_stat'].keys():
                self.dct_prod_up_prod_stat[prod][
                    's_stat'][up_prod]['set_disrupt_firm'] = set()

    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]['p_stat'][-1][0] == 'N':
                return True
            else:
                return False
        else:
            return False