# This file may not be shared/redistributed without permission. Please read copyright notice in the git repo. If this file contains other copyright notices disregard this text.
from gymnasium import logger
from irlc.ex01.agent import Agent
import time
import sys
import gymnasium as gym
import os
import asyncio
import numpy as np

HUMAN_REQUEST_RESET = 'reset the environment'

try:
    # Using this backend apparently clash with scientific mode. Not sure why it was there in the first place so
    # disabling it for now.
    # matplotlib.use('TkAgg')
    import matplotlib.pyplot as plt
except ImportError as e:
    logger.warn('failed to set matplotlib backend, plotting will not work: %s' % str(e))
    plt = None

try:
    import pygame
    from pygame.event import Event
except ImportError as e:
    logger.warn('failed to import pygame. Many of the interactive visualizations will not work: %s' % str(e))
    Event = None

class AgentWrapper(Agent):
    """Wraps the environment to allow a modular transformation.

    This class is the base class for all wrappers. The subclass could override
    some methods to change the behavior of the original environment without touching the
    original code.

    .. note::

        Don't forget to call ``super().__init__(env)`` if the subclass overrides :meth:`__init__`.

    """
    def __init__(self, agent, env):
        self.agent = agent
        self.env = env

    def __getattr__(self, name):
        if name.startswith('_'):
            raise AttributeError("attempted to get missing private attribute '{}'".format(name))
        return getattr(self.agent, name)

    @classmethod
    def class_name(cls):
        return cls.__name__

    def pi(self, state, k, info=None):
        return self.agent.pi(state, k, info)
        # return self.env.step(action)

    def train(self, *args, **kwargs):
        return self.agent.train(*args, **kwargs)

    def __str__(self):
        return '<{}{}>'.format(type(self).__name__, self.agent)

    def __repr__(self):
        return str(self)

    @property
    def unwrapped(self):
        return self.agent.unwrapped

PAUSE_KEY = ord('p')
SPACEBAR = "_SPACE_BAR_PRESSED_"
HUMAN_DEMAND_RESET = "reset env."

async def _webassembly_interactive(env, agent, autoplay=False):
    from pygame import gfxdraw
    import numpy as np

    def aapolygon(surface, points, color):
        pygame.draw.polygon(surface, color, points)

    def filled_polygon(surface, points, color):
        pygame.draw.polygon(surface, color, points, width=0)

    def aacircle(surface, x, y, r, color):
        pygame.draw.circle(surface, color, (x, y), r, width=1)

    def filled_circle(surface, x, y, r, color):
        pygame.draw.circle(surface, color, (x, y), r, width=0)


    gfxdraw.aapolygon = aapolygon
    gfxdraw.filled_polygon = filled_polygon
    gfxdraw.aacircle = aacircle
    gfxdraw.filled_circle = filled_circle


    # from irlc.utils.player_wrapper import AsyncPlayerWrapperPygame
    agent = AsyncPlayerWrapperPygame(agent, env, autoplay=autoplay)

    s, info = env.reset()
    k = 0
    await asyncio.sleep(0.2)  # set up graphics.
    running = True
    step = 0
    while running:
        pi_action = agent.pi(s, k, info)
        a = None
        while True:
            for event in pygame.event.get():
                # if event.type == pygame.KEYDOWN and event.key == pygame.K_r:
                #     print("Resetting the environment")
                #     a = HUMAN_DEMAND_RESET
                #     break

                a = agent._resolve_event_into_action(event=event, pi_action=pi_action, info=info)
                if a is not None:
                    break
            if agent.human_demand_autoplay:
                a = pi_action
                break
            if a is not None:
                break
            await asyncio.sleep(0.05)

        # if not isinstance(a, np.ndarray) and isinstance(a, str) and a == HUMAN_DEMAND_RESET:
        #     s, info = env.reset()
        #     continue

        sp, reward, done, truncated, info_sp = (await env.async_step(a)) if hasattr(env, 'async_step') else env.step(a)

        agent.train(s, a, reward, sp, done=done, info_s=info, info_sp=info_sp)

        step = step + 1
        k = k + 1
        if done or truncated:
            sp, info_sp = env.reset()
            k = 0

        s = sp
        info = info_sp
        await asyncio.sleep(0.01)  # release loop.


