import collections
import functools
import multiprocessing
import os
from typing import Set

import numpy as np
import tensorflow as tf


def is_image(tensor):
    return len(tensor.shape) == 3 and tensor.shape[-1] in [1, 3, 4]


def batch_to_seq(tensor_batch, n_batch, n_steps, flat=False):
    if flat:
        tensor_batch = tf.reshape(tensor_batch, [n_batch, n_steps])
    else:
        tensor_batch = tf.reshape(tensor_batch, [n_batch, n_steps, -1])
    return [
        tf.squeeze(v, [1])
        for v in tf.split(axis=1, num_or_size_splits=n_steps, value=tensor_batch)
    ]


def seq_to_batch(tensor_sequence, flat=False):
    shape = tensor_sequence[0].get_shape().as_list()
    if not flat:
        assert len(shape) > 1
        n_hidden = tensor_sequence[0].get_shape()[-1].value
        return tf.reshape(tf.concat(axis=1, values=tensor_sequence), [-1, n_hidden])
    else:
        return tf.reshape(tf.stack(values=tensor_sequence, axis=1), [-1])


def check_shape(tensors, shapes):
    i = 0
    for (tensor, shape) in zip(tensors, shapes):
        assert tensor.get_shape().as_list() == shape, (
            "id " + str(i) + " shape " + str(tensor.get_shape()) + str(shape)
        )
        i += 1


def huber_loss(tensor, delta=1.0):
    return tf.compat.v1.where(
        tf.abs(tensor) < delta,
        tf.square(tensor) * 0.5,
        delta * (tf.abs(tensor) - 0.5 * delta),
    )


def sample(logits):
    noise = tf.random.uniform(tf.shape(input=logits))
    return tf.argmax(input=logits - tf.math.log(-tf.math.log(noise)), axis=1)


def calc_entropy(logits):
    a_0 = logits - tf.reduce_max(input_tensor=logits, axis=1, keepdims=True)
    exp_a_0 = tf.exp(a_0)
    z_0 = tf.reduce_sum(input_tensor=exp_a_0, axis=1, keepdims=True)
    p_0 = exp_a_0 / z_0
    return tf.reduce_sum(input_tensor=p_0 * (tf.math.log(z_0) - a_0), axis=1)


def mse(pred, target):
    return tf.reduce_mean(input_tensor=tf.square(pred - target))


def avg_norm(tensor):
    return tf.reduce_mean(
        input_tensor=tf.sqrt(tf.reduce_sum(input_tensor=tf.square(tensor), axis=-1))
    )


def gradient_add(grad_1, grad_2, param, verbose=0):
    if verbose > 1:
        print([grad_1, grad_2, param.name])
    if grad_1 is None and grad_2 is None:
        return None
    elif grad_1 is None:
        return grad_2
    elif grad_2 is None:
        return grad_1
    else:
        return grad_1 + grad_2


def q_explained_variance(q_pred, q_true):
    _, var_y = tf.nn.moments(x=q_true, axes=[0, 1])
    _, var_pred = tf.nn.moments(x=q_true - q_pred, axes=[0, 1])
    check_shape([var_y, var_pred], [[]] * 2)
    return 1.0 - (var_pred / var_y)


def make_session(num_cpu=None, make_default=False, graph=None):
    if num_cpu is None:
        num_cpu = int(os.getenv('RCALL_NUM_CPU', multiprocessing.cpu_count()))
    tf_config = tf.compat.v1.ConfigProto(
        allow_soft_placement=True,
        inter_op_parallelism_threads=num_cpu,
        intra_op_parallelism_threads=num_cpu,
    )
    tf_config.gpu_options.allow_growth = True
    if make_default:
        return tf.compat.v1.InteractiveSession(config=tf_config, graph=graph)
    else:
        return tf.compat.v1.Session(config=tf_config, graph=graph)


def single_threaded_session(make_default=False, graph=None):
    return make_session(num_cpu=1, make_default=make_default, graph=graph)


def in_session(func):
    @functools.wraps(func)
    def newfunc(*args, **kwargs):
        with tf.compat.v1.Session():
            func(*args, **kwargs)

    return newfunc


ALREADY_INITIALIZED = set()


