# This file is part of NPFL138 <http://github.com/ufal/npfl138/>.
#
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
import os
import sys

import gymnasium as gym
import numpy as np
import torch


############################
# Gym Environment Wrappers #
############################

class EvaluationEnv(gym.Wrapper):
    """A wrapper over gym environments capable of performing evaluation on demant.

    In addition to a standard gym envrironment, it provides the following features:

    - The `reset` method can be called as `reset(start_evaluation=False)` to start the evaluation.
    - The `env.episode` property returns the number of completed episodes.
    """
    def __init__(self, env, seed=None, render_each=0, evaluate_for=100, report_each=10):
        super().__init__(env)
        self._render_each = render_each
        self._evaluate_for = evaluate_for
        self._report_each = report_each
        self._report_verbose = os.getenv("VERBOSE", "1") not in ["", "0"]

        gym.Env.reset(self.unwrapped, seed=seed)
        self.action_space.seed(seed)
        self.observation_space.seed(seed)

        self._episode_running = False
        self._episode_returns = []
        self._evaluating_from = None
        self._original_render_mode = env.render_mode
        self._pygame = __import__("pygame") if self._render_each else None

    @property
    def episode(self):
        return len(self._episode_returns)

    def reset(self, *, start_evaluation=False, logging=True, seed=None, options=None):
        start_evaluation = start_evaluation or (options or {}).get("start_evaluation", False)
        logging = logging and (options or {}).get("logging", True)

        if seed is not None:
            raise RuntimeError("The EvaluationEnv cannot be reseeded")
        if self._evaluating_from is not None and self._episode_running:
            raise RuntimeError("Cannot reset a running episode after `start_evaluation=True`")
        if start_evaluation and self._evaluating_from is None:
            self._evaluating_from = self.episode

        if logging and self._render_each and (self.episode + 1) % self._render_each == 0:
            self.unwrapped.render_mode = "human"
        elif self._render_each:
            self.unwrapped.render_mode = self._original_render_mode
        self._episode_running = True
        self._episode_return = 0 if logging or self._evaluating_from is not None else None
        return super().reset(options=options)

    def step(self, action):
        if not self._episode_running:
            raise RuntimeError("Cannot run `step` on environments without an active episode, run `reset` first")

        observation, reward, terminated, truncated, info = super().step(action)
        done = terminated or truncated

        self._episode_running = not done
        if self._episode_return is not None:
            self._episode_return += reward
        if self._episode_return is not None and done:
            self._episode_returns.append(self._episode_return)

            if self._report_each and self.episode % self._report_each == 0:
                print("Episode {}, mean {}-episode return {:.2f} +-{:.2f}{}".format(
                    self.episode, self._evaluate_for, np.mean(self._episode_returns[-self._evaluate_for:]),
                    np.std(self._episode_returns[-self._evaluate_for:]), "" if not self._report_verbose else
                    ", returns " + " ".join(map("{:g}".format, self._episode_returns[-self._report_each:]))),
                    file=sys.stderr, flush=True)
            if self._evaluating_from is not None and self.episode >= self._evaluating_from + self._evaluate_for:
                print("The mean {}-episode return after evaluation {:.2f} +-{:.2f}".format(
                    self._evaluate_for, np.mean(self._episode_returns[-self._evaluate_for:]),
                    np.std(self._episode_returns[-self._evaluate_for:])), flush=True)
                self.close()
                sys.exit(0)

        if self._pygame and self.unwrapped.render_mode == "human" and self._pygame.get_init():
            if self._pygame.event.get(self._pygame.QUIT):
                self.unwrapped.render_mode = self._original_render_mode

        return observation, reward, terminated, truncated, info


#############
# Utilities #
#############


def typed_torch_function(device, *types, via_np=True):
    """Typed Torch function decorator.

    The positional input arguments are converted to torch Tensors of the given
    types and on the given device; for NumPy arrays on the same device,
    the conversion should not copy the data.

    The torch Tensors generated by the wrapped function are converted back
    to Numpy arrays before returning (while keeping original tuples, lists,
    and dictionaries).
    """
    def check_typed_torch_function(wrapped, args):
        if len(types) != len(args):
            while hasattr(wrapped, "__wrapped__"):
                wrapped = wrapped.__wrapped__
            raise AssertionError("The typed_torch_function decorator for {} expected {} arguments, but got {}".format(
                wrapped, len(types), len(args)))

    def structural_map(value):
        if isinstance(value, torch.Tensor):
            return value.numpy(force=True)
        if isinstance(value, tuple):
            return tuple(structural_map(element) for element in value)
        if isinstance(value, list):
            return [structural_map(element) for element in value]
        if isinstance(value, dict):
            return {key: structural_map(element) for key, element in value.items()}
        return value

    class TypedTorchFunctionWrapper:
        def __init__(self, func):
            self.__wrapped__ = func

        def __call__(self, *args, **kwargs):
            check_typed_torch_function(self.__wrapped__, args)
            return structural_map(self.__wrapped__(
                *[torch.as_tensor(np.asarray(arg) if via_np else arg, dtype=typ, device=device)
                  for arg, typ in zip(args, types)], **kwargs))

        def __get__(self, instance, cls):
            return TypedTorchFunctionWrapper(self.__wrapped__.__get__(instance, cls))

    return TypedTorchFunctionWrapper
