# Path: oscar/shared/env.py
# class Env:
# 
# 	timestamp = 0
# 
# Path: oscar/shared/screen.py
# class Screen(object):
# 
#     def __init__(self):
#         """ cache data """
#         self._scanned_units = None
#         self._scanned_units_timestamp = None
# 
#     @staticmethod
#     def width(obs):
#         return len(obs.observation[SCREEN][0])
# 
#     @staticmethod
#     def height(obs):
#         return len(obs.observation[SCREEN][0][0])
# 
#     @staticmethod
#     def random_point(obs, margin=0):
#         loc = Point()
#         loc.x = np.random.randint(margin, Screen.width(obs)-margin-1)
#         loc.y = np.random.randint(margin, Screen.height(obs)-margin-1)
#         return loc
# 
#     def scan_units(self, obs, shared, unit_ids, player_id):
#         scan = self.scan(obs, shared)
#         return [u for u in scan if u.unit_id in unit_ids and u.player_id == player_id]
# 
#     def scan(self, obs, shared):
#         """ Returns a list of the units on screen """
# 
#         if not self._scanned_units_timestamp or self._scanned_units_timestamp != shared['env'].timestamp:
#             def _explore_contiguous(x, y, unit_id, player_relative, scanned):
#                 """
#                 Explores contiguous pixels corresponding to the same unit
#                 and returns approximate center of the unit.
#                 """
# 
#                 def _try_explore(x, y, queue, scanned):
#                     if x < screen_width and y < screen_height \
#                             and scanned[x, y] == 0 \
#                             and obs.observation[SCREEN][SCREEN_PLAYER_RELATIVE][x, y] == player_relative \
#                             and obs.observation[SCREEN][SCREEN_UNIT_TYPE][x, y] == unit_id:
#                         queue.append((x, y))
#                         scanned[x, y] = 1
#                         return True
#                     else:
#                         return False
# 
#                 sum_x, sum_y, nb_scanned = x, y, 1
#                 queue = []
#                 queue.append((x, y))
# 
#                 while len(queue) > 0:
#                     cur = queue.pop()
#                     if _try_explore(cur[0]+1, cur[1], queue, scanned):
#                         sum_x += cur[0]+1
#                         sum_y += cur[1]
#                         nb_scanned += 1
#                     if _try_explore(cur[0], cur[1]+1, queue, scanned):
#                         sum_x += cur[0]
#                         sum_y += cur[1]+1
#                         nb_scanned += 1
#                     if _try_explore(cur[0], cur[1]-1, queue, scanned):
#                         sum_x += cur[0]
#                         sum_y += cur[1]-1
#                         nb_scanned += 1
# 
#                 return round(sum_x/nb_scanned), round(sum_y/nb_scanned)
# 
#             screen_width = Screen.width(obs)
#             screen_height = Screen.height(obs)
#             scanned = np.zeros((screen_width, screen_height))
#             self._scanned_units = []
# 
#             for x in range(screen_width):
#                 for y in range(screen_height):
#                     if scanned[x, y] == 0 and obs.observation[SCREEN][SCREEN_UNIT_TYPE][x, y] != 0:
#                         center_x, center_y = _explore_contiguous(x, y,
#                                                                  obs.observation[SCREEN][SCREEN_UNIT_TYPE][x, y],
#                                                                  obs.observation[SCREEN][SCREEN_PLAYER_RELATIVE][x, y],
#                                                                  scanned)
# 
#                         self._scanned_units.append(Unit(Location(screen_loc=Point(center_x, center_y)),
#                                                         obs.observation[SCREEN][SCREEN_UNIT_TYPE][x, y],
#                                                         obs.observation[SCREEN][SCREEN_PLAYER_RELATIVE][x, y]))
# 
#             self._scanned_units_timestamp = shared['env'].timestamp
# 
#         return self._scanned_units


import random
import numpy as np
from scipy.signal import convolve2d
from oscar.constants import *
from oscar.meta_action.meta_action_error import *
from oscar.shared.env import Env
from oscar.shared.screen import Screen

RANDOM_CENTER_JUMP_LIMIT = 3
RANDOM_CENTER_MOVE_FACTOR = 2
RANDOM_CENTER_ITERATION_LIMIT = 5


def find_position(obs, unit_type_id, select_method="random_center", player_relative=PLAYER_SELF, exception=NoUnitError):
    # unit_type_id is either an iterable, or a single int. In case it is a single int, it is embedded in a list,
    # so it will be process as an iterable afterward.
    try:
        unit_type_id_list = list(unit_type_id)
    except TypeError:
        unit_type_id_list = [unit_type_id]

    unit_type_map = obs.observation[SCREEN][SCREEN_UNIT_TYPE]
    player_relative_map = obs.observation[SCREEN][SCREEN_PLAYER_RELATIVE]
    correct_unit_type_array = np.isin(unit_type_map, unit_type_id_list)
    correct_player_relative_array = (player_relative_map == player_relative)
    unit_y, unit_x = (correct_unit_type_array & correct_player_relative_array).nonzero()
    if len(unit_x) == 0:
        raise exception("Unit of id {0} for the player_relative {1} is not on the screen."
                        .format(unit_type_id_list, player_relative))

    if select_method == "random_center":
        return find_random_center(unit_type_map, unit_x, unit_y, unit_type_id_list)
    elif select_method == "random":
        return random.choice(list(zip(unit_x, unit_y)))
    elif select_method == "mean":
        return int(unit_x.mean()), int(unit_y.mean())
    elif select_method == "all":

