import agentpy as ap


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

        self.code = code
        self.name = name
        self.type_region = type_region
        self.revenue_log = revenue_log
        self.a_list_product = a_list_product
        self.dct_prod_capacity = dict.fromkeys(self.a_list_product)

        self.a_list_up_product_removed = ap.AgentList(self.model, [])
        self.a_list_product_disrupted = ap.AgentList(self.model, [])
        self.a_list_product_removed = ap.AgentList(self.model, [])

        self.dct_num_trial_up_product_removed = {}
        self.dct_request_prod_from_firm = {}

    def remove_edge_to_cus_remove_cus_up_prod(self, remove_product):
        list_out_edges = list(
            self.firm_network.graph.out_edges(
                self.firm_network.positions[self], keys=True, data='Product'))
        for n1, n2, key, product_code in list_out_edges:
            if product_code == remove_product.code:
                # remove edge
                # print(n1, n2, key, product_code)
                self.firm_network.graph.remove_edge(n1, n2, key)
                # remove customer up product if does not have alternative
                customer = ap.AgentIter(self.model, n2).to_list()[0]
                list_in_edges = list(
                    self.firm_network.graph.in_edges(n2,
                                                     keys=True,
                                                     data='Product'))
                select_edges = [
                    edge for edge in list_in_edges
                    if edge[-1] == remove_product.code
                ]
                # print(select_edges)
                if len(select_edges) == 0:
                    print(self.name, remove_product.code, 'affect', customer.name)
                    if remove_product not in \
                            customer.a_list_up_product_removed:
                        customer.a_list_up_product_removed.append(
                            remove_product)
                        customer.dct_num_trial_up_product_removed[
                            remove_product] = 0

                # # disrupt customer
                # customer = ap.AgentIter(self.model, n2).to_list()[0]
                # for product in customer.a_list_product:
                #     if product in remove_product.a_successors():
                #         if product not in customer.a_list_product_disrupted:
                #             customer.a_list_product_disrupted.append(
                #                 product)
                # print(customer.a_list_product_disrupted.code)

    def seek_alt_supply(self):
        for product in self.a_list_up_product_removed:
            print(f"{self.name} seek alt supply for {product.code}")
            if self.dct_num_trial_up_product_removed[
                    product] <= self.model.int_n_max_trial:
                # select a list of candidate firm that has the product
                candidate_alt_supply = self.model.a_list_total_firms.select([
                    product in firm.a_list_product
                    and product not in firm.a_list_product_removed
                    for firm in self.model.a_list_total_firms
                ])
                # print(candidate_alt_supply)
                # print(candidate_alt_supply.name)
                # print(candidate_alt_supply.a_list_product.code)
                if not candidate_alt_supply:
                    continue
                # select based on size
                list_prob = [
                    size / sum(candidate_alt_supply.revenue_log)
                    for size in candidate_alt_supply.revenue_log
                ]
                select_alt_supply = self.model.nprandom.choice(
                    candidate_alt_supply, p=list_prob)
                print(f"{self.name} selct alt supply for {product.code} from {select_alt_supply.name}")
                assert product in select_alt_supply.a_list_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_num_trial_up_product_removed[product] += 1

    def handle_request(self):
        print(self.name, 'handle_request')
        for product, list_firm in self.dct_request_prod_from_firm.items():
            if self.dct_prod_capacity[product] > 0:
                if len(list_firm) == 0:
                    continue
                elif len(list_firm) == 1:
                    self.accept_request(list_firm[0], product)
                elif len(list_firm) > 1:
                    # handling based on connection
                    # TBC
                    # handling based on size
                    list_size = [firm.revenue_log for firm in list_firm]
                    list_prob = [size / sum(list_size) for size in list_size]
                    select_customer = self.model.nprandom.choice(list_firm,
                                                                 p=list_prob)
                    self.accept_request(select_customer, product)
            # print(product.code, [firm.name for firm in list_firm])

    def accept_request(self, down_firm, product):
        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)
        down_firm.a_list_up_product_removed.remove(product)
        print(f"{self.name} accept {product.code} request from {down_firm.name}")

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

    def clean_before_time_step(self):
        self.dct_num_trial_up_product_removed = {}
        self.a_list_up_product_removed = ap.AgentList(self.model, [])