def initialize(sess=None):
    if sess is None:
        sess = tf.compat.v1.get_default_session()
    new_variables = set(tf.compat.v1.global_variables()) - ALREADY_INITIALIZED
    sess.run(tf.compat.v1.variables_initializer(new_variables))
    ALREADY_INITIALIZED.update(new_variables)


def function(inputs, outputs, updates=None, givens=None):
    if isinstance(outputs, list):
        return _Function(inputs, outputs, updates, givens=givens)
    elif isinstance(outputs, (dict, collections.OrderedDict)):
        func = _Function(inputs, outputs.values(), updates, givens=givens)
        return lambda *args, **kwargs: type(outputs)(
            zip(outputs.keys(), func(*args, **kwargs))
        )
    else:
        func = _Function(inputs, [outputs], updates, givens=givens)
        return lambda *args, **kwargs: func(*args, **kwargs)[0]


class _Function(object):
    def __init__(self, inputs, outputs, updates, givens):
        for inpt in inputs:
            if not hasattr(inpt, 'make_feed_dict') and not (
                isinstance(inpt, tf.Tensor) and len(inpt.op.inputs) == 0
            ):
                assert (
                    False
                ), "inputs should all be placeholders, constants, or have a make_feed_dict method"
        self.inputs = inputs
        updates = updates or []
        self.update_group = tf.group(*updates)
        self.outputs_update = list(outputs) + [self.update_group]
        self.givens = {} if givens is None else givens

    @classmethod
    def _feed_input(cls, feed_dict, inpt, value):
        if hasattr(inpt, 'make_feed_dict'):
            feed_dict.update(inpt.make_feed_dict(value))
        else:
            feed_dict[inpt] = value

    def __call__(self, *args, sess=None, **kwargs):
        assert len(args) <= len(self.inputs), "Too many arguments provided"
        if sess is None:
            sess = tf.compat.v1.get_default_session()
        feed_dict = {}
        for inpt, value in zip(self.inputs, args):
            self._feed_input(feed_dict, inpt, value)
        for inpt in self.givens:
            feed_dict[inpt] = feed_dict.get(inpt, self.givens[inpt])
        results = sess.run(self.outputs_update, feed_dict=feed_dict, **kwargs)[:-1]
        return results


def var_shape(tensor):
    out = tensor.get_shape().as_list()
    assert all(
        isinstance(a, int) for a in out
    ), "shape function assumes that shape is fully known"
    return out


def numel(tensor):
    return intprod(var_shape(tensor))


def intprod(tensor):
    return int(np.prod(tensor))


def flatgrad(loss, var_list, clip_norm=None):
    grads = tf.gradients(ys=loss, xs=var_list)
    if clip_norm is not None:
        grads = [tf.clip_by_norm(grad, clip_norm=clip_norm) for grad in grads]
    return tf.concat(
        axis=0,
        values=[
            tf.reshape(grad if grad is not None else tf.zeros_like(v), [numel(v)])
            for (v, grad) in zip(var_list, grads)
        ],
    )


class SetFromFlat(object):
    def __init__(self, var_list, dtype=tf.float32, sess=None):
        shapes = list(map(var_shape, var_list))
        total_size = np.sum([intprod(shape) for shape in shapes])
        self.theta = theta = tf.compat.v1.placeholder(dtype, [total_size])
        start = 0
        assigns = []
        for (shape, _var) in zip(shapes, var_list):
            size = intprod(shape)
            assigns.append(
                tf.compat.v1.assign(
                    _var, tf.reshape(theta[start : start + size], shape)
                )
            )
            start += size
        self.operation = tf.group(*assigns)
        self.sess = sess

    def __call__(self, theta):
        if self.sess is None:
            return tf.compat.v1.get_default_session().run(
                self.operation, feed_dict={self.theta: theta}
            )
        else:
            return self.sess.run(self.operation, feed_dict={self.theta: theta})


class GetFlat(object):
    def __init__(self, var_list, sess=None):
        self.operation = tf.concat(
            axis=0, values=[tf.reshape(v, [numel(v)]) for v in var_list]
        )
        self.sess = sess

    def __call__(self):
        if self.sess is None:
            return tf.compat.v1.get_default_session().run(self.operation)
        else:
            return self.sess.run(self.operation)


