import random

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

from PatchType import PatchType
from ForkliftAgent import ForkliftAgent
from PatchAgent import PatchAgent


class ForkliftModel(Model):

    def __init__(self, width, height):
        # self.num_agents = 5
        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}
        # )

        agent = ForkliftAgent(0, self)
        self.schedule.add(agent)
        # Add the agent to a random grid cell
        x = 3
        y = 2
        self.grid.place_agent(agent, (x, y))

        agent = PatchAgent(1, self, PatchType.pickUp)
        self.schedule.add(agent)
        self.grid.place_agent(agent, (0, self.grid.height-1))

        agent = PatchAgent(2, self, PatchType.dropOff)
        self.schedule.add(agent)
        self.grid.place_agent(agent, (self.grid.width-1, self.grid.height-1))

        for i in range(3):
            a = PatchAgent(i+3, self, PatchType.item)
            self.schedule.add(a)
            self.grid.place_agent(a, (i, 0))

        # Create patch agents
        # for i in range(self.num_agents):
        #     a = PatchAgent(i, 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()