import agentpy as ap


class ProductAgent(ap.Agent):
    def setup(self, code, name):
        self.product_network = self.model.product_network

        self.code = code
        self.name = name

    def a_successors(self):
        # find successors of a product, return in AgentList (ProductAgent)
        nodes = self.product_network.graph.successors(
            self.product_network.positions[self])
        return ap.AgentList(
            self.model,
            [ap.AgentIter(self.model, node).to_list()[0] for node in nodes])

    def a_predecessors(self):
        # find predecessors of a product, return in AgentList (ProductAgent)
        nodes = self.product_network.graph.predecessors(
            self.product_network.positions[self])
        return ap.AgentList(
            self.model,
            [ap.AgentIter(self.model, node).to_list()[0] for node in nodes])