from collections import deque
from gym import spaces
import cv2
cv2.ocl.setUseOpenCL(False)

import torch

import numpy as np
import gym
from gym.wrappers import TimeLimit, Monitor

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

verbose=False

class NoopResetEnv(gym.Wrapper):
    def __init__(self, env, noop_max=30):
        assert False
        """Sample initial states by taking random number of no-ops on reset.
        No-op is assumed to be action 0.
        """
        gym.Wrapper.__init__(self, env)
        self.noop_max = noop_max
        self.override_num_noops = None
        self.noop_action = 0
        assert env.unwrapped.get_action_meanings()[0] == 'NOOP'

    def reset(self, **kwargs):
        if verbose:
            print("noop reset")
            print(self.env.unwrapped.ale.lives())
        """ Do no-op action for a number of steps in [1, noop_max]."""
        self.env.reset(**kwargs)
        if self.override_num_noops is not None:
            noops = self.override_num_noops
        else:
            noops = self.unwrapped.np_random.randint(1, self.noop_max + 1) #pylint: disable=E1101
        assert noops > 0
        obs = None
        for _ in range(noops):
            obs, _, done, _ = self.env.step(self.noop_action)
            if done:
                obs = self.env.reset(**kwargs)
        return obs

    def step(self, ac):
        return self.env.step(ac)

class FireResetEnv(gym.Wrapper):
    def __init__(self, env):
        assert False
        """Take action on reset for environments that are fixed until firing."""
        gym.Wrapper.__init__(self, env)
        assert env.unwrapped.get_action_meanings()[1] == 'FIRE'
        assert len(env.unwrapped.get_action_meanings()) >= 3

    def reset(self, **kwargs):
        if verbose:
            print("fire reset")
            print(self.env.unwrapped.ale.lives())
        self.env.reset(**kwargs)
        obs, _, done, _ = self.env.step(1)
        if done:
            self.env.reset(**kwargs)
        obs, _, done, _ = self.env.step(2)
        if done:
            self.env.reset(**kwargs)
        return obs

    def step(self, ac):
        return self.env.step(ac)

class EpisodicLifeEnv(gym.Wrapper):
    def __init__(self, env):
        """Make end-of-life == end-of-episode, but only reset on true game over.
        Done by DeepMind for the DQN and co. since it helps value estimation.
        """
        gym.Wrapper.__init__(self, env)
        self.lives = 0
        self.was_real_done  = True

    def step(self, action):
        obs, reward, done, info = self.env.step(action)

        self.was_real_done = done
        # check current lives, make loss of life terminal,
        # then update lives to handle bonus lives
        lives = self.env.unwrapped.ale.lives()
        if lives < self.lives and lives > 0:
            # for Qbert sometimes we stay in lives == 0 condition for a few frames
            # so it's important to keep lives > 0, so that we only reset once
            # the environment advertises done.
            done = True
        self.lives = lives
        return obs, reward, done, info

    def reset(self, **kwargs):
        """Reset only when lives are exhausted.
        This way all states are still reachable even though lives are episodic,
        and the learner need not know about any of this behind-the-scenes.
        """
        if self.was_real_done:
            obs = self.env.reset(**kwargs)
        else:
            # no-op step to advance from terminal/lost life state
            obs, _, _, _ = self.env.step(0)
        self.lives = self.env.unwrapped.ale.lives()
        return obs

class MaxAndSkipEnv(gym.Wrapper):
    def __init__(self, env, skip=4, sticky_prob=0):
        """Return only every `skip`-th frame"""
        gym.Wrapper.__init__(self, env)
        # most recent raw observations (for max pooling across time steps)
        self._obs_buffer = np.zeros((2,)+env.observation_space.shape, dtype=np.uint8)
        self._skip       = skip
        self.sticky_prob = sticky_prob
        self.last_action = None

    def step(self, action):
        """Repeat action, sum reward, and max over last observations."""
        # print("step max skip")
        total_reward = 0.0
        done = None
        for i in range(self._skip):
            if self.sticky_prob>0 and self.last_action is not None:
                if self.unwrapped.np_random.random_sample()<self.sticky_prob:
                    action_selected = self.last_action
                else:
                    action_selected = action
                obs, reward, done, info = self.env.step(action_selected)
                self.last_action = action_selected
            else:
                obs, reward, done, info = self.env.step(action)
                self.last_action = action
            # print(info)
            if i == self._skip - 2: self._obs_buffer[0] = obs
            if i == self._skip - 1: self._obs_buffer[1] = obs
            total_reward += reward
            if done:
                break
        # print(f"done maxskip {done}")
        # Note that the observation on the done=True frame
        # doesn't matter
        max_frame = self._obs_buffer.max(axis=0)

        return max_frame, total_reward, done, info

    def reset(self):
        self.last_action = None
        ob = self.env.reset()
        return ob

