import sys
import numpy as np
from absl import flags
from pysc2.lib import actions
from pysc2.lib import features
from pysc2.lib import protocol
from pysc2.lib import named_array
from pysc2.env.environment import StepType
from acme import specs
from acme.sc2_types import SC2Space, SC2FuncIdSpace, Space
from acme.sc2_spec import Spec
import dm_env
from pysc2.env import sc2_env
import copy

ACTIONS_MINIGAMES, ACTIONS_MINIGAMES_ALL, ACTIONS_ALL = [
    'minigames', 'minigames_all', 'all']


class Pysc2Wrapper(dm_env.Environment):
    """
    'minigames' action set is enough to solve all minigames listed in SC2LE
    'minigames_all' expands that set with actions that may improve end result, but will drop performance
    'all' is the full action set, only necessary for generic agent playing full game with all three races

    You can also specify your own action set in the gin config file under SC2Env.action_ids
    Full list of available actions https://github.com/deepmind/pysc2/blob/master/pysc2/lib/actions.py#L447-L1008
    """

    def __init__(
        self,
        environment: sc2_env.SC2Env,
        spatial_dim=16,
        action_ids=ACTIONS_MINIGAMES,
        obs_features=None
    ):
        self._env = environment
        self.spatial_dim = spatial_dim

        # sensible action set for all minigames
        # 24 actions
        if not action_ids or action_ids in [ACTIONS_MINIGAMES, ACTIONS_MINIGAMES_ALL]:
            action_ids = [0, 1, 2, 3, 4, 6, 7, 12, 13, 42, 
                        44, 50, 79, 91, 183, 234, 309, 331, 332, 333, 
                        334, 451, 477, 452, 490]

        # some additional actions for minigames (not necessary to solve)
        # additional 18 actions
        if action_ids == ACTIONS_MINIGAMES_ALL:
            action_ids += [11, 71, 72, 73, 74, 140, 168,
                           239, 261, 264, 269, 274, 318, 335, 336, 453, 477]

        # full action space, including outdated / unusable to current race / usable only in certain cases
        if action_ids == ACTIONS_ALL:
            action_ids = [f.id for f in actions.FUNCTIONS]

        # by default use majority of obs features, except for some that are unnecessary for minigames
        # e.g. race-specific like creep and shields or redundant like player_id
        if not obs_features:
            obs_features = {
                'screen': ['player_relative', 'selected', 'visibility_map', 'unit_hit_points_ratio', 'unit_density'],
                'minimap': ['player_relative', 'selected', 'visibility_map', 'camera'],
                # available actions should always be present and in first position
                'non-spatial': ['available_actions', 'player']}

        self.act_wrapper = ActionWrapper(spatial_dim, action_ids)
        self.obs_wrapper = ObservationWrapper(obs_features, action_ids)

    def step(self, action) -> dm_env.TimeStep:
        try:
            timestep = self.obs_wrapper(self._env.step(self.act_wrapper(action)))
            return timestep

        # hacky fix from websocket timeout issue...
        except protocol.ConnectionError:
            self.restart()

    def reset(self) -> dm_env.TimeStep:
        try:
            timestep = self.obs_wrapper(self._env.reset())
            return timestep

        # hacky fix from websocket timeout issue...
        except protocol.ConnectionError:
            self.restart()

    def stop(self):
        self._env.close()

    # Restart the environment and return a timestep
    def restart(self) -> dm_env.TimeStep:
        self.stop()
        old_environment = self._env
        new_environment = sc2_env.SC2Env(
            map_name=old_environment._map_name,
            players=old_environment._players,
            agent_interface_format=old_environment._agent_interface_format,
            step_mul=old_environment._step_mul,
            save_replay_episodes=old_environment._save_replay_episodes,
            replay_dir=old_environment._replay_dir)
        self._env = new_environment
        return self.reset()

    def observation_spec(self) -> Spec:
        if not self.obs_wrapper.spec:
            self.make_specs()
        return self.obs_wrapper.spec

    def action_spec(self) -> Spec:
        if not self.act_wrapper.spec:
            self.make_specs()
        return self.act_wrapper.spec

    def reward_spec(self) -> Space:
        return Space(shape=(), dtype=np.int64, domain=(0,float('inf')), name='reward')

    def discount_spec(self) -> Space:
        return Space(shape=(), dtype=np.float, domain=(0,1), name='discount')

    def make_specs(self):
        # importing here to lazy-load
        from pysc2.env import mock_sc2_env
        mock_env = mock_sc2_env.SC2TestEnv(agent_interface_format=[
            features.parse_agent_interface_format(feature_screen=self.spatial_dim, feature_minimap=self.spatial_dim)])
        self.act_wrapper.make_spec(mock_env.action_spec())
        self.obs_wrapper.make_spec(mock_env.observation_spec())
        mock_env.close()