# Path: oscar/shared/env.py
# class Env:
# 
# 	timestamp = 0
# 
# Path: oscar/shared/screen.py
# class Screen(object):
# 
#     def __init__(self):
#         """ cache data """
#         self._scanned_units = None
#         self._scanned_units_timestamp = None
# 
#     @staticmethod
#     def width(obs):
#         return len(obs.observation[SCREEN][0])
# 
#     @staticmethod
#     def height(obs):
#         return len(obs.observation[SCREEN][0][0])
# 
#     @staticmethod
#     def random_point(obs, margin=0):
#         loc = Point()
#         loc.x = np.random.randint(margin, Screen.width(obs)-margin-1)
#         loc.y = np.random.randint(margin, Screen.height(obs)-margin-1)
#         return loc
# 
#     def scan_units(self, obs, shared, unit_ids, player_id):
#         scan = self.scan(obs, shared)
#         return [u for u in scan if u.unit_id in unit_ids and u.player_id == player_id]
# 
#     def scan(self, obs, shared):
#         """ Returns a list of the units on screen """
# 
#         if not self._scanned_units_timestamp or self._scanned_units_timestamp != shared['env'].timestamp:
#             def _explore_contiguous(x, y, unit_id, player_relative, scanned):
#                 """
#                 Explores contiguous pixels corresponding to the same unit
#                 and returns approximate center of the unit.
#                 """
# 
#                 def _try_explore(x, y, queue, scanned):
#                     if x < screen_width and y < screen_height \
#                             and scanned[x, y] == 0 \
#                             and obs.observation[SCREEN][SCREEN_PLAYER_RELATIVE][x, y] == player_relative \
#                             and obs.observation[SCREEN][SCREEN_UNIT_TYPE][x, y] == unit_id:
#                         queue.append((x, y))
#                         scanned[x, y] = 1
#                         return True
#                     else:
#                         return False
# 
#                 sum_x, sum_y, nb_scanned = x, y, 1
#                 queue = []
#                 queue.append((x, y))
# 
#                 while len(queue) > 0:
#                     cur = queue.pop()
#                     if _try_explore(cur[0]+1, cur[1], queue, scanned):
#                         sum_x += cur[0]+1
#                         sum_y += cur[1]
#                         nb_scanned += 1
#                     if _try_explore(cur[0], cur[1]+1, queue, scanned):
#                         sum_x += cur[0]
#                         sum_y += cur[1]+1
#                         nb_scanned += 1
#                     if _try_explore(cur[0], cur[1]-1, queue, scanned):
#                         sum_x += cur[0]
#                         sum_y += cur[1]-1
#                         nb_scanned += 1
# 
#                 return round(sum_x/nb_scanned), round(sum_y/nb_scanned)
# 
#             screen_width = Screen.width(obs)
#             screen_height = Screen.height(obs)
#             scanned = np.zeros((screen_width, screen_height))
#             self._scanned_units = []
# 
#             for x in range(screen_width):
#                 for y in range(screen_height):
#                     if scanned[x, y] == 0 and obs.observation[SCREEN][SCREEN_UNIT_TYPE][x, y] != 0:
#                         center_x, center_y = _explore_contiguous(x, y,
#                                                                  obs.observation[SCREEN][SCREEN_UNIT_TYPE][x, y],
#                                                                  obs.observation[SCREEN][SCREEN_PLAYER_RELATIVE][x, y],
#                                                                  scanned)
# 
#                         self._scanned_units.append(Unit(Location(screen_loc=Point(center_x, center_y)),
#                                                         obs.observation[SCREEN][SCREEN_UNIT_TYPE][x, y],
#                                                         obs.observation[SCREEN][SCREEN_PLAYER_RELATIVE][x, y]))
# 
#             self._scanned_units_timestamp = shared['env'].timestamp
# 
#         return self._scanned_units


# Path: oscar/util/selection.py
import random
import numpy as np
from scipy.signal import convolve2d
from oscar.constants import *
from oscar.meta_action.meta_action_error import *
from oscar.shared.env import Env
from oscar.shared.screen import Screen


RANDOM_CENTER_JUMP_LIMIT = 3
RANDOM_CENTER_MOVE_FACTOR = 2
RANDOM_CENTER_ITERATION_LIMIT = 5


def find_position(obs, unit_type_id, select_method="random_center", player_relative=PLAYER_SELF, exception=NoUnitError):
    # unit_type_id is either an iterable, or a single int. In case it is a single int, it is embedded in a list,
    # so it will be process as an iterable afterward.
    try:
        unit_type_id_list = list(unit_type_id)
    except TypeError:
        unit_type_id_list = [unit_type_id]

    unit_type_map = obs.observation[SCREEN][SCREEN_UNIT_TYPE]
    player_relative_map = obs.observation[SCREEN][SCREEN_PLAYER_RELATIVE]
    correct_unit_type_array = np.isin(unit_type_map, unit_type_id_list)
    correct_player_relative_array = (player_relative_map == player_relative)
    unit_y, unit_x = (correct_unit_type_array & correct_player_relative_array).nonzero()
    if len(unit_x) == 0:
        raise exception("Unit of id {0} for the player_relative {1} is not on the screen."
                        .format(unit_type_id_list, player_relative))

    if select_method == "random_center":
        return find_random_center(unit_type_map, unit_x, unit_y, unit_type_id_list)
    elif select_method == "random":
        return random.choice(list(zip(unit_x, unit_y)))
    elif select_method == "mean":
        return int(unit_x.mean()), int(unit_y.mean())
    elif select_method == "all":

        return unit_x, unit_y