class ClipRewardEnv(gym.RewardWrapper):
    def __init__(self, env):
        gym.RewardWrapper.__init__(self, env)

    def reward(self, reward):
        """Bin reward to {+1, 0, -1} by its sign."""
        return np.sign(reward)

class WarpFrame(gym.ObservationWrapper):
    def __init__(self, env, width=84, height=84, grayscale=True, dict_space_key=None):
        """
        Warp frames to 84x84 as done in the Nature paper and later work.
        If the environment uses dictionary observations, `dict_space_key` can be specified which indicates which
        observation should be warped.
        """
        super().__init__(env)
        self._width = width
        self._height = height
        self._grayscale = grayscale
        self._key = dict_space_key
        if self._grayscale:
            num_colors = 1
        else:
            num_colors = 3

        new_space = gym.spaces.Box(
            low=0,
            high=255,
            shape=(self._height, self._width, num_colors),
            dtype=np.uint8,
        )
        if self._key is None:
            original_space = self.observation_space
            self.observation_space = new_space
        else:
            original_space = self.observation_space.spaces[self._key]
            self.observation_space.spaces[self._key] = new_space
        assert original_space.dtype == np.uint8 and len(original_space.shape) == 3

    def observation(self, obs):
        if self._key is None:
            frame = obs
        else:
            frame = obs[self._key]

        if self._grayscale:
            frame = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
        frame = cv2.resize(
            frame, (self._width, self._height), interpolation=cv2.INTER_AREA
        )
        if self._grayscale:
            frame = np.expand_dims(frame, -1)

        if self._key is None:
            obs = frame
        else:
            obs = obs.copy()
            obs[self._key] = frame
        return obs

class FrameStack(gym.Wrapper):
    def __init__(self, env, k):
        """Stack k last frames.
        Returns lazy array, which is much more memory efficient.
        See Also
        --------
        baselines.common.atari_wrappers.LazyFrames
        """
        gym.Wrapper.__init__(self, env)
        self.k = k
        self.frames = deque([], maxlen=k)
        shp = env.observation_space.shape
        self.observation_space = spaces.Box(low=0, high=255, shape=(shp[:-1] + (shp[-1] * k,)), dtype=env.observation_space.dtype)

    def reset(self):
        if verbose:
            print("framestack reset")
            print(self.env.unwrapped.ale.lives())
        # print("franestack reset") #TODO remove
        ob = self.env.reset()
        for _ in range(self.k):
            self.frames.append(ob)
        return self._get_ob()

    def step(self, action):
        if verbose:
            print("step max skip")
        ob, reward, done, info = self.env.step(action)
        self.frames.append(ob)
        if verbose:
            print(f"done frame stack {done}")
        return self._get_ob(), reward, done, info

    def _get_ob(self):
        assert len(self.frames) == self.k
        return LazyFrames(list(self.frames))

class PacManReset(gym.Wrapper):
    """Do 60 steps of no-op (Pacman wait around 60 steps before starting) and take a certain number of random actions
    to randomize the initial state
    """
    def __init__(self, env, max_nb_random_actions = 0, use_only_first_action = True):
        gym.Wrapper.__init__(self, env)
        self.max_nb_random_actions = max_nb_random_actions
        self.use_only_first_action = use_only_first_action


    def reset(self):
        if verbose:
            print("pacman reset")
            print(self.env.unwrapped.ale.lives())
        ob = self.env.reset()
        for _ in range(60):
            ob, reward, done, info = self.env.step(0)
            if done:
                assert False
        if self.max_nb_random_actions>0:
            nb_random_actions = self.unwrapped.np_random.randint(1, self.max_nb_random_actions + 1)
            for _ in range(nb_random_actions):
                if not self.use_only_first_action:
                    action = self.env.action_space.sample()
                else:
                    action = 0
                ob, reward, done, info = self.env.step(action)
                if done:
                    assert False
        return ob

