from stable_baselines3 import PPO

from coopihc import (
    State,
    Bundle,
    RLPolicy,
    BasePolicy,
    discrete_array_element,
)
from coopihc.bundle.wrappers.Train import TrainGym

from coopihczoo import SimplePointingTask, ConstantCDGain, CarefulPointer


import numpy
import gym
from gym.wrappers import FilterObservation
import copy


# wrappers used during training
TEN_EPSILON32 = 10 * numpy.finfo(numpy.float32).eps


class NormalizeActionWrapper(gym.ActionWrapper):
    def __init__(self, env):
        super().__init__(env)

        self.action_space = gym.spaces.Box(
            low=-1,
            high=1,
            shape=(1,),
            dtype=numpy.float32,
        )

    def action(self, action):
        return {
            "user_action": int(
                numpy.around(action * 11 / 2 - TEN_EPSILON32, decimals=0)
            )
        }

    def reverse_action(self, action):
        return numpy.array((action["user_action"] - 5.0) / 11.0 * 2).astype(
            numpy.float32
        )


# action_state used during training
action_state = State()
action_state["action"] = discrete_array_element(low=-5, high=5)

# Wrapperless Env used for training
task = SimplePointingTask(gridsize=31, number_of_targets=8)
unitcdgain = ConstantCDGain(1)

action_state = State()
action_state["action"] = discrete_array_element(low=-5, high=5)

user = CarefulPointer(override_policy=(BasePolicy, {"action_state": action_state}))
bundle = Bundle(task=task, user=user, assistant=unitcdgain)
observation = bundle.reset(go_to=1)
env = TrainGym(
    bundle,
    train_user=True,
    train_assistant=False,
)
from gym.wrappers import FilterObservation, FlattenObservation

# [start-load-policy]
model_path = "saved_model.zip"
learning_algorithm = "PPO"


class FilledFilterObservation(FilterObservation):
    def __init__(self, env):
        super().__init__(env, filter_keys=("position", "goal"))


wrappers = {
    "observation_wrappers": [FilledFilterObservation, FlattenObservation],
    "action_wrappers": [NormalizeActionWrapper],
}
library = "stable_baselines3"

trained_policy = RLPolicy(
    action_state, model_path, learning_algorithm, env, wrappers, library
)
# [end-load-policy]
# exit()
# [start-viz-policy]
distance_to_goal = []
actions = []
for i in range(1000):
    bundle.reset(go_to=3)
    obs = bundle.user.observation
    distance_to_goal.append(
        obs["user_state"]["goal"][0] - obs["task_state"]["position"][0]
    )
    action, reward = trained_policy.sample(observation=obs)
    actions.append(action[0])

import matplotlib.pyplot as plt

fig = plt.figure()
ax = fig.add_subplot(111)
ax.plot(distance_to_goal, actions, "o", label="action")
ax.set_xlabel("algebraic distance to goal (steps)")
ax.set_ylabel("actions selected by the RLPolicy")
ax.legend()
plt.tight_layout()
# plt.show()
# [end-viz-policy]

# [start-play-policy]
task = SimplePointingTask(gridsize=31, number_of_targets=8)
unitcdgain = ConstantCDGain(1)

user = CarefulPointer(override_policy=(trained_policy, {}))
bundle = Bundle(task=task, user=user, assistant=unitcdgain)

bundle.reset()
bundle.render("plotext")
while True:
    obs, rew, flag = bundle.step()
    bundle.render("plotext")
    if flag:
        break
# [end-play-policy]
