"""Behavioural Cloning (BC).

Trains policy by applying supervised learning to a fixed dataset of (observation,
action) pairs generated by some expert demonstrator.
"""

import collections
import os
from typing import Any, Callable, List, Mapping, Optional, Type, Union

import cloudpickle
import gym
import numpy as np
import tensorflow as tf
from stable_baselines import logger
from stable_baselines.common.policies import ActorCriticPolicy, BasePolicy
from tqdm.autonotebook import trange

from imitation.data import datasets, types
from imitation.policies.base import FeedForward32Policy


def set_tf_vars(
    *,
    values: List[np.ndarray],
    scope: Optional[str] = None,
    tf_vars: Optional[List[tf.Variable]] = None,
    sess: Optional[tf.Session] = None,
):
    """Set a collection of variables to take the values in `values`.

    Variables can be either specified by scope or passed directly into the
    function as a list. Variables and values will be matched based on the order
    in which they appear in their respective collections, so there must be as
    many values as variables.

    Args:
        values: list of values to load into variables.
        scope: scope to collect variables from. Either this argument xor
          `tf_vars` must be given.
        tf_vars: explicit list of TF variables to write to. Mutex with `scope`.
        sess: TF session to use, if not the default.
    """
    if scope is not None:
        assert tf_vars is None, "must give either `tf_vars` xor `scope` kwargs"
        tf_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope=scope)
    else:
        assert tf_vars is not None, "must give either `tf_vars` xor `scope` kwargs"
    assert len(tf_vars) == len(values), (
        f"{len(tf_vars)} tf variables but " f"{len(values)} values supplied"
    )
    sess = sess or tf.get_default_session()
    assert sess is not None, "must supply session or have one in context"
    placeholders = [tf.placeholder(shape=v.shape, dtype=v.dtype) for v in tf_vars]
    assign_ops = [tf.assign(var, ph) for var, ph in zip(tf_vars, placeholders)]
    sess.run(
        assign_ops, feed_dict={ph: value for ph, value in zip(placeholders, values)}
    )