def get_trainable_vars(name):
    return tf.compat.v1.get_collection(
        tf.compat.v1.GraphKeys.TRAINABLE_VARIABLES, scope=name
    )


def get_globals_vars(name):
    return tf.compat.v1.get_collection(
        tf.compat.v1.GraphKeys.GLOBAL_VARIABLES, scope=name
    )


def outer_scope_getter(scope, new_scope=""):
    def _getter(getter, name, *args, **kwargs):
        name = name.replace(scope + "/", new_scope, 1)
        val = getter(name, *args, **kwargs)
        return val

    return _getter


def total_episode_reward_logger(rew_acc, rewards, masks, writer, steps):
    with tf.compat.v1.variable_scope("environment_info", reuse=True):
        for env_idx in range(rewards.shape[0]):
            dones_idx = np.sort(np.argwhere(masks[env_idx]))
            if len(dones_idx) == 0:
                rew_acc[env_idx] += sum(rewards[env_idx])
            else:
                rew_acc[env_idx] += sum(rewards[env_idx, : dones_idx[0, 0]])
                summary = tf.compat.v1.Summary(
                    value=[
                        tf.compat.v1.Summary.Value(
                            tag="episode_reward", simple_value=rew_acc[env_idx]
                        )
                    ]
                )
                writer.add_summary(summary, steps + dones_idx[0, 0])
                for k in range(1, len(dones_idx[:, 0])):
                    rew_acc[env_idx] = sum(
                        rewards[env_idx, dones_idx[k - 1, 0] : dones_idx[k, 0]]
                    )
                    summary = tf.compat.v1.Summary(
                        value=[
                            tf.compat.v1.Summary.Value(
                                tag="episode_reward", simple_value=rew_acc[env_idx]
                            )
                        ]
                    )
                    writer.add_summary(summary, steps + dones_idx[k, 0])
                rew_acc[env_idx] = sum(rewards[env_idx, dones_idx[-1, 0] :])
    return rew_acc

import collections
import functools
import multiprocessing
import os
from typing import Set

import numpy as np
import tensorflow as tf


def is_image(tensor):
    return len(tensor.shape) == 3 and tensor.shape[-1] in [1, 3, 4]


def batch_to_seq(tensor_batch, n_batch, n_steps, flat=False):
    if flat:
        tensor_batch = tf.reshape(tensor_batch, [n_batch, n_steps])
    else:
        tensor_batch = tf.reshape(tensor_batch, [n_batch, n_steps, -1])
    return [
        tf.squeeze(v, [1])
        for v in tf.split(axis=1, num_or_size_splits=n_steps, value=tensor_batch)
    ]


def seq_to_batch(tensor_sequence, flat=False):
    shape = tensor_sequence[0].get_shape().as_list()
    if not flat:
        assert len(shape) > 1
        n_hidden = tensor_sequence[0].get_shape()[-1].value
        return tf.reshape(tf.concat(axis=1, values=tensor_sequence), [-1, n_hidden])
    else:
        return tf.reshape(tf.stack(values=tensor_sequence, axis=1), [-1])


def check_shape(tensors, shapes):
    i = 0
    for (tensor, shape) in zip(tensors, shapes):
        assert tensor.get_shape().as_list() == shape, (
            "id " + str(i) + " shape " + str(tensor.get_shape()) + str(shape)
        )
        i += 1


def huber_loss(tensor, delta=1.0):
    return tf.compat.v1.where(
        tf.abs(tensor) < delta,
        tf.square(tensor) * 0.5,
        delta * (tf.abs(tensor) - 0.5 * delta),
    )


def sample(logits):
    noise = tf.random.uniform(tf.shape(input=logits))
    return tf.argmax(input=logits - tf.math.log(-tf.math.log(noise)), axis=1)


def calc_entropy(logits):
    a_0 = logits - tf.reduce_max(input_tensor=logits, axis=1, keepdims=True)
    exp_a_0 = tf.exp(a_0)
    z_0 = tf.reduce_sum(input_tensor=exp_a_0, axis=1, keepdims=True)
    p_0 = exp_a_0 / z_0
    return tf.reduce_sum(input_tensor=p_0 * (tf.math.log(z_0) - a_0), axis=1)


