from typing import Dict, Any, Union, List, Callable

import gym


def prefix_dict(input_dict: Dict[str, Any], prefix: str):
    return {prefix + key: value for key, value in input_dict.items()}


def is_gym_space_suitable_for_table_idx(space: Union[List[int], gym.Space]):
    return isinstance(space, int) or isinstance(space, list) or isinstance(space, gym.spaces.Discrete) or isinstance(
        space, gym.spaces.Discrete)


def convert_gym_space_to_q_shape(space: Union[List[int], gym.Space]):
    if isinstance(space, int):
        return [space]
    elif isinstance(space, list):
        return space
    elif isinstance(space, gym.spaces.Discrete):
        return [space.n]
    elif isinstance(space, gym.spaces.MultiDiscrete):
        return space.nvec
    else:
        assert False, "Can only use discrete or multidiscrete for q-table agent"


class TrainingProgress:
    global_step_count: int
    global_episode_count: int

    def __init__(self, global_step_count=0, global_episode_count=0):
        self.global_step_count = global_step_count
        self.global_episode_count = global_episode_count


class Updater:
    """A neat little class to perform tasks periodically without maintaining separate counters"""

    def __init__(self, do_update: Callable[[], None]):
        self.do_update = do_update
        self.steps_since_update = 0

    def update_now(self):
        self.do_update()
        self.steps_since_update = 0

    def update_every(self, steps: int) -> bool:
        """
        Note that update_now resets the counter since the last update
        :return: Whether an update was performed
        """
        if self.steps_since_update >= steps - 1:
            self.update_now()
            return True
        else:
            self.steps_since_update += 1
            return False