class PacManDone(gym.Wrapper):
    def __init__(self, env):
        assert False
        """Reset the environment after only one life
        """
        gym.Wrapper.__init__(self, env)
        self.lives = None

    def step(self, action):
        if verbose:
            print("step episodic life")
        obs, reward, done, info = self.env.step(action)
        lives = self.env.unwrapped.ale.lives()
        if verbose:
            print(f"done episodic {done}")
        if lives < self.lives:
            if verbose:
                print("lives")
            assert lives==2
            done = True
        self.lives = lives
        return obs, reward, done, info

    def reset(self, **kwargs):
        assert self.lives is None or self.lives==2
        if verbose:
            print("episodic life reset")
            print(self.env.unwrapped.ale.lives())
        obs = self.env.reset(**kwargs)
        self.lives = self.env.unwrapped.ale.lives()
        assert self.lives==3
        return obs

class ScaledFloatFrame(gym.ObservationWrapper):
    def __init__(self, env):
        assert False
        gym.ObservationWrapper.__init__(self, env)
        self.observation_space = gym.spaces.Box(low=0, high=1, shape=env.observation_space.shape, dtype=np.float32)

    def observation(self, observation):
        # careful! This undoes the memory optimization, use
        # with smaller replay buffers only.
        return np.array(observation).astype(np.float32) / 255.0

class LazyFrames(object):
    def __init__(self, frames):
        """This object ensures that common frames between the observations are only stored once.
        It exists purely to optimize memory usage which can be huge for DQN's 1M frames replay
        buffers.
        This object should only be converted to numpy array before being passed to the model.
        You'd not believe how complex the previous solution was."""
        self._frames = frames
        self._out = None

    def _force(self):
        if self._out is None:
            self._out = np.concatenate(self._frames, axis=-1)
            self._frames = None
        return self._out

    def __array__(self, dtype=None):
        out = self._force()
        if dtype is not None:
            out = out.astype(dtype)
        return out

    def __len__(self):
        return len(self._force())

    def __getitem__(self, i):
        return self._force()[i]

    def count(self):
        frames = self._force()
        return frames.shape[frames.ndim - 1]

    def frame(self, i):
        return self._force()[..., i]

from multiprocessing import Process, Pipe
import gym

def worker(conn, env):
    while True:
        cmd, data = conn.recv()
        if cmd == "step":
            obs, reward, done, info = env.step(data)
            if done:
                obs = env.reset()
            conn.send((obs, reward, done, info))
        # elif cmd == "stepnoreset":
        #     action, init_state, init_obs = data
        #     obs, reward, done, info = env.step(action)
        #     if done:
        #         obs = init_obs
        #         env.restore_state(init_state)
        #     conn.send((obs, reward, done, info))
        elif cmd == "stepnoresetfull":
            action, init_state, init_obs = data
            obs, reward, done, info = env.step(action)
            if done:
                obs = init_obs
                env.restore_full_state(init_state)
            conn.send((obs, reward, done, info))
        elif cmd == "reset":
            obs = env.reset()
            conn.send(obs)
        # elif cmd == "clone_state":
        #     state = env.clone_state()
        #     conn.send(state)
        # elif cmd == "restore_state":
        #     env.restore_state(data)
        #     conn.send(None)
        # elif cmd == "clone_full_state":
        #     state = env.clone_full_state()
        #     conn.send(state)
        elif cmd == "restore_full_state":
            env.restore_full_state(data)
            conn.send(None)
        else:
            raise NotImplementedError