def mse(pred, target):
    return tf.reduce_mean(input_tensor=tf.square(pred - target))


def avg_norm(tensor):
    return tf.reduce_mean(
        input_tensor=tf.sqrt(tf.reduce_sum(input_tensor=tf.square(tensor), axis=-1))
    )


def gradient_add(grad_1, grad_2, param, verbose=0):
    if verbose > 1:
        print([grad_1, grad_2, param.name])
    if grad_1 is None and grad_2 is None:
        return None
    elif grad_1 is None:
        return grad_2
    elif grad_2 is None:
        return grad_1
    else:
        return grad_1 + grad_2


def q_explained_variance(q_pred, q_true):
    _, var_y = tf.nn.moments(x=q_true, axes=[0, 1])
    _, var_pred = tf.nn.moments(x=q_true - q_pred, axes=[0, 1])
    check_shape([var_y, var_pred], [[]] * 2)
    return 1.0 - (var_pred / var_y)


def make_session(num_cpu=None, make_default=False, graph=None):
    if num_cpu is None:
        num_cpu = int(os.getenv('RCALL_NUM_CPU', multiprocessing.cpu_count()))
    tf_config = tf.compat.v1.ConfigProto(
        allow_soft_placement=True,
        inter_op_parallelism_threads=num_cpu,
        intra_op_parallelism_threads=num_cpu,
    )
    tf_config.gpu_options.allow_growth = True
    if make_default:
        return tf.compat.v1.InteractiveSession(config=tf_config, graph=graph)
    else:
        return tf.compat.v1.Session(config=tf_config, graph=graph)


def single_threaded_session(make_default=False, graph=None):
    return make_session(num_cpu=1, make_default=make_default, graph=graph)


def in_session(func):
    @functools.wraps(func)
    def newfunc(*args, **kwargs):
        with tf.compat.v1.Session():
            func(*args, **kwargs)

    return newfunc


ALREADY_INITIALIZED = set()


def initialize(sess=None):
    if sess is None:
        sess = tf.compat.v1.get_default_session()
    new_variables = set(tf.compat.v1.global_variables()) - ALREADY_INITIALIZED
    sess.run(tf.compat.v1.variables_initializer(new_variables))
    ALREADY_INITIALIZED.update(new_variables)


def function(inputs, outputs, updates=None, givens=None):
    if isinstance(outputs, list):
        return _Function(inputs, outputs, updates, givens=givens)
    elif isinstance(outputs, (dict, collections.OrderedDict)):
        func = _Function(inputs, outputs.values(), updates, givens=givens)
        return lambda *args, **kwargs: type(outputs)(
            zip(outputs.keys(), func(*args, **kwargs))
        )
    else:
        func = _Function(inputs, [outputs], updates, givens=givens)
        return lambda *args, **kwargs: func(*args, **kwargs)[0]


class _Function(object):
    def __init__(self, inputs, outputs, updates, givens):
        for inpt in inputs:
            if not hasattr(inpt, 'make_feed_dict') and not (
                isinstance(inpt, tf.Tensor) and len(inpt.op.inputs) == 0
            ):
                assert (
                    False
                ), "inputs should all be placeholders, constants, or have a make_feed_dict method"
        self.inputs = inputs
        updates = updates or []
        self.update_group = tf.group(*updates)
        self.outputs_update = list(outputs) + [self.update_group]
        self.givens = {} if givens is None else givens

    @classmethod
    def _feed_input(cls, feed_dict, inpt, value):
        if hasattr(inpt, 'make_feed_dict'):
            feed_dict.update(inpt.make_feed_dict(value))
        else:
            feed_dict[inpt] = value

    def __call__(self, *args, sess=None, **kwargs):
        assert len(args) <= len(self.inputs), "Too many arguments provided"
        if sess is None:
            sess = tf.compat.v1.get_default_session()
        feed_dict = {}
        for inpt, value in zip(self.inputs, args):
            self._feed_input(feed_dict, inpt, value)
        for inpt in self.givens:
            feed_dict[inpt] = feed_dict.get(inpt, self.givens[inpt])
        results = sess.run(self.outputs_update, feed_dict=feed_dict, **kwargs)[:-1]
        return results


