import math
import random

from mesa.visualization.ModularVisualization import ModularServer
from mesa.visualization.modules import CanvasGrid

from ForkliftAgent import ForkliftAgent
from GameModel import GameModel
from InitialStateFactory import InitialStateFactory
from PatchAgent import PatchAgent
from PatchType import PatchType
from PictureVisualizationAgent import PictureVisualizationAgent
from data.Order import Order
from data.enum.Direction import Direction
from tensorflow import keras

from data.enum.Priority import Priority
from genetic_order.GeneticOrder import GeneticOrder
from util.PathDefinitions import GridWithWeights
from visualization.DisplayAttributeElement import DisplayAttributeElement
from visualization.DisplayItemListAttribute import DisplayItemListAttributeElement
from visualization.DisplayOrderList import DisplayOrderList

colors = [
    'blue', 'cyan', 'orange', 'yellow', 'magenta', 'purple', '#103d3e', '#9fc86c',
    '#b4c2ed', '#31767d', '#31a5fa', '#ba96e0', '#fef3e4', '#6237ac', '#f9cacd', '#1e8123'
]


def agent_portrayal(agent):
    if isinstance(agent, ForkliftAgent):
        shape = ""
        if agent.current_rotation == Direction.top:
            shape = "img/image_top.png"
        elif agent.current_rotation == Direction.right:
            shape = "img/image_right.png"
        elif agent.current_rotation == Direction.down:
            shape = "img/image_down.png"
        elif agent.current_rotation == Direction.left:
            shape = "img/image_left.png"

        portrayal = {"Shape": shape, "scale": 1.0, "Layer": 0}

    if isinstance(agent, PatchAgent):
        color = colors[0]
        if agent.patch_type == PatchType.wall:
            portrayal = {"Shape": "img/brick.webp", "scale": 1.0, "Layer": 0}
        elif agent.patch_type == PatchType.dropOff:
            portrayal = {"Shape": "img/truck.png", "scale": 1.0, "Layer": 0}
        elif agent.patch_type == PatchType.pickUp:
            portrayal = {"Shape": "img/okB00mer.png", "scale": 1.0, "Layer": 0}
        elif agent.patch_type == PatchType.diffTerrain:
            portrayal = {"Shape": "img/puddle.png", "scale": 1.0, "Layer": 0}
        elif agent.patch_type == PatchType.packingShelf:
            portrayal = {"Shape": "img/shelf_s.jpg", "scale": 1.0, "Layer": 0}
        elif agent.patch_type == PatchType.packingRefrigerator:
            portrayal = {"Shape": "img/fridge_f.jpg", "scale": 1.0, "Layer": 0}
        elif agent.patch_type == PatchType.packingDoor:
            portrayal = {"Shape": "img/door_d.jpg", "scale": 1.0, "Layer": 0}

        elif agent.patch_type == PatchType.divider:
            portrayal = \
                {"Shape": "rect",
                 "Filled": "true",
                 "Layer": 0,
                 "Color": "black",
                 "w": 1,
                 "h": 1}
        else:
            color = colors[random.randrange(13) + 3]
            portrayal = {"Shape": "rect",
                         "Filled": "true",
                         "Layer": 0,
                         "Color": color,
                         "w": 1,
                         "h": 1}

    if isinstance(agent, PictureVisualizationAgent):
        portrayal = {"Shape": f"{agent.img}", "scale": 3.0, "Layer": 0}

    return portrayal


if __name__ == '__main__':
    base = 512
    gridWidth = 10
    gridHeight = 13
    scale = base / gridWidth

    diagram = GridWithWeights(gridWidth, gridHeight)
    diagram.walls = [(6, 5), (6, 6), (6, 7), (6, 8), (2, 3), (2, 4), (3, 4), (4, 4), (6, 4)]
    diagram.puddles = [(2, 2), (2, 5), (2, 6), (5, 4)]
    diagram.packingStations = [(PatchType.packingShelf, (4, 8)), (PatchType.packingRefrigerator, (4, 6)), (PatchType.packingDoor, (4, 2))]

    grid = CanvasGrid(agent_portrayal, gridWidth, gridHeight, scale * gridWidth, scale * gridHeight)

    readyText = DisplayAttributeElement("phase")
    # current_item = DisplayPictureElement("current_item_recognition")
    provided_itesm = DisplayItemListAttributeElement("provided_items")
    recognised_items = DisplayItemListAttributeElement("recognised_items")
    ordersText = DisplayOrderList("orderList")
    fulfilled_orders = DisplayOrderList("fulfilled_orders")

    model = keras.models.load_model("imageClasification/my_model")

    server = ModularServer(GameModel,
                           [grid, readyText, provided_itesm, recognised_items, ordersText,
                            fulfilled_orders],
                           "Automatyczny Wózek Widłowy",
                           dict(width=gridHeight, height=gridWidth, graph=diagram, items=50, orders=20, classificator=model))

    # def order_rule(order: Order):
    #     return order.id
    #
    #
    # punish_low = 500
    # punish_med = 300
    # def sum_wrong(member: [Order]) -> int:
    #     last_high = 0
    #     last_med = 0
    #
    #     sum_high = 0
    #     sum_med = 0
    #     sum_low = 0
    #
    #     counter = 0
    #
    #     for i in range(len(member)):
    #         o: Order = member[i]
    #         if o.priority == Priority.HIGH :
    #             sum_high += 1
    #             last_high = i
    #         elif o.priority == Priority.MEDIUM:
    #             sum_med += 1
    #             last_med = i
    #         else:
    #             sum_low += 1
    #
    #     for i in range(last_high):
    #         o: Order = member[i]
    #         if o.priority == Priority.MEDIUM:
    #             counter += punish_med
    #         elif o.priority == Priority.LOW:
    #             counter += punish_low
    #
    #     for i in range(last_med):
    #         o: Order = member[i]
    #         if o.priority == Priority.LOW:
    #             counter += punish_low
    #
    #     print("sum: high:", sum_high, "med:", sum_med, "low:", sum_low)
    #     print("last_high:", last_high, "last_med:", last_med, "sum:", counter)
    #     return counter
    #
    # orders = InitialStateFactory.generate_order_list_XD(20)
    # test: GeneticOrder = GeneticOrder(orders)
    # punish_low = test.punish_low
    # punish_med = test.punish_med
    #
    # print("SIEMA before: ")
    # sum_wrong(orders)
    # for i in orders:
    #     print("id:", i.id, "priority:", i.priority, "sum/time:", i.sum/i.time)
    #     # print("id:", i.id, "priority:", i.priority)
    #
    # newOrders = test.get_orders_sorted(orders)
    #
    # print("NAURA after:")
    # sum_wrong(newOrders)
    # for i in newOrders:
    #     print("id:", i.id, "priority:", i.priority, "sum/time:", i.sum/i.time)
    #     # print("id:", i.id, "priority:", i.priority)

    server.port = 8888
    server.launch()