from pysc2.agents import base_agent
from pysc2.env import sc2_env
from pysc2.lib import actions, features, units
from absl import app
import time
from datetime import datetime

# tensorforce imports for environment
import numpy as np
from tensorforce import TensorForceError
from tensorforce.environments import Environment



SCREEN_SIZE = 84
MINIMAP_SIZE = 64

class minigame_agent(base_agent.BaseAgent):
  def __init__(self):
    super(minigame_agent, self).__init__()

    self.attack_coordinates = None

  def unit_type_is_selected(self, obs, unit_type):
    if (len(obs.observation.single_select) > 0 and
            obs.observation.single_select[0].unit_type == unit_type):
      return True

    if (len(obs.observation.multi_select) > 0 and
            obs.observation.multi_select[0].unit_type == unit_type):
      return True

    return False

  def get_units_by_type(self, obs, unit_type):
    return [unit for unit in obs.observation.feature_units
            if unit.unit_type == unit_type]

  def can_do(self, obs, action):
    return action in obs.observation.available_actions

  def step(self, obs):
    super(minigame_agent, self).step(obs)

    marines = self.get_units_by_type(obs, units.Terran.Marine)
    roaches = self.get_units_by_type(obs, units.Zerg.Roach)

    env = marines + roaches

    time.sleep(.5)
    return actions.FUNCTIONS.no_op()

def main(unused_argv):
  agent = minigame_agent()
  try:
    while True:
      try:
        with sc2_env.SC2Env(
          map_name="DefeatRoaches",
          #players=[sc2_env.Agent(sc2_env.Race.Terran),
          #         ],
          agent_interface_format=features.AgentInterfaceFormat(
            feature_dimensions=features.Dimensions(screen=SCREEN_SIZE,
                                                   minimap=MINIMAP_SIZE),
            use_feature_units=True),
          step_mul=16,
          game_steps_per_episode=0,
          visualize=True,
          ) as env:

          agent.setup(env.observation_spec(), env.action_spec())
          timesteps = env.reset()
          agent.reset()

          while True:
            step_actions = [agent.step(timesteps[0])]
            timesteps = env.step(step_actions)

      except KeyboardInterrupt:
        raise
      except Exception as e:
        with open('log.txt', 'a') as myfile:
          myfile.write('At %s recorded exception: %s.\n' %
                       (datetime.now(), repr(e)))

  except KeyboardInterrupt:
    pass

if __name__ == "__main__":
  app.run(main)