def var_shape(tensor):
    out = tensor.get_shape().as_list()
    assert all(
        isinstance(a, int) for a in out
    ), "shape function assumes that shape is fully known"
    return out


def numel(tensor):
    return intprod(var_shape(tensor))


def intprod(tensor):
    return int(np.prod(tensor))


def flatgrad(loss, var_list, clip_norm=None):
    grads = tf.gradients(ys=loss, xs=var_list)
    if clip_norm is not None:
        grads = [tf.clip_by_norm(grad, clip_norm=clip_norm) for grad in grads]
    return tf.concat(
        axis=0,
        values=[
            tf.reshape(grad if grad is not None else tf.zeros_like(v), [numel(v)])
            for (v, grad) in zip(var_list, grads)
        ],
    )


class SetFromFlat(object):
    def __init__(self, var_list, dtype=tf.float32, sess=None):
        shapes = list(map(var_shape, var_list))
        total_size = np.sum([intprod(shape) for shape in shapes])
        self.theta = theta = tf.compat.v1.placeholder(dtype, [total_size])
        start = 0
        assigns = []
        for (shape, _var) in zip(shapes, var_list):
            size = intprod(shape)
            assigns.append(
                tf.compat.v1.assign(
                    _var, tf.reshape(theta[start : start + size], shape)
                )
            )
            start += size
        self.operation = tf.group(*assigns)
        self.sess = sess

    def __call__(self, theta):
        if self.sess is None:
            return tf.compat.v1.get_default_session().run(
                self.operation, feed_dict={self.theta: theta}
            )
        else:
            return self.sess.run(self.operation, feed_dict={self.theta: theta})


class GetFlat(object):
    def __init__(self, var_list, sess=None):
        self.operation = tf.concat(
            axis=0, values=[tf.reshape(v, [numel(v)]) for v in var_list]
        )
        self.sess = sess

    def __call__(self):
        if self.sess is None:
            return tf.compat.v1.get_default_session().run(self.operation)
        else:
            return self.sess.run(self.operation)


def get_trainable_vars(name):
    return tf.compat.v1.get_collection(
        tf.compat.v1.GraphKeys.TRAINABLE_VARIABLES, scope=name
    )


def get_globals_vars(name):
    return tf.compat.v1.get_collection(
        tf.compat.v1.GraphKeys.GLOBAL_VARIABLES, scope=name
    )


def outer_scope_getter(scope, new_scope=""):
    def _getter(getter, name, *args, **kwargs):
        name = name.replace(scope + "/", new_scope, 1)
        val = getter(name, *args, **kwargs)
        return val

    return _getter


def total_episode_reward_logger(rew_acc, rewards, masks, writer, steps):
    with tf.compat.v1.variable_scope("environment_info", reuse=True):
        for env_idx in range(rewards.shape[0]):
            dones_idx = np.sort(np.argwhere(masks[env_idx]))
            if len(dones_idx) == 0:
                rew_acc[env_idx] += sum(rewards[env_idx])
            else:
                rew_acc[env_idx] += sum(rewards[env_idx, : dones_idx[0, 0]])
                summary = tf.compat.v1.Summary(
                    value=[
                        tf.compat.v1.Summary.Value(
                            tag="episode_reward", simple_value=rew_acc[env_idx]
                        )
                    ]
                )
                writer.add_summary(summary, steps + dones_idx[0, 0])
                for k in range(1, len(dones_idx[:, 0])):
                    rew_acc[env_idx] = sum(
                        rewards[env_idx, dones_idx[k - 1, 0] : dones_idx[k, 0]]
                    )
                    summary = tf.compat.v1.Summary(
                        value=[
                            tf.compat.v1.Summary.Value(
                                tag="episode_reward", simple_value=rew_acc[env_idx]
                            )
                        ]
                    )
                    writer.add_summary(summary, steps + dones_idx[k, 0])
                rew_acc[env_idx] = sum(rewards[env_idx, dones_idx[-1, 0] :])
    return rew_acc
