# Path: oscar/env/envs/general_learning_env.py
# class GeneralLearningEnv(gym.Env):
#     def __init__(self, configuration_file=DEFAULT_CONFIGURATION, enable_visualisation=True, game_steps_per_update=8,
#                  log_file_path=None, publish_stats=True):
#         self.env_thread = Pysc2EnvRunner(configuration_file=configuration_file,
#                                          enable_visualisation=enable_visualisation,
#                                          game_steps_per_update=game_steps_per_update)
#         self.shared_memory = self.env_thread.shared_memory
#         self.log_file = log_file_path
#         self.publish_stats = publish_stats
#         # get semaphores from shared memory
#         self.semaphore_obs_ready = self.shared_memory.semaphore_obs_ready
#         self.semaphore_action_set = self.shared_memory.semaphore_action_set
#         # set action/observation space from shared memory
#         self.action_space = self.shared_memory.action_space
#         self.observation_space = self.shared_memory.observation_space
#         # start thread
#         self.env_thread.start()
# 
#     def step(self, action):
#         info_dict = {}
#         # set action into shared memory
#         self.shared_memory.shared_action = action
#         # let the thread run again
#         self.semaphore_action_set.release()
#         # increase the number of steps performed by the learning agent
#         self.env_thread.learning_agent_step += 1
#         # wait a call to the target agent
#         self.semaphore_obs_ready.acquire(blocking=True, timeout=None)
#         # get the obs of the current state as seen by the agent
#         obs = self.shared_memory.shared_obs
#         # get the reward of the current run and reset it to 0 for next step
#         reward = self.env_thread.reward  # / self.env_thread.step_count
#         self.env_thread.reward = 0
#         # get done state from thread
#         done = self.env_thread.was_done
#         if done:
#             self.env_thread.was_done = False  # reset was_done to false for next loop
#             if self.publish_stats:
#                 info_dict["stats"] = self.env_thread.stats
#             # log result to disk if asked
#             if self.log_file is not None and self.env_thread.stats is not None:
#                 if os.path.isfile(self.log_file):
#                     self.env_thread.stats.to_csv(self.log_file, sep=',', mode='a', header=False)
#                 else:
#                     self.env_thread.stats.to_csv(self.log_file, sep=',', mode='w', header=True)
#         # return current obs
#         return obs.copy(), reward, done, info_dict
# 
#     def reset(self):
#         # wait for the env to stop as waiting action
#         self.semaphore_obs_ready.acquire(blocking=True, timeout=None)
#         return self.shared_memory.shared_obs
# 
#     def close(self):
#         self.env_thread.stop = True
#         self.semaphore_action_set.release()  # to be sure that the thread go to stop condition
# 
#     def render(self, mode='human', close=False):
#         pass
# 
#     def seed(self, seed=None):
#         pass
# 
#     def get_action_mask(self):
#         # every action is supposed to be playable
#         # return np.ones(shape=self.action_space.n)
#         # get the actions available according to the learning agent
#         return self.shared_memory.available_action_mask


import numpy as np
import gym
from keras.models import Sequential
from keras.layers import Dense, Activation, Flatten
from keras.optimizers import Adam
from rl.agents.dqn import DQNAgent
from rl.policy import BoltzmannQPolicy, LinearAnnealedPolicy
from rl.memory import SequentialMemory
from oscar.env.envs.general_learning_env import GeneralLearningEnv
from ML_homework.policy_iteration.complex_policy_iteration import policy_iteration_iterator, state_from_obs
print(model.summary())

# Finally, we configure and compile our agent. You can use every built-in Keras optimizer and
# even the metrics!
memory = SequentialMemory(limit=50000, window_length=1)
boltzmann_policy = BoltzmannQPolicy(tau=20.0, clip=(0.0, 500.0))
policy = LinearAnnealedPolicy(boltzmann_policy, 'tau', 1.0, 20.0, 1.0, 10000)
# enable the dueling network
# you can specify the dueling_type to one of {'avg','max','naive'}
dqn = DQNAgent(model=model, nb_actions=nb_actions, memory=memory, nb_steps_warmup=10,
               enable_dueling_network=True, dueling_type='avg', target_model_update=1e-2, policy=policy)
dqn.compile(Adam(lr=1e-3), metrics=['mae'])

# warm up
pi = None
for p in policy_iteration_iterator(10, 0.5, file_path="/tmp/state_table.csv", save_path="/tmp/OSCAR/"):
    pi = p

for i in range(20):
    obs = env.reset()
    while True:
        s = state_from_obs(obs)
        a = pi[s.id()]
        old_obs = obs
        obs, r, done, debug_dict = env.step(a)
        memory.append(old_obs, a, r, done, False)
        if done:
            break

env.close()