class ParallelEnv(gym.Env):
    """A concurrent execution of environments in multiple processes."""

    def __init__(self, envs):
        assert len(envs) >= 1, "No environment given."

        self.envs = envs
        self.observation_space = self.envs[0].observation_space
        self.action_space = self.envs[0].action_space
        self.init_state = None
        self.init_obs = None

        self.locals = []
        # import pdb; pdb.set_trace()
        for env in self.envs[1:]:
            local, remote = Pipe()
            self.locals.append(local)
            p = Process(target=worker, args=(remote, env))
            p.daemon = True
            p.start()
            remote.close()

    def reset(self):
        for local in self.locals:
            local.send(("reset", None))
        results = [self.envs[0].reset()] + [local.recv() for local in self.locals]
        return results

    # def reset_one(self, i):
    #     if i==0:
    #         self.envs[0].reset()
    #     for local in self.locals:
    #         local.send(("reset", None))
    #     results = [self.envs[0].reset()] + [local.recv() for local in self.locals]
    #     return results

    def step(self, actions):
        # import pdb; pdb.set_trace()

        obs, reward, done, info = self.envs[0].step(actions[0])

        if done:
            if self.init_state is None:
                obs = self.envs[0].reset()
            else:
                # if len(self.locals)>0:
                #     assert False
                obs = self.init_obs
                try:
                    self.envs[0].restore_full_state(self.init_state)
                except:
                    import pdb; pdb.set_trace()

        if len(self.locals)==0:
            return (obs,), (reward,), (done,), (info,)
        for local, action in zip(self.locals, actions[1:]):
            # if self.init_state is not None:
            #     import pdb;
            #     pdb.set_trace()
            local.send(("step" if self.init_state is None else "stepnoresetfull",)+ (action if self.init_state is None else (action, self.init_state, self.init_obs),))

        # import pdb;pdb.set_trace()
        p = [local.recv() for local in self.locals]
        # try:
        results = zip(*[(obs, reward, done, info)] + p)
        # except:
        #     import pdb; pdb.set_trace()
        # import pdb; pdb.set_trace()
        return results

    # def clone_state(self):
    #     for local in self.locals:
    #         local.send(("clone_state", None))
    #     states = [self.envs[0].clone_state()] + [local.recv() for local in self.locals]
    #     return states
    #
    # def restore_state(self, states):
    #     self.envs[0].restore_state(states[0])
    #     for local, state in zip(self.locals, states[1:]):
    #         local.send(("restore_state", state))
    #     for local in self.locals:
    #         local.recv()

    def set_init_state(self, init_state):
        self.init_state = init_state

    def set_init_obs(self, init_obs):
        self.init_obs = init_obs

    def clone_full_state(self):
        assert len(self.locals)==0
        return self.envs[0].clone_full_state()
        # for local in self.locals:
        #     local.send(("clone_full_state", None))
        # states = [self.envs[0].clone_full_state()] + [local.recv() for local in self.locals]
        # return states

    # def restore_full_state(self, states):
    #     print("restore full state")
    #     self.envs[0].restore_full_state(states[0])
    #     for local, state in zip(self.locals, states[1:]):
    #         local.send(("restore_full_state", state))
    #     for local in self.locals:
    #         print(local)
    #         local.recv()
    #     print("restore done")

    def restore_full_state(self):
        assert self.init_state is not None
        # print("restore full state")
        # import pdb; pdb.set_trace()
        self.envs[0].restore_full_state(self.init_state)
        for local in self.locals:
            # print(local)
            local.send(("restore_full_state", self.init_state))
            # local.recv()
        for local in self.locals:
            # print(local)
            local.recv()
        # print("restore done")

    def restore_full_state_one(self,i):
        assert self.init_state is not None
        if i==0:
            self.envs[0].restore_full_state(self.init_state)
        else:
            self.locals[i-1].send(("restore_full_state", self.init_state))
            self.locals[i-1].recv()

    def render(self, mode='rgb_array', caption=None):
        if mode == 'human':
            return self.envs[0].render(mode='human')
        elif mode == 'rgb_array':
            img = self.envs[0].render(mode='rgb_array')
            if caption is None:
                return img
            else:
                # setup text
                font = cv2.FONT_HERSHEY_SIMPLEX
                text = caption

                # get boundary of this text
                textsize = cv2.getTextSize(text, font, 1, 2)[0]

                # get coords based on boundary
                textX = int((img.shape[1] - textsize[0]) / 2)
                textY = int(textsize[1])

                # add text centered on image
                cv2.putText(img, text, (textX, textY), font, 0.8, (255, 255, 255), 2)
                return img

