#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.

import logging
from typing import Any, Callable, List, Tuple

import gym
import numpy as np
from reagent.gym.types import Transition
from reagent.replay_memory.circular_replay_buffer import ReplayBuffer


logger = logging.getLogger(__name__)


try:
    from recsim.simulator.recsim_gym import RecSimGymEnv

    HAS_RECSIM = True
except ImportError:
    HAS_RECSIM = False
    logger.warning(f"ReplayBuffer.create_from_env() will not recognize RecSim env")


# Arguments: replay_buffer, obs, action, reward, terminal, log_prob
ReplayBufferInserter = Callable[[ReplayBuffer, Transition], None]


def make_replay_buffer_inserter(env: gym.Env) -> ReplayBufferInserter:
    if HAS_RECSIM and isinstance(env.unwrapped, RecSimGymEnv):
        return RecSimReplayBufferInserter.create_for_env(env)
    return BasicReplayBufferInserter()


class BasicReplayBufferInserter:
    def __call__(self, replay_buffer: ReplayBuffer, transition: Transition):
        replay_buffer.add(**transition.asdict())


class RecSimReplayBufferInserter:
    def __init__(
        self,
        *,
        num_docs: int,
        num_responses: int,
        discrete_keys: List[str],
        box_keys: List[str],
        response_discrete_keys: List[Tuple[str, int]],
        response_box_keys: List[Tuple[str, Tuple[int]]],
        augmentation_discrete_keys: List[str],
        augmentation_box_keys: List[str],
    ):
        self.num_docs = num_docs
        self.num_responses = num_responses
        self.discrete_keys = discrete_keys
        self.box_keys = box_keys
        self.response_discrete_keys = response_discrete_keys
        self.response_box_keys = response_box_keys
        self.augmentation_discrete_keys = augmentation_discrete_keys
        self.augmentation_box_keys = augmentation_box_keys

    @classmethod
    def create_for_env(cls, env: gym.Env):
        obs_space = env.observation_space
        assert isinstance(obs_space, gym.spaces.Dict)
        user_obs_space = obs_space["user"]
        if not isinstance(user_obs_space, gym.spaces.Box):
            raise NotImplementedError(
                f"User observation space {type(user_obs_space)} is not supported"
            )

        doc_obs_space = obs_space["doc"]
        if not isinstance(doc_obs_space, gym.spaces.Dict):
            raise NotImplementedError(
                f"Doc space {type(doc_obs_space)} is not supported"
            )

        # Assume that all docs are in the same space

        discrete_keys: List[str] = []
        box_keys: List[str] = []

        key_0 = next(iter(doc_obs_space.spaces))
        doc_0_space = doc_obs_space[key_0]

        if isinstance(doc_0_space, gym.spaces.Dict):
            for k, v in doc_0_space.spaces.items():
                if isinstance(v, gym.spaces.Discrete):
                    if v.n > 0:
                        discrete_keys.append(k)
                elif isinstance(v, gym.spaces.Box):
                    shape_dim = len(v.shape)
                    if shape_dim <= 1:
                        box_keys.append(k)
                    else:
                        raise NotImplementedError
                else:
                    raise NotImplementedError(
                        f"Doc feature {k} with the observation space of {type(v)}"
                        " is not supported"
                    )
        elif isinstance(doc_0_space, gym.spaces.Box):
            pass
        else:
            raise NotImplementedError(f"Unknown space {doc_0_space}")

        augmentation_discrete_keys: List[str] = []
        augmentation_box_keys: List[str] = []
        augmentation = obs_space.spaces.get("augmentation", None)
        if augmentation is not None:
            aug_0_space = list(augmentation.spaces.values())[0]
            for k, v in aug_0_space.spaces.items():
                if isinstance(v, gym.spaces.Discrete):
                    if v.n > 0:
                        augmentation_discrete_keys.append(k)
                elif isinstance(v, gym.spaces.Box):
                    shape_dim = len(v.shape)
                    if shape_dim <= 1:
                        augmentation_box_keys.append(k)
                    else:
                        raise NotImplementedError
                else:
                    raise NotImplementedError(
                        f"Augmentation {k} with the observation space "
                        f" of {type(v)} is not supported"
                    )

        response_space = obs_space["response"][0]
        assert isinstance(response_space, gym.spaces.Dict)
        response_box_keys: List[Tuple[str, Tuple[int]]] = []
        response_discrete_keys: List[Tuple[str, int]] = []
        for k, v in response_space.spaces.items():
            if isinstance(v, gym.spaces.Discrete):
                response_discrete_keys.append((k, v.n))
            elif isinstance(v, gym.spaces.Box):
                response_box_keys.append((k, v.shape))
            else:
                raise NotImplementedError

        return cls(
            num_docs=len(doc_obs_space.spaces),
            num_responses=len(obs_space["response"]),
            discrete_keys=discrete_keys,
            box_keys=box_keys,
            response_box_keys=response_box_keys,
            response_discrete_keys=response_discrete_keys,
            augmentation_box_keys=augmentation_box_keys,
            augmentation_discrete_keys=augmentation_discrete_keys,
        )

    def __call__(self, replay_buffer: ReplayBuffer, transition: Transition):
        transition_dict = transition.asdict()
        obs = transition_dict.pop("observation")
        user = obs["user"]

        kwargs = {}

        if self.box_keys or self.discrete_keys:
            doc_obs = obs["doc"]
            for k in self.box_keys:
                kwargs[f"doc_{k}"] = np.stack([v[k] for v in doc_obs.values()])
            for k in self.discrete_keys:
                kwargs[f"doc_{k}"] = np.array([v[k] for v in doc_obs.values()])
        else:
            kwargs["doc"] = np.stack(list(obs["doc"].values()))

        # Augmentation

        if self.augmentation_box_keys or self.augmentation_discrete_keys:
            aug_obs = obs["augmentation"]
            for k in self.augmentation_box_keys:
                kwargs[f"augmentation_{k}"] = np.stack([v[k] for v in aug_obs.values()])
            for k in self.augmentation_discrete_keys:
                kwargs[f"augmentation_{k}"] = np.array([v[k] for v in aug_obs.values()])

        # Responses

        response = obs["response"]
        # We need to handle None below because the first state won't have response
        for k, d in self.response_box_keys:
            if response is not None:
                kwargs[f"response_{k}"] = np.stack([v[k] for v in response])
            else:
                kwargs[f"response_{k}"] = np.zeros(
                    (self.num_responses, *d), dtype=np.float32
                )
        for k, _n in self.response_discrete_keys:
            if response is not None:
                kwargs[f"response_{k}"] = np.array([v[k] for v in response])
            else:
                kwargs[f"response_{k}"] = np.zeros(
                    (self.num_responses,), dtype=np.int64
                )

        transition_dict.update(kwargs)
        replay_buffer.add(observation=user, **transition_dict)