class PlayWrapperPygame(AgentWrapper):
    render_after_train = True
    ACTION_FORCE_RESET = "force reset of environment."

    def __init__(self, agent : Agent, env : gym.Env, keys_to_action=None, autoplay=False):
        super().__init__(agent, env)
        if keys_to_action is None:
            if hasattr(env, 'get_keys_to_action'):
                keys_to_action = env.get_keys_to_action()
            elif hasattr(env, "env") and hasattr(env.env, 'get_keys_to_action'):
                keys_to_action = env.env.get_keys_to_action()
            elif hasattr(env.unwrapped, 'get_keys_to_action'):
                keys_to_action = env.unwrapped.get_keys_to_action()
            else:
                keys_to_action = dict()
                # print(env.spec.id +" does not have explicit key to action mapping, please specify one manually")
                # assert False, env.spec.id + " does not have explicit key to action mapping, " + \
                #               "please specify one manually"
                # keys_to_action = dict()
        self.keys_to_action = keys_to_action
        self.env = env
        self.human_wants_restart = False
        self.human_sets_pause = False
        self.human_agent_action = -1
        self.human_demand_autoplay = autoplay

        # Now for the more fucky stuff. Collect internal statistics and such.
        self.env._interactive_data = dict(trajectories=[], in_terminal_state=False)
        # Now fix the train function
        train2 = agent.train
        reset2 = env.reset

        def reset_(**kwargs):
            from irlc.ex01.agent import Trajectory
            s, info = reset2(**kwargs)
            env._interactive_data['trajectories'].append(Trajectory(state=[s], time=[0], action=[], reward=[], env_info=[info]))
            env._interactive_data['in_terminal_state'] = False
            return s, info

        def train_(s, a, r, sp, done, info_s, info_sp):
            if not isinstance(a, str) or a != PlayWrapperPygame.ACTION_FORCE_RESET:
                train2(s, a, r, sp, done, info_s, info_sp)
                env._interactive_data['trajectories'][-1].state.append(sp)
                env._interactive_data['trajectories'][-1].reward.append(r)
                env._interactive_data['trajectories'][-1].action.append(a)
                env._interactive_data['trajectories'][-1].env_info.append(info_sp)

                if done:
                    env._interactive_data['in_terminal_state'] = True
                    env._interactive_data['avg_reward_per_episode'] = np.mean( [sum(t.reward) for t in env._interactive_data['trajectories']] )
                    env._interactive_data['completed_episodes'] = len(env._interactive_data['trajectories'] )

            if self.render_after_train: # always true.
                env.render()


        step2 = env.step
        def step_(a):
            if isinstance(a, str) and a == PlayWrapperPygame.ACTION_FORCE_RESET:
                print("Resetting interaction data.")
                self.env._interactive_data = dict(trajectories=[], in_terminal_state=False)
                return "done", 3, True, True, {}
            else:
                sp, r, terminated, truncated, info = step2(a)
            return sp, r, terminated, truncated, info


        agent.train = train_
        env.agent = agent
        env.unwrapped.agent = agent

        env.reset = reset_
        env.step = step_


    def _resolve_event_into_action(self, event : pygame.event.Event, pi_action, info : dict):
        # assert False # Unclear if used.
        if event.type == pygame.QUIT:
            if hasattr(self, 'env'):
                self.env.close()
            time.sleep(0.1)
            pygame.display.quit()
            time.sleep(0.1)
            pygame.quit()
            time.sleep(0.1)
            sys.exit()

        # checking if keydown event happened or not
        if event.type == pygame.KEYDOWN:
            # if event.key == pygame.K_f:
            #     print("Pressing f")
            #     self.env.render()
            #     return None
            if event.key == pygame.K_SPACE:
                a = pi_action
                return a
            elif (event.key,) in self.keys_to_action:
                a = self.keys_to_action[(event.key,)]
                if info is not None and 'mask' in info:
                    from irlc.utils.common import DiscreteTextActionSpace
                    if isinstance(self.env.action_space, DiscreteTextActionSpace):
                        aint = self.env.action_space.actions.index(a)
                    else:
                        aint = a

                    if info['mask'][aint] == 0:
                        # The action was masked. This means that this action is unavailable, and we should select another.
                        # The default is to select one of the available actions from the mask.
                        from irlc.pacman.gamestate import GameState
                        from irlc.pacman.pacman_environment import PacmanEnvironment
                        if isinstance(self.env, PacmanEnvironment):
                            a = "Stop"
                        else:
                            a = info['mask'].argmax()
                            if isinstance(self.env.action_space, DiscreteTextActionSpace):
                                a = self.env.action_space.actions[a]
                    return a
                else:
                    return a
            elif event.key == pygame.K_r:
                print("Pressing r")
                if hasattr(self, 'reset'):
                    self.reset()
                return PlayWrapperPygame.ACTION_FORCE_RESET

            elif event.key == pygame.K_f:
                print("Pressing f")
                self.env.render()

            elif event.key == pygame.K_p:
                # unpause
                # print("Unpausing game")
                self.human_demand_autoplay = not self.human_demand_autoplay
                # print(f"Unpausing game {self.human_demand_autoplay=}")

                if self.human_demand_autoplay:
                    # print("Returning", pi_action)
                    a = pi_action
                    return a
            else:
                # try to pass event on to the game.
                if hasattr(self.env, 'keypress'):
                    self.env.keypress(event)
        return None

    def pi(self, state, k, info=None):
        # print("Entering pi")
        pi_action = super().pi(state, k, info) # make sure super class pi method is called in case it has side effects.
        a = None
        while True:
            # print(f"while loop {self.human_demand_autoplay=}")
            for event in pygame.event.get():
                # print("resolving", event)
                a = self._resolve_event_into_action(event, pi_action, info)
                # print("Resolved", a)
                if a is not None:
                    # print("Breaking", a)
                    break
                # if False:
                #     if event.type == pygame.QUIT:
                #         if hasattr(self, 'env'):
                #             self.env.close()
                #         time.sleep(0.1)
                #         pygame.display.quit()
                #         time.sleep(0.1)
                #         pygame.quit()
                #         time.sleep(0.1)
                #         sys.exit()
                #
                #     # checking if keydown event happened or not
                #     if event.type == pygame.KEYDOWN:
                #         if event.key == pygame.K_SPACE:
                #             a = pi_action
                #             break
                #         elif (event.key,) in self.keys_to_action:
                #             a = self.keys_to_action[(event.key,)]
                #             if info is not None and 'mask' in info:
                #                 from irlc.utils.common import DiscreteTextActionSpace
                #                 if isinstance(self.env.action_space, DiscreteTextActionSpace):
                #                     aint = self.env.action_space.actions.index(a)
                #                 else:
                #                     aint = a
                #
                #                 if info['mask'][aint] == 0:
                #                     # The action was masked. This means that this action is unavailable, and we should select another.
                #                     # The default is to select one of the available actions from the mask.
                #                     a = info['mask'].argmax()
                #                     if isinstance(self.env.action_space, DiscreteTextActionSpace):
                #                         a = self.env.action_space.actions[a]
                #                 break
                #             else:
                #                 break
                #         elif event.key == pygame.K_r:
                #             print("Pressing r")
                #             if hasattr(self, 'reset'):
                #                 return PlayWrapperPygame.ACTION_FORCE_RESET
                #             #
                #             #     self.reset()
                #             #     self.env.reset()
                #             #
                #             # self.env.render()
                #         elif event.key == pygame.K_f:
                #             print("Pressing f")
                #             self.env.render()
                #
                #         elif event.unicode == 'p':
                #             # unpause
                #             self.human_demand_autoplay = not self.human_demand_autoplay
                #             break
                #         else:
                #             # try to pass event on to the game.
                #             if hasattr(self.env, 'keypress'):
                #                 self.env.keypress(event)

            if self.human_demand_autoplay:
                a = pi_action

            if a is not None:
                try:
                    from irlc.pacman.gamestate import GameState
                    if isinstance(state, GameState):
                        if a not in state.A():
                            a = "Stop"
                except Exception as e:
                    pass
                # print(f"{a=}")
                return a
            time.sleep(0.1)