class PyTorchWrapper(gym.Wrapper):
    def __init__(self, env):
        gym.Wrapper.__init__(self, env)

    def reset(self):
        if verbose:
            print("pytorchwrap reset")
            print(self.env.unwrapped.ale.lives())
        obs = self.env.reset()
        if isinstance(obs,np.ndarray):
            obs = np.expand_dims(obs, axis=0)
            obs = torch.from_numpy(obs).float().to(device)
            return obs
        else:
            assert False, 'obs should be an array'
            obs = {key:np.expand_dims(obs[key], axis=0) for key in obs}
            obs = {key:torch.from_numpy(obs[key]).float().to(device) for key in obs}
            return obs

    def step(self, actions):
        if verbose:
            print("step pytorch wrap")
        # actions = actions.cpu().numpy()
        actions = actions.item()
        obs, reward, done, info = self.env.step(actions)
        if isinstance(obs, np.ndarray):
            obs = np.expand_dims(obs, axis=0)
            obs = torch.from_numpy(obs).float().to(device)
        else:
            assert False, 'obs should be an array'
            obs = {key: np.expand_dims(obs[key], axis=0) for key in obs}
            obs = {key: torch.from_numpy(obs[key]).float().to(device) for key in obs}
        reward = torch.from_numpy(np.array([reward])).float()
        if verbose:
            print(f"done pyt wrap {done}")
        return obs, reward, [[done]], info

    def change_seed(self, seed):
        self.env.seed(seed)
        self.env.action_space.seed(seed)
        self.env.observation_space.seed(seed)

def wrap_atari(env, max_episode_steps=None):
    # import pdb; pdb.set_trace()
    assert 'NoFrameskip' in env.spec.id
    # env = NoopResetEnv(env, noop_max=30)
    env = MaxAndSkipEnv(env, skip=4)

    assert max_episode_steps is None

    return env

def wrap_deepmind(env, episode_life=True, clip_rewards=True, frame_stack=False, scale=False):
    """Configure environment for DeepMind-style Atari.
    """
    env = PacManReset(env)
    if episode_life:
        env = EpisodicLifeEnv(env)
    if 'FIRE' in env.unwrapped.get_action_meanings():
        env = FireResetEnv(env)
    env = WarpFrame(env)
    if scale:
        env = ScaledFloatFrame(env)
    if clip_rewards:
        env = ClipRewardEnv(env)
    if frame_stack:
        env = FrameStack(env, 4)

    return env

class ImageToPyTorch(gym.ObservationWrapper):
    """
    Image shape to channels x weight x height
    """

    def __init__(self, env):
        super(ImageToPyTorch, self).__init__(env)
        old_shape = self.observation_space.shape
        self.observation_space = gym.spaces.Box(
            low=0,
            high=255,
            shape=(old_shape[-1], old_shape[0], old_shape[1]),
            dtype=np.uint8,
        )

    def observation(self, observation):
        return np.transpose(observation, axes=(2, 0, 1))

def wrap_pytorch(env):
    return ImageToPyTorch(env)