# Path: oscar/env/envs/general_learning_env.py
# class GeneralLearningEnv(gym.Env):
#     def __init__(self, configuration_file=DEFAULT_CONFIGURATION, enable_visualisation=True, game_steps_per_update=8,
#                  log_file_path=None, publish_stats=True):
#         self.env_thread = Pysc2EnvRunner(configuration_file=configuration_file,
#                                          enable_visualisation=enable_visualisation,
#                                          game_steps_per_update=game_steps_per_update)
#         self.shared_memory = self.env_thread.shared_memory
#         self.log_file = log_file_path
#         self.publish_stats = publish_stats
#         # get semaphores from shared memory
#         self.semaphore_obs_ready = self.shared_memory.semaphore_obs_ready
#         self.semaphore_action_set = self.shared_memory.semaphore_action_set
#         # set action/observation space from shared memory
#         self.action_space = self.shared_memory.action_space
#         self.observation_space = self.shared_memory.observation_space
#         # start thread
#         self.env_thread.start()
# 
#     def step(self, action):
#         info_dict = {}
#         # set action into shared memory
#         self.shared_memory.shared_action = action
#         # let the thread run again
#         self.semaphore_action_set.release()
#         # increase the number of steps performed by the learning agent
#         self.env_thread.learning_agent_step += 1
#         # wait a call to the target agent
#         self.semaphore_obs_ready.acquire(blocking=True, timeout=None)
#         # get the obs of the current state as seen by the agent
#         obs = self.shared_memory.shared_obs
#         # get the reward of the current run and reset it to 0 for next step
#         reward = self.env_thread.reward  # / self.env_thread.step_count
#         self.env_thread.reward = 0
#         # get done state from thread
#         done = self.env_thread.was_done
#         if done:
#             self.env_thread.was_done = False  # reset was_done to false for next loop
#             if self.publish_stats:
#                 info_dict["stats"] = self.env_thread.stats
#             # log result to disk if asked
#             if self.log_file is not None and self.env_thread.stats is not None:
#                 if os.path.isfile(self.log_file):
#                     self.env_thread.stats.to_csv(self.log_file, sep=',', mode='a', header=False)
#                 else:
#                     self.env_thread.stats.to_csv(self.log_file, sep=',', mode='w', header=True)
#         # return current obs
#         return obs.copy(), reward, done, info_dict
# 
#     def reset(self):
#         # wait for the env to stop as waiting action
#         self.semaphore_obs_ready.acquire(blocking=True, timeout=None)
#         return self.shared_memory.shared_obs
# 
#     def close(self):
#         self.env_thread.stop = True
#         self.semaphore_action_set.release()  # to be sure that the thread go to stop condition
# 
#     def render(self, mode='human', close=False):
#         pass
# 
#     def seed(self, seed=None):
#         pass
# 
#     def get_action_mask(self):
#         # every action is supposed to be playable
#         # return np.ones(shape=self.action_space.n)
#         # get the actions available according to the learning agent
#         return self.shared_memory.available_action_mask


# Path: learning_tools/keras-rl/dqn/dqn_complex.py
import numpy as np
import gym
from keras.models import Sequential
from keras.layers import Dense, Activation, Flatten
from keras.optimizers import Adam
from rl.agents.dqn import DQNAgent
from rl.policy import BoltzmannQPolicy, LinearAnnealedPolicy
from rl.memory import SequentialMemory
from oscar.env.envs.general_learning_env import GeneralLearningEnv
from ML_homework.policy_iteration.complex_policy_iteration import policy_iteration_iterator, state_from_obs

print(model.summary())

# Finally, we configure and compile our agent. You can use every built-in Keras optimizer and
# even the metrics!
memory = SequentialMemory(limit=50000, window_length=1)
boltzmann_policy = BoltzmannQPolicy(tau=20.0, clip=(0.0, 500.0))
policy = LinearAnnealedPolicy(boltzmann_policy, 'tau', 1.0, 20.0, 1.0, 10000)
# enable the dueling network
# you can specify the dueling_type to one of {'avg','max','naive'}
dqn = DQNAgent(model=model, nb_actions=nb_actions, memory=memory, nb_steps_warmup=10,
               enable_dueling_network=True, dueling_type='avg', target_model_update=1e-2, policy=policy)
dqn.compile(Adam(lr=1e-3), metrics=['mae'])

# warm up
pi = None
for p in policy_iteration_iterator(10, 0.5, file_path="/tmp/state_table.csv", save_path="/tmp/OSCAR/"):
    pi = p

for i in range(20):
    obs = env.reset()
    while True:
        s = state_from_obs(obs)
        a = pi[s.id()]
        old_obs = obs
        obs, r, done, debug_dict = env.step(a)
        memory.append(old_obs, a, r, done, False)
        if done:
            break

env.close()

env = GeneralLearningEnv(CONFIG_FILE, False, log_file_path=LOG_FILE, publish_stats=False)