class AsyncPlayerWrapperPygame(PlayWrapperPygame):
    # render_after_train = False

    def pi(self, state, k, info=None):
        pi_action = self.agent.pi(state, k, info)  # make sure super class pi method is called in case it has side effects.
        return pi_action

    def _post_process_action_for_visualization(self):
        pass



[docs]
def interactive(env : gym.Env, agent: Agent, autoplay=False) -> (gym.Env, Agent):
    """
    This function is used for visualizations. It can

    - Allow you to input keyboard commands to an environment
    - Allow you to save results
    - Visualize reinforcement-learning agents in the gridworld environment.

    by adding a single extra line ``env, agent = interactive(env,agent)``.
    The following shows an example:

        >>> from irlc.gridworld.gridworld_environments import BookGridEnvironment
        >>> from irlc import train, Agent, interactive
        >>> env = BookGridEnvironment(render_mode="human", zoom=0.8) # Pass render_mode='human' for visualization.
        >>> env, agent = interactive(env, Agent(env))               # Make the environment interactive. Note that it needs an agent.
        >>> train(env, agent, num_episodes=2)                     # You can train and use the agent and environment as usual.
        >>> env.close()

    It also enables you to visualize the environment at a matplotlib figure or save it as a pdf file using ``env.plot()`` and ``env.savepdf('my_file.pdf)``.

    All demos and figures in the notes are made using this function.

    :param env: A gym environment (an instance of the ``Env`` class)
    :param agent: An agent (an instance of the ``Agent`` class)
    :param autoplay: Whether the simulation should be unpaused automatically
    :return: An environment and agent which have been slightly updated to make them interact with each other. You can use them as usual with the ``train``-function.
    """
    from PIL import Image # Let's put this one here in case we run the code in headless mode.

    agent = PlayWrapperPygame(agent, env, autoplay=autoplay)

    pass

    def plot():
        env.render_mode, rmt = 'rgb_array', env.render_mode
        frame = env.render()
        env.render_mode = rmt
        im = Image.fromarray(frame)
        plt.imshow(im)
        plt.axis('off')
        plt.axis('off')
        plt.tight_layout()

    def savepdf(file):
        env.render_mode, rmt = 'rgb_array', env.render_mode
        frame = env.render()
        env.render_mode = rmt

        im = Image.fromarray(frame)
        snapshot_base = file
        if snapshot_base.endswith(".png"):
            sf = snapshot_base[:-4]
            fext = 'png'
        else:
            fext = 'pdf'
            if snapshot_base.endswith(".pdf"):
                sf = snapshot_base[:-4]
            else:
                sf = snapshot_base

        sf = f"{sf}.{fext}"
        dn = os.path.dirname(sf)
        if len(dn) > 0 and not os.path.isdir(dn):
            os.makedirs(dn)
        print("Saving snapshot of environment to", os.path.abspath(sf))
        if fext == 'png':
            im.save(sf)
            from irlc import _move_to_output_directory
            _move_to_output_directory(sf)
        else:
            plt.figure(figsize=(16, 16))
            plt.imshow(im)
            plt.axis('off')
            plt.tight_layout()
            from irlc import savepdf
            savepdf(sf, verbose=True)
            plt.show()
    env.plot = plot
    env.savepdf = savepdf
    return env, agent



def main():
    from irlc.ex11.q_agent import QAgent
    from irlc.gridworld.gridworld_environments import BookGridEnvironment  
    from irlc import train, Agent
    env = BookGridEnvironment(render_mode="human", zoom=0.8)  # Pass render_mode='human' for visualization.
    env, agent = interactive(env, Agent(env))  # Make th
    env.reset()     # We always need to call reset
    env.plot()      # Plot the environment.
    env.close()  

    # Interaction with a random agent.
    from irlc.gridworld.gridworld_environments import BookGridEnvironment 
    from irlc import train, Agent
    env = BookGridEnvironment(render_mode="human", zoom=0.8) # Pass render_mode='human' for visualization.
    env, agent = interactive(env, Agent(env))               # Make the environment interactive. Note that it needs an agent.
    train(env, agent, num_episodes=100)                      # You can train and use the agent and environment as usual. 
    env.close()


if __name__ == "__main__":
    main()