class VecPyTorch(gym.Wrapper):
    def __init__(self, env, device):
        super(VecPyTorch, self).__init__(env)
        self.device = device

    def reset(self):
        if verbose:
            print("vec reset")
            print(self.env.envs[0].unwrapped.ale.lives())
        obs = self.env.reset()
        nobs = np.array(obs)
        tensor_obs = torch.tensor(nobs, device=self.device, dtype=torch.float)
        return tensor_obs, obs[0]

    def step(self, actions):
        try:
            actions = actions.cpu().numpy()
        except:
            # import pdb; pdb.set_trace()
            pass
        # try:
        obs, reward, done, info = self.env.step(actions)
        # except:
        #     import pdb; pdb.set_trace()
        nobs = np.array(obs)
        try:
            tensor_obs = torch.tensor(nobs, device=self.device, dtype=torch.float)
        except:
            import pdb; pdb.set_trace()
        reward = np.array(reward)
        reward = torch.from_numpy(reward).unsqueeze(dim=1).float()

        if info is None:
            info = [{"obs": obs[0]}]
        elif len(info)==1:
            info[0]["obs"]=obs[0]
        return tensor_obs, reward, done, info
    #
    # def change_seed(self, seeds):
    #     for index,seed in enumerate(seeds):
    #         # import pdb; pdb.set_trace()
    #         self.venv.envs[index].seed(seed)
    #         self.venv.envs[index].action_space.seed(seed)
    #         self.venv.envs[index].observation_space.seed(seed)

class CaptionWrapper(gym.Wrapper):
    def __init__(self, env):
        assert False
        gym.Wrapper.__init__(self, env)

    def render(self, mode='rgb_array', caption=None):
        if mode == 'human':
            return self.env.render(mode='human')
        elif mode == 'rgb_array':
            img = self.env.render(mode='rgb_array')
            if caption is None:
                return img
            else:
                if isinstance(caption, list):
                    nb_str = len(caption)
                    # setup text
                    font = cv2.FONT_HERSHEY_PLAIN
                    texts = caption
                    # get boundary of this text
                    textsizes = [cv2.getTextSize(text, font, 0.8, 1)[0] for text in texts]
                    textXs = [int((img.shape[1] - textsize[0]) / 2) for textsize in textsizes]
                    textYs = [img.shape[0]-(len(textsizes)-j-1)*textsize[1] for j,textsize in enumerate(textsizes)]
                    for i,text in enumerate(texts):
                        cv2.putText(img, text, (textXs[i], textYs[i]), font, 0.8, (255,255,255), 1)
                else:
                    # setup text
                    font = cv2.FONT_HERSHEY_PLAIN
                    text = caption

                    # get boundary of this text
                    textsize = cv2.getTextSize(text, font, 0.8, 1)[0]

                    # get coords based on boundary
                    textX = int((img.shape[1] - textsize[0]) / 2)
                    textY = int(textsize[1])

                    # add text centered on image
                    cv2.putText(img, text, (textX, textY), font, 0.8, (255,255,255), 1)
                return img

def make_env(gym_id, seed=None, sticky_prob=0, max_nb_random_actions =0, use_only_first_action=True, reward_clipping = False):
    env = gym.make(gym_id)

    env = MaxAndSkipEnv(env, sticky_prob=sticky_prob) #Step, Modifies Observation and Action
    env = gym.wrappers.RecordEpisodeStatistics(env)
    env = PacManReset(env,max_nb_random_actions =max_nb_random_actions, use_only_first_action=use_only_first_action) #Reset, Modifies Observation and Action Sequence
    env = EpisodicLifeEnv(env) #Step, Modifies Done
    # env = gym.wrappers.RecordEpisodeStatistics(env)

    env = WarpFrame(env) #Modify Observation Format
    if reward_clipping:
        env = ClipRewardEnv(env)
    env = FrameStack(env, 4) #Modify Observation and action seqience
    env = ImageToPyTorch(env) #Modifies Observation format

    if seed is not None:
        env.seed(seed)
        env.action_space.seed(seed)
        env.observation_space.seed(seed)
    return env

def make_one_env(gym_id, idx, capture_video, experiment_name, seed=None):
    env = gym.make(gym_id)
    # import pdb; pdb.set_trace()
    env = wrap_atari(env)
    env = gym.wrappers.RecordEpisodeStatistics(env)
    if capture_video:
        if idx == 0:
            env = Monitor(env, f'videos/{experiment_name}')
    # env = ProbsVisualizationWrapper(env)
    env = wrap_pytorch(
        wrap_deepmind(
            env,
            clip_rewards=True,
            frame_stack=True,
            scale=False,
        )
    )

    if seed is not None:
        env.seed(seed)
        env.action_space.seed(seed)
        env.observation_space.seed(seed)

    env = PyTorchWrapper(env)

    return env
