from mesa import Model
from mesa.space import MultiGrid
from mesa.time import RandomActivation

from ForkliftAgent import ForkliftAgent
from PatchAgent import PatchAgent
from PatchType import PatchType
from util.PathDefinitions import inverse_y
from util.PathVisualiser import draw_grid, reconstruct_path
from util.Pathfinder import a_star_search


class GameModel(Model):

    def __init__(self, width, height, graph):
        # self.num_agents = 5
        self.running = True
        self.grid = MultiGrid(height, width, True)
        self.schedule = RandomActivation(self)
        self.agents = []

        self.forklift_agent = ForkliftAgent(0, self)
        self.schedule.add(self.forklift_agent)
        self.agents.append(self.forklift_agent)

        # Add the agent to a random grid cell
        x = 5
        y = 5
        self.grid.place_agent(self.forklift_agent, (x, y))
        self.forklift_agent.current_position = (x, y)

        start, goal = (x, y), (2, 1)
        came_from, cost_so_far = a_star_search(graph, start, goal)
        draw_grid(graph, point_to=came_from, start=start, goal=goal)

        path = map(lambda t: (t[0], inverse_y(height, t[1])),
                   reconstruct_path(came_from=came_from, start=start, goal=goal))

        print("cam from: {}".format(came_from))
        print("costerino: {}".format(cost_so_far))
        draw_grid(graph, path=reconstruct_path(came_from, start=start, goal=goal))

        self.forklift_agent.assign_new_movement_task(path)

        agent = PatchAgent(1, self, PatchType.pickUp)
        # self.schedule.add(agent)
        self.grid.place_agent(agent, (self.grid.width - 1, self.grid.height - 1))
        self.agents.append(agent)

        agent = PatchAgent(2, self, PatchType.dropOff)
        # self.schedule.add(agent)
        self.grid.place_agent(agent, (0, self.grid.height - 1))
        self.agents.append(agent)

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

    def step(self):
        self.schedule.step()

        print("update multiGrid")

        self.grid.remove_agent(self.forklift_agent)
        self.grid.place_agent(self.forklift_agent, self.forklift_agent.current_position)