#!/usr/bin/env python3

import copy
from enum import Enum

class Elevator:
	def __init__(self, e_id, floor, goal):
		self.id = e_id
		self.floor = floor
		self.goal = goal
		self.pickup_direction = None # the direction we're handling pickup requests for
		self.dropoffs = []

	def direction(self):
		' the direction the elevator is currently moving in '
		if self.goal is None:
			return None
		if self.goal > self.floor:
			return Direction.UP
		else:
			return Direction.DOWN

	def dropoff(self, floor):
		if floor in self.dropoffs:
			assert self.goal is not None
			return
		assert floor != self.floor
		assert self.pickup_direction is None

		if self.goal is None:
			self.goal = floor
			self.dropoffs.append(floor)
		else:
			direction = self.direction()
			# TODO: handle opposite direction dropoffs
			if direction == Direction.UP:
				if floor > self.goal:
					self.goal = floor
				self.dropoffs.append(floor)
				self.dropoffs.sort()
			elif direction == Direction.DOWN:
				if floor < self.goal:
					self.goal = floor
				self.dropoffs.append(floor)
				self.dropoffs.sort(reverse=True)

	def estimate_pickup_time(self, floor, direction):
		# while dropping off, can handle requests in same elevator dir in between current and goal
		# while picking up, can handle requests in same dir as request dir
		#    when requested floor is in dir of elevator
		our_dir = self.direction()
		if our_dir is None:
			return abs(floor - self.floor)

		if self.dropoffs: # we are moving and have dropoffs
			if our_dir == Direction.UP:
				if self.floor <= floor and direction == our_dir:
					return floor - self.floor
				else:
					to_goal = self.goal - self.floor
					to_request = abs(self.goal - floor)
					return to_goal + to_request
			else:
				if self.floor >= floor and direction == our_dir:
					return self.floor - floor
				else:
					to_goal = self.floor - self.goal
					to_request = abs(self.goal - floor)
					return to_goal + to_request
		else: # we are moving to pick up a request
			assert self.pickup_direction is not None
			if our_dir == Direction.UP:
				if self.floor <= floor <= self.goal and direction == self.pickup_direction:
					return floor - self.floor
				else:
					return float('inf')
			else:
				if self.floor >= floor >= self.goal and direction == self.pickup_direction:
					return self.floor - floor
				else:
					return float('inf')

	def __repr__(self):
		attrs = ', '.join('%s=%s' % t for t in self.__dict__.items())
		return 'Elevator(%s)' % attrs

class Direction(Enum):
	UP = 1
	DOWN = 2

class ElevatorControlSystem:
	def __init__(self, num_elevators):
		self.elevators = {i: Elevator(i, 1, None) for i in range(num_elevators)}
		self.pickup_requests = []

	def pickup(self, floor, direction):
		self.pickup_requests.append((floor, direction))

	def dropoff(self, elevator_id, floor):
		self.elevators[elevator_id].dropoff(floor)

	def step(self):
		# serve pickup requests
		for floor, direction in copy.copy(self.pickup_requests):
			min_time = float('inf')
			min_elevator = None
			for elevator in self.elevators.values():
				time = elevator.estimate_pickup_time(floor, direction)
				if time < min_time:
					min_time = time
					min_elevator = elevator
			if min_elevator:
				# only handle the request now if the min_elevator is right here or doing nothing. if
				# it's moving, min_time should improve by at least 1 per step. handle this request
				# later (in case some other elevator finishes a dropoff and can pickup sooner)
				if min_time == 0:
					assert min_elevator.floor == floor
					self.pickup_requests.remove((floor, direction))
				elif min_elevator.goal is None:
					min_elevator.goal = floor
					min_elevator.pickup_direction = direction
					self.pickup_requests.remove((floor, direction))

		# move elevators to requested floors
		for elevator in self.elevators.values():
			if elevator.goal is not None:
				if elevator.goal > elevator.floor:
					elevator.floor += 1
				elif elevator.goal < elevator.floor:
					elevator.floor -= 1
				# check dropoffs and goal after moving
				if elevator.dropoffs and elevator.dropoffs[0] == elevator.floor:
					elevator.dropoffs = elevator.dropoffs[1:]
				if elevator.goal == elevator.floor:
					elevator.goal = elevator.pickup_direction = None

if __name__ == '__main__':
	ecs = ElevatorControlSystem(2)
	print(' '*4, ecs.elevators)

	print('requesting pickup on floor 2 down')
	ecs.pickup(2, Direction.DOWN)
	ecs.step()
	print(' '*4, ecs.elevators)

	print('requesting dropoff on floor 1 and pickup on floor 1 up')
	ecs.dropoff(0, 1)
	ecs.pickup(1, Direction.UP)
	ecs.step()
	print(' '*4, ecs.elevators)

	print('requesting dropoff on floor 5 and 3')
	ecs.dropoff(1, 5)
	ecs.dropoff(1, 3)
	ecs.step()
	print(' '*4, ecs.elevators)

	print('requesting pickup on floor 3 up')
	ecs.pickup(3, Direction.UP)
	ecs.step()
	print(' '*4, ecs.elevators)

	print('requesting pickup on floor 5 down')
	ecs.pickup(5, Direction.DOWN)
	ecs.step()
	print(' '*8, ecs.pickup_requests)
	print(' '*4, ecs.elevators)

	ecs.step()
	print(' '*8, ecs.pickup_requests)
	print(' '*4, ecs.elevators)

	ecs.step()
	print(' '*8, ecs.pickup_requests)
	print(' '*4, ecs.elevators)