class ObservationWrapper:
    def __init__(self, _features=None, action_ids=None):
        self.spec = None
        self.features = _features
        self.action_ids = action_ids

        screen_feature_to_idx = {feat: idx for idx, feat in enumerate(
            features.SCREEN_FEATURES._fields)}
        minimap_feature_to_idx = {feat: idx for idx, feat in enumerate(
            features.MINIMAP_FEATURES._fields)}
        
        # use feature masks to select necessary features for minigames.
        self.feature_masks = {
            'screen': [screen_feature_to_idx[f] for f in _features['screen']],
            'minimap': [minimap_feature_to_idx[f] for f in _features['minimap']]
        }

    def __call__(self, timestep) -> dm_env.TimeStep:
        timestep = timestep[0]

        obs = timestep.observation

        obs_wrapped = [
            obs['feature_screen'][self.feature_masks['screen']],
            obs['feature_minimap'][self.feature_masks['minimap']]
        ]
        for feat_name in self.features['non-spatial']:
            if feat_name == 'available_actions':
                fn_ids_idxs = [i for i, fn_id in enumerate(
                    self.action_ids) if fn_id in obs[feat_name]]
                mask = np.zeros((len(self.action_ids),), dtype=np.int32)
                mask[fn_ids_idxs] = 1
                obs[feat_name] = mask
            obs_wrapped.append(obs[feat_name])

        wrapped_timestep = dm_env.TimeStep(observation=obs_wrapped,
                                        reward=timestep.reward,
                                        step_type=timestep.step_type, 
                                        discount=timestep.discount)

        return wrapped_timestep

    def make_spec(self, spec):
        spec = spec[0]

        default_dims = {
            'available_actions': (len(self.action_ids), ),
        }

        screen_shape = (
            len(self.features['screen']), *spec['feature_screen'][1:])
        minimap_shape = (
            len(self.features['minimap']), *spec['feature_minimap'][1:])
        screen_dims = get_spatial_dims(
            self.features['screen'], features.SCREEN_FEATURES)
        minimap_dims = get_spatial_dims(
            self.features['minimap'], features.MINIMAP_FEATURES)

        spaces = [
            SC2Space(screen_shape, 'screen',
                     self.features['screen'], screen_dims),
            SC2Space(minimap_shape, 'minimap',
                     self.features['minimap'], minimap_dims),
        ]

        for feat in self.features['non-spatial']:
            if 0 in spec[feat]:
                spec[feat] = default_dims[feat]

            spaces.append(Space(spec[feat], name=feat))

        self.spec = Spec(spaces, 'Observation')
        


class ActionWrapper:
    def __init__(self, spatial_dim, action_ids, args=None):
        self.spec = None
        if not args:
            args = [
                'screen',
                'minimap',
                'screen2',
                'queued',
                'control_group_act',
                'control_group_id',
                'select_add',
                'select_point_act',
                'select_unit_act',
                # 'select_unit_id'
                'select_worker',
                'build_queue_id',
                # 'unload_id'
            ]
        self.func_ids = action_ids
        self.args, self.spatial_dim = args, spatial_dim

    def __call__(self, action):
        # create a deepcoy of the action list so that it won't affect the original one
        action_to_func_call = copy.deepcopy(action)

        default = {
            'control_group_act': 0,
            'control_group_id': 0,
            'select_point_act': 0,
            'select_unit_act': 0,
            'select_unit_id': 0,
            'build_queue_id': 0,
            'unload_id': 0,
        }
        # fn_id_idx is the index position of the chosen function_id in the list "func_ids"
        fn_id_idx, args = action_to_func_call.pop(0), []

        # get the real function id defined in pysc2 library
        fn_id = self.func_ids[fn_id_idx]

        # create the argument list for the function id chosen
        for arg_type in actions.FUNCTIONS[fn_id].args:
            arg_name = arg_type.name
            if arg_name in self.args:
                arg = action_to_func_call[self.args.index(arg_name)]
                # pysc2 expects all args in their separate list
                if type(arg) not in [list, tuple]:
                    arg = [arg]
                # pysc2 expects spatial coords, but we have flattened => attempt to fix
                # In random agent under the spatial coordinates are not flattened.
                # if len(arg_type.sizes) > 1 and len(arg) == 1:
                #     arg = [arg[0] % self.spatial_dim,
                #            arg[0] // self.spatial_dim]
                args.append(arg)
            else:
                arg = [default[arg_name]]
                args.append(arg)
    
        return [actions.FunctionCall(fn_id, args)]

    def make_spec(self, spec):
        spec = spec[0]

        spaces = [SC2FuncIdSpace(self.func_ids, self.args)]
        for arg_name in self.args:
            arg = getattr(spec.types, arg_name)
            if len(arg.sizes) > 1:
                # two dimensional of the array to describe the spatial information
                # example:
                # Space "minimap", value: (2, 3)
                spaces.append(
                    Space(shape=(2,), domain=(0, self.spatial_dim), categorical=True, name=arg_name))
            else:
                spaces.append(
                    Space(shape=(), domain=(0, arg.sizes[0]), categorical=True, name=arg_name))

        self.spec = Spec(spaces, "Action")
        
def get_spatial_dims(feat_names, feats):
    feats_dims = []
    for feat_name in feat_names:
        feat = getattr(feats, feat_name)
        feats_dims.append(1)
        if feat.type == features.FeatureType.CATEGORICAL:
            feats_dims[-1] = feat.scale
    return feats_dims


