import random

from mesa import Agent, Model
from mesa.time import RandomActivation
from mesa.space import MultiGrid
from mesa.datacollection import DataCollector

from ForkliftAgent import ForkliftAgent


class ForkliftModel(Model):

    def __init__(self, width, height):
        self.num_agents = 1
        self.running = True
        self.grid = MultiGrid(height, width, True)
        self.schedule = RandomActivation(self)
        # self.datacollector = DataCollector(
        #     model_reporters={"Gini": compute_gini},
        #     agent_reporters={"Wealth": lambda a: a.wealth}
        # )

        # Create agents
        for i in range(self.num_agents):
            a = ForkliftAgent(1, self)
            self.schedule.add(a)
            # Add the agent to a random grid cell
            x = random.randrange(self.grid.width)
            y = random.randrange(self.grid.height)
            self.grid.place_agent(a, (x, y))

    def step(self):
        # self.datacollector.collect(self)
        self.schedule.step()