from functools import partial
import gym
import numpy as np
from DeepRTS.Engine import Constants
from DeepRTS.python import util
from coding.util import LimitedDiscrete


class ScenarioData:

    def __init__(self):
        self.previous_statistic_gathered_gold = 0

    def reset(self):
        self.__init__()


class Scenario(gym.Env):
    DEFAULTS = dict(
        updates_per_action=1,
        flatten=False
    )

    def __init__(self, config, game, *scenarios):
        self.config = util.dict_update(Scenario.DEFAULTS.copy(), config)
        self.game = game

        # Struct that holds all data that needs to be stored during a episode
        self.data = ScenarioData()

        # Binds scenario requirements to *this*
        self.scenarios = [partial(scenario, self) for scenario in scenarios]

        # Define the action space
        self.action_space = LimitedDiscrete(Constants.action_min, Constants.action_max)

        # Define the observation space, here we assume that max is 255 (image) # TODO
        self.observation_space = gym.spaces.Box(0, 255, shape=self.get_state().shape, dtype=np.float32)

    def evaluate(self):
        success, reward = zip(*[scenario() for scenario in self.scenarios])
        return all(success), sum(reward)

    def _optimal_play_sequence(self):