from mesa.agent import Agent


class DemandAgent(Agent):
    """Demand agent holding monthly demand per city and product."""

    def __init__(self, model, city, region, product_monthly_demand):
        super().__init__(model)
        self.city = city
        self.region = region
        # product_monthly_demand: {product: {month: demand_units}}
        self.product_monthly_demand = product_monthly_demand
        self.fulfilled = 0.0
        self.total_demand = sum(sum(m.values()) for m in product_monthly_demand.values())

    def step(self):
        month = self.model.current_month
        total = 0
        for product, monthly in self.product_monthly_demand.items():
            total += monthly.get(month, 0)
        if total > 0:
            self.model.record_demand(self.region, total)