from warehouse import Coordinates, Tile, Pack
from queue import PriorityQueue
from math import sqrt
from attributes import TURN_LEFT_DIRECTIONS, TURN_RIGHT_DIRECTIONS, PackStatus
import pygame
import sys
import pdb
from package_location_classifier.classifier import PackageLocationClassifier

class Node:
	def __init__(self, coord_x, coord_y, package=None, is_rack=False):
		self.x = coord_x
		self.y = coord_y
		self.parent = None
		self.package = package
		self.is_rack = is_rack
		self.g_cost = 0
		self.h_cost = 0
	def __eq__(self, other):
		if isinstance(other, Node):
			return self.x == other.x and self.y == self.y
		return False

	def __lt__(self, other):
		return isinstance(other, Node) and self.g_cost < other.g_cost

	def __repr__(self):
		return "Node:{}x{}".format(self.x, self.y)

class Agent:
	def __init__(self, start_x, start_y, assigned_warehouse, radius=5):
		self.x = start_x
		self.y = start_y
		self.radius = radius
		self.warehouse = assigned_warehouse
		self.is_loaded = False
		self.transported_package = None
		self.direction = "up"
		self.dest = Node(1, 1)
		self.closed = list()
		self.open = PriorityQueue()
		self.path = list()
		self.location_classifier = PackageLocationClassifier()
		self.check_packages_locations()

	def check_packages_locations(self):
		for pack in self.warehouse.packages:
			if pack.lays_on_field.category.name in self.warehouse.storage_types:
				can_place = self.location_classifier.check_if_can_place(pack, pack.lays_on_field)
				pack.status = PackStatus.STORED if (can_place and pack.lays_on_field.capacity >= 0) else PackStatus.STORED_BAD_LOCATION


	def find_path(self):
		self.closed = []
		self.path = []
		self.open = PriorityQueue()
		if self.is_loaded:
			rack = self.find_nearest_rack_for(self.transported_package)
			self.dest = Node(rack.x_position, rack.y_position, is_rack=True)
		else:
			package = self.find_nearest_package()
			self.dest = Node(package.lays_on_field.x_position, package.lays_on_field.y_position, package=package)

		start_node = Node(self.x, self.y)

		self.open.put((0, start_node))
		while self.open:
			_, current_node = self.open.get()
			self.closed.append(current_node)
			if current_node.x == self.dest.x and current_node.y == self.dest.y:
				while current_node.x != start_node.x or current_node.y != start_node.y:
					self.path.append(current_node)
					current_node = current_node.parent
				return True
			neighbour_list = self.get_neighbours(current_node)
			for neighbour in neighbour_list:
				cost = current_node.g_cost + self.heuristic(current_node, neighbour)
				if self.check_if_closed(neighbour):
					continue
				if self.check_if_open(neighbour):
					if neighbour.g_cost > cost:
						neighbour.g_cost = cost
						neighbour.parent = current_node
				else:
					neighbour.g_cost = cost
					neighbour.h_cost = self.heuristic(neighbour, self.dest)
					neighbour.parent = current_node
					self.open.put((neighbour.g_cost, neighbour))
		return False

	def turn_left(self):
		new_direction = TURN_LEFT_DIRECTIONS.get(self.direction, self.direction)
		self.direction = new_direction

	def turn_right(self):
		new_direction = TURN_RIGHT_DIRECTIONS.get(self.direction, self.direction)
		self.direction = new_direction

	def heuristic(self, start: Node, goal: Node):
		diff_x = pow(goal.x - start.x, 2)
		diff_y = pow(goal.y - start.y, 2)
		additional_cost = 0
		return round(sqrt(diff_x + diff_y), 3) + float(10*additional_cost)
		
	def check_if_open(self, node: Node):
		return (node.x, node.y) in [(n.x, n.y) for (_,n) in self.open.queue]

	def check_if_closed(self, node: Node):
		return (node.x, node.y) in [(n.x, n.y) for n in self.closed]

	def get_neighbours(self, node: Node):
		neighbours = []
		if self.check_if_can_move(Coordinates(node.x + 1, node.y)):
			neighbours.append(Node(node.x + 1, node.y))
		if self.check_if_can_move(Coordinates(node.x - 1, node.y)):
			neighbours.append(Node(node.x - 1, node.y))
		if self.check_if_can_move(Coordinates(node.x, node.y + 1)):
			neighbours.append(Node(node.x, node.y + 1))
		if self.check_if_can_move(Coordinates(node.x, node.y - 1)):
			neighbours.append(Node(node.x, node.y - 1))
		return neighbours
    	
	def move(self):
		dest_coords = (self.dest.x, self.dest.y)
		if not self.path:
			if not self.find_path():
				return
		else:
			next = self.path.pop()
			star_dir = self.direction
			if self.x > next.x and not self.direction == 'left':
				if self.direction == 'down':
					self.turn_right()
				else:
					self.turn_left()
			elif self.x < next.x and not self.direction == 'right':
				if self.direction == 'down':
					self.turn_left()
				else:
					self.turn_right()
			elif self.y > next.y and not self.direction == 'up':
				if self.direction == 'left':
					self.turn_right()
				else:
					self.turn_left()
			elif self.y < next.y and not self.direction == 'down':
				if self.direction == 'right':
					self.turn_right()
				else:
					self.turn_left()

			if (next.x, next.y) == dest_coords:
				if self.dest.package:
					self.pick_up_package(self.dest.package)
					return
				elif self.dest.is_rack:
					self.unload_package(self.dest)
					return

			if star_dir == self.direction:
				self.x = next.x
				self.y = next.y
			else:
				self.path.append(next)
		self.closed = []
		
	def check_if_can_move(self, next_coords: Coordinates):
		tile_on_map = 0 <= next_coords.x < self.warehouse.width and 0 <= next_coords.y < self.warehouse.height
		tile_passable = True
		if not tile_on_map:
			return False
		next_tile = self.warehouse.tiles[next_coords.x][next_coords.y]
		if (next_coords.x, next_coords.y) != (self.dest.x, self.dest.y):
			tile_passable = isinstance(next_tile, Tile) and next_tile.category.passable
		return tile_passable

	def find_nearest_package(self):
		packages_costs = []
		start_node = Node(self.x, self.y)
		if not self.warehouse.packages:
			return None
		for package in self.warehouse.packages:
			if package.status == PackStatus.STORED:
				continue
			new_node = Node(package.lays_on_field.x_position, package.lays_on_field.y_position)
			cost = self.heuristic(start_node, new_node)
			if cost > 0:
				packages_costs.append((package, cost))
		if not packages_costs:
			return
			# pygame.quit()
			# sys.exit()

		package = min(packages_costs, key=lambda l: l[1])[0]
		return package

	def rack_heuristics(self, start, goal, can_place):
		heur_can_place = not can_place
		diff_x = pow(goal.x - start.x, 2)
		diff_y = pow(goal.y - start.y, 2)
		place_cost = 100 * float(heur_can_place)
		return round(sqrt(diff_x + diff_y), 3) + float(place_cost)

	def find_nearest_rack_for(self, package, expand_box=0):
		weight = package.size
		storage = "Rack"
		if package.category == "freezed":
			storage = "Fridge"
		start_node = Node(self.x, self.y)
		quarter_x = int(self.warehouse.width/4) + expand_box
		quarter_y = int(self.warehouse.height/4) + expand_box
		start_quarter_x = self.x - quarter_x if self.x - quarter_x > 0 else 0
		end_quarter_x = self.x + quarter_x if self.x + quarter_x < self.warehouse.width else self.warehouse.width - 1
		start_quarter_y = self.y - quarter_y if self.y - quarter_y > 0 else 0
		end_quarter_y = self.y + quarter_y if self.y + quarter_y < self.warehouse.height else self.warehouse.height - 1
		quarter = [row[start_quarter_y:end_quarter_y] for row in self.warehouse.tiles[start_quarter_x:end_quarter_x]]
		quarter_racks = [[t for t in row if t.category.name == storage and t.capacity >= weight] for row in quarter]
		quarter_racks = [t for row in quarter_racks for t in row]
		racks_costs = []
		for rack in quarter_racks:
			new_node = Node(rack.x_position, rack.y_position)
			can_place = self.location_classifier.check_if_can_place(package, rack)
			cost = self.rack_heuristics(start_node, new_node, can_place)
			if cost > 0:
				racks_costs.append((rack, cost))

		rack = self.find_nearest_rack_for(package, expand_box + 1) if not racks_costs else min(racks_costs, key=lambda l: l[1])[0]
		return rack


	def pick_up_package(self, pack):
		self.warehouse.packages.remove(pack)
		self.is_loaded = True
		if pack.lays_on_field.category.name in ['Rack', 'Fridge']:
			pack.lays_on_field.capacity += pack.size

		self.dest.package = None
		pack.lays_on_field = None
		self.transported_package = pack

	def unload_package(self, rack):
		pack = self.transported_package
		tile = self.warehouse.tiles[rack.x][rack.y]
		self.transported_package = None
		self.is_loaded = False
		pack.lays_on_field = tile
		pack.lays_on_field.capacity -= pack.size
		pack.status = PackStatus.STORED
		self.warehouse.packages.append(pack)
		# print(tile.air_temperature, tile.humidity)