class BC:
    def __init__(
        self,
        observation_space: gym.Space,
        action_space: gym.Space,
        *,
        policy_class: Type[ActorCriticPolicy] = FeedForward32Policy,
        policy_kwargs: Optional[Mapping[str, Any]] = None,
        expert_data: Union[
            types.TransitionsMinimal,
            datasets.Dataset[types.TransitionsMinimal],
            None,
        ] = None,
        batch_size: int = 32,
        optimizer_cls: Type[tf.train.Optimizer] = tf.train.AdamOptimizer,
        optimizer_kwargs: Optional[dict] = None,
        ent_weight: float = 1e-3,
        l2_weight: float = 0.0,
    ):
        """Behavioral cloning (BC).

        Recovers a policy via supervised learning on a Dataset of observation-action
        pairs.

        Args:
            env: environment to train on.
            policy_class: used to instantiate imitation policy.
            expert_data: If not None, then immediately call
                  `self.set_expert_dataset(expert_data)` during initialization.
            batch_size: batch size used for training.
            optimizer_cls: optimiser to use for supervised training.
            optimizer_kwargs: keyword arguments for optimiser construction.
        """
        self.action_space = action_space
        self.observation_space = observation_space
        self.policy_class = policy_class
        policy_kwargs_defaults = dict(
            ob_space=self.observation_space,
            ac_space=self.action_space,
            n_batch=None,
            n_env=1,
            n_steps=1000,
        )
        self.policy_kwargs = collections.ChainMap(
            policy_kwargs or {}, policy_kwargs_defaults
        )

        assert batch_size >= 1
        self.batch_size = batch_size
        self.expert_dataset: Optional[datasets.Dataset[types.TransitionsMinimal]] = None
        self.sess = tf.get_default_session()
        assert self.sess is not None, "need to construct this within a session scope"
        self.ent_weight = ent_weight
        self.l2_weight = l2_weight

        self.optimizer_cls = optimizer_cls
        self.optimizer_kwargs = optimizer_kwargs or {}

        self._build_tf_graph()
        self.sess.run(tf.global_variables_initializer())
        if expert_data is not None:
            self.set_expert_dataset(expert_data)

    def set_expert_dataset(
        self,
        expert_data: Union[
            types.TransitionsMinimal,
            datasets.Dataset[types.TransitionsMinimal],
        ],
    ):
        """Replace the current expert dataset with a new one.

        Useful for DAgger and other interactive algorithms.

        Args:
             expert_data: Either a `Dataset[types.TransitionsMinimal]` for which
                 `.size()` is not None, or a instance of `TransitionsMinimal`, which
                 is automatically converted to a shuffled, epoch-order
                 `Dataset[types.TransitionsMinimal]`.
        """
        if isinstance(expert_data, types.TransitionsMinimal):
            trans = expert_data
            expert_dataset = datasets.TransitionsDictDatasetAdaptor(
                trans, datasets.EpochOrderDictDataset
            )
        else:
            assert isinstance(expert_data, datasets.Dataset)
            expert_dataset = expert_data
        assert expert_dataset.size() is not None
        self.expert_dataset = expert_dataset

    def _build_tf_graph(self):
        with tf.variable_scope("bc_supervised_loss"):
            with tf.variable_scope("model") as model_scope:
                self.policy = self.policy_class(
                    sess=self.sess, **self.policy_kwargs
                )  # pytype: disable=not-instantiable
                inner_scope = tf.get_variable_scope().name
                self.policy_variables = tf.get_collection(
                    tf.GraphKeys.TRAINABLE_VARIABLES, scope=inner_scope
                )
            self._true_acts_ph = self.policy.pdtype.sample_placeholder(
                [None], name="ref_acts_ph"
            )
            self._entropy = tf.reduce_mean(self.policy.proba_distribution.entropy())
            self._ent_loss = -self.ent_weight * self._entropy
            self._neglogp = tf.reduce_mean(
                self.policy.proba_distribution.neglogp(self._true_acts_ph)
            )

            _l2_norms = [tf.nn.l2_loss(w) for w in model_scope.trainable_variables()]
            self._l2_norm = tf.reduce_sum(_l2_norms)
            self._l2_loss = self._l2_norm * self.l2_weight

            self._loss = self._neglogp + self._ent_loss + self._l2_loss
            opt = self.optimizer_cls(**self.optimizer_kwargs)
            self._train_op = opt.minimize(self._loss)

            self._prob_true_act = tf.reduce_mean(
                tf.exp(-self.policy.proba_distribution.neglogp(self._true_acts_ph))
            )
            self._stats_dict = dict(
                neglopp=self._neglogp,
                loss=self._loss,
                entropy=self._entropy,
                ent_loss=self._ent_loss,
                prob_true_act=self._prob_true_act,
                l2_norm=self._l2_norm,
                l2_loss=self._l2_loss,
            )

    def train(
        self,
        n_epochs: int = 100,
        *,
        on_epoch_end: Callable[[dict], None] = None,
        log_interval: int = 100,
    ):
        """Train with supervised learning for some number of epochs.

        Here an 'epoch' is just a complete pass through the expert transition
        dataset.

        Args:
          n_epochs: number of complete passes made through dataset.
          on_epoch_end: optional callback to run at
            the end of each epoch. Will receive all locals from this function as
            dictionary argument (!!).
        """
        assert self.batch_size >= 1
        samples_so_far = 0
        batch_num = 0
        for epoch_num in trange(n_epochs, desc="BC epoch"):
            while samples_so_far < (epoch_num + 1) * self.expert_dataset.size():
                batch_num += 1
                trans = self.expert_dataset.sample(self.batch_size)
                assert len(trans) == self.batch_size
                samples_so_far += self.batch_size
                feed_dict = {
                    self._true_acts_ph: trans.acts,
                    self.policy.obs_ph: trans.obs,
                }
                _, stats_dict = self.sess.run(
                    [self._train_op, self._stats_dict], feed_dict=feed_dict
                )
                stats_dict["epoch_num"] = epoch_num
                stats_dict["n_updates"] = batch_num
                stats_dict["batch_size"] = self.batch_size

                if batch_num % log_interval == 0:
                    for k, v in stats_dict.items():
                        logger.logkv(k, v)
                    logger.dumpkvs()
                batch_num += 1

            if on_epoch_end is not None:
                on_epoch_end(locals())

    def save_policy(self, policy_path: str):
        """Save a policy to a pickle. Can be reloaded by `.reconstruct_policy()`.

        Args:
            policy_path: path to save policy to.
        """
        policy_params = self.sess.run(self.policy_variables)
        data = {
            "class": self.policy_class,
            "kwargs": self.policy_kwargs,
            "params": policy_params,
        }
        dirname = os.path.dirname(policy_path)
        if dirname:
            os.makedirs(dirname, exist_ok=True)
        with open(policy_path, "wb") as fp:
            cloudpickle.dump(data, fp)

    @staticmethod
    def reconstruct_policy(
        policy_path: str,
        sess: Optional[tf.Session] = None,
    ) -> BasePolicy:
        """Reconstruct a saved policy.

        Args:
            policy_path: path a policy produced by `.save_policy()`.
            sess: optional session to construct policy under,
              if not the default session.

        Returns:
            policy: policy with reloaded weights.
        """
        if sess is None:
            sess = tf.get_default_session()
            assert sess is not None, "must supply session via kwarg or context mgr"

        # re-read data from dict
        with open(policy_path, "rb") as fp:
            loaded_pickle = cloudpickle.load(fp)

        # construct the policy class
        klass = loaded_pickle["class"]
        kwargs = loaded_pickle["kwargs"]
        with tf.variable_scope("reconstructed_policy"):
            rv_pol = klass(sess=sess, **kwargs)
            inner_scope = tf.get_variable_scope().name

        # set values for the new policy's parameters
        param_values = loaded_pickle["params"]
        set_tf_vars(values=param_values, scope=inner_scope, sess=sess)

        return rv_pol
