"""
Adapted from OpenAI Baselines.
"""

import numpy as np
import tensorflow as tf  # pylint: ignore-module
import random
import copy
import os
import functools
import collections
import multiprocessing

def switch(condition, then_expression, else_expression):
    """Switches between two operations depending on a scalar value (int or bool).
    Note that both `then_expression` and `else_expression`
    should be symbolic tensors of the *same shape*.

    # Arguments
        condition: scalar tensor.
        then_expression: TensorFlow operation.
        else_expression: TensorFlow operation.
    """
    x_shape = copy.copy(then_expression.get_shape())
    x = tf.cond(tf.cast(condition, 'bool'),
                lambda: then_expression,
                lambda: else_expression)
    x.set_shape(x_shape)
    return x

# ================================================================
# Extras
# ================================================================

def lrelu(x, leak=0.2):
    f1 = 0.5 * (1 + leak)
    f2 = 0.5 * (1 - leak)
    return f1 * x + f2 * abs(x)

# ================================================================
# Mathematical utils
# ================================================================

def huber_loss(x, delta=1.0):
    """Reference: https://en.wikipedia.org/wiki/Huber_loss"""
    return tf.where(
        tf.abs(x) < delta,
        tf.square(x) * 0.5,
        delta * (tf.abs(x) - 0.5 * delta)
    )

# ================================================================
# Global session
# ================================================================

def make_session(num_cpu=None, make_default=False):
    """Returns a session that will use <num_cpu> CPU's only"""
    if num_cpu is None:
        num_cpu = int(os.getenv('RCALL_NUM_CPU', multiprocessing.cpu_count()))
    tf_config = tf.ConfigProto(
        inter_op_parallelism_threads=num_cpu,
        intra_op_parallelism_threads=num_cpu)
    tf_config.gpu_options.allocator_type = 'BFC'
    if make_default:
        return tf.InteractiveSession(config=tf_config)
    else:
        return tf.Session(config=tf_config)

def single_threaded_session():
    """Returns a session which will only use a single CPU"""
    return make_session(num_cpu=1)

def in_session(f):
    @functools.wraps(f)
    def newfunc(*args, **kwargs):
        with tf.Session():
            f(*args, **kwargs)
    return newfunc

ALREADY_INITIALIZED = set()

def initialize():
    """Initialize all the uninitialized variables in the global scope."""
    new_variables = set(tf.global_variables()) - ALREADY_INITIALIZED
    tf.get_default_session().run(tf.variables_initializer(new_variables))
    ALREADY_INITIALIZED.update(new_variables)

# ================================================================
# Saving variables and setting up experiment directories
# ================================================================

def load_state(fname):
    saver = tf.train.Saver()
    saver.restore(tf.get_default_session(), fname)

def save_state(fname):
    os.makedirs(os.path.dirname(fname), exist_ok=True)
    saver = tf.train.Saver()
    saver.save(tf.get_default_session(), fname)

def load(fname):
    import cloudpickle
    with open(fname, 'rb') as f:
        return cloudpickle.load(f)

def save(fname, obj):
    import cloudpickle
    os.makedirs(os.path.dirname(fname), exist_ok=True)
    with open(fname, 'wb') as fh:
        cloudpickle.dump(obj, fh)

class Experiment(object):
    def __init__(self, logdir):
        self.logdir = logdir
        os.makedirs(os.path.join(logdir, 'checkpoints'), exist_ok=True)

    def load(self, timestep=None):
        if timestep is None:
            # get latest ckpt
            import glob
            fs = glob.glob(os.path.join(self.logdir, 'checkpoints/*'))
            timesteps = []
            for f in fs:
                try: timesteps.append(int(os.path.basename(f)))
                except: pass
            if len(timesteps) == 0:
                return 0
            timestep = max(timesteps)
        fname = os.path.join(self.logdir, 'checkpoints', str(timestep), 'model')
        load_state(fname)
        return timestep

    def save(self, timestep):
        fname = os.path.join(self.logdir, 'checkpoints', str(timestep), 'model')
        save_state(fname)

    def load_model_fn(self):
        fname = os.path.join(self.logdir, 'checkpoints/model_fn.pkl')
        assert os.path.exists(fname), "No model file saved."
        return load(fname)

    def save_model_fn(self, model_fn):
        fname = os.path.join(self.logdir, 'checkpoints/model_fn.pkl')
        save(fname, model_fn)


# ================================================================
# Model components
# ================================================================

def batch_to_seq(h, nbatch, nsteps, flat=False):
    """
    Assumes Time major data!!
    x.shape = [nsteps, nbatch, *obs_shape]
    h = x.reshape([-1, *x.shape[2:]]))
    """
    if flat:
        h = tf.reshape(h, [nsteps, nbatch])
    else:
        h = tf.reshape(h, [nsteps, nbatch, -1])
    return [tf.squeeze(v, [0]) for v in tf.split(axis=0, num_or_size_splits=nsteps, value=h)]

def seq_to_batch(h, flat = False):
    """
    Assumes Time major data!!
    x.shape = [nsteps, nbatch, *obs_shape]
    x = output.reshape(nsteps, nbatch, *obs_shape), where output is the output of this function.
    """
    shape = h[0].get_shape().as_list()
    if not flat:
        assert(len(shape) > 1)
        nh = h[0].get_shape()[-1].value
        return tf.reshape(tf.concat(axis=0, values=h), [-1, nh])
    else:
        return tf.reshape(tf.stack(values=h, axis=0), [-1])

def ortho_init(scale=1.0):
    def _ortho_init(shape, dtype, partition_info=None):
        #lasagne ortho init for tf
        shape = tuple(shape)
        if len(shape) == 2:
            flat_shape = shape
        elif len(shape) == 4: # assumes NHWC
            flat_shape = (np.prod(shape[:-1]), shape[-1])
        else:
            raise NotImplementedError
        a = np.random.normal(0.0, 1.0, flat_shape)
        u, _, v = np.linalg.svd(a, full_matrices=False)
        q = u if u.shape == flat_shape else v # pick the one with the correct shape
        q = q.reshape(shape)
        return (scale * q[:shape[0], :shape[1]]).astype(np.float32)
    return _ortho_init

def normc_initializer(std=1.0):
    def _initializer(shape, dtype=None, partition_info=None):  # pylint: disable=W0613
        out = np.random.randn(*shape).astype(np.float32)
        out *= std / np.sqrt(np.square(out).sum(axis=0, keepdims=True))
        return tf.constant(out)
    return _initializer

def lstm(xs, ms, s, scope, nh, init_scale=1.0):
    nbatch, nin = [v.value for v in xs[0].get_shape()]
    nsteps = len(xs)
    with tf.variable_scope(scope):
        wx = tf.get_variable("wx", [nin, nh*4], initializer=ortho_init(init_scale))
        wh = tf.get_variable("wh", [nh, nh*4], initializer=ortho_init(init_scale))
        b = tf.get_variable("b", [nh*4], initializer=tf.constant_initializer(0.0))

    c, h = tf.split(axis=1, num_or_size_splits=2, value=s)
    for idx, (x, m) in enumerate(zip(xs, ms)):
        c = c*(1-m)
        h = h*(1-m)
        z = tf.matmul(x, wx) + tf.matmul(h, wh) + b
        i, f, o, u = tf.split(axis=1, num_or_size_splits=4, value=z)
        i = tf.nn.sigmoid(i)
        f = tf.nn.sigmoid(f)
        o = tf.nn.sigmoid(o)
        u = tf.tanh(u)
        c = f*c + i*u
        h = o*tf.tanh(c)
        xs[idx] = h
    s = tf.concat(axis=1, values=[c, h])
    return xs, s


# ================================================================
# Theano-like Function
# ================================================================

def function(inputs, outputs, updates=None, givens=None):
    """Just like Theano function. Take a bunch of tensorflow placeholders and expressions
    computed based on those placeholders and produces f(inputs) -> outputs. Function f takes
    values to be fed to the input's placeholders and produces the values of the expressions
    in outputs.

    Input values can be passed in the same order as inputs or can be provided as kwargs based
    on placeholder name (passed to constructor or accessible via placeholder.op.name).

    Example:
        x = tf.placeholder(tf.int32, (), name="x")
        y = tf.placeholder(tf.int32, (), name="y")
        z = 3 * x + 2 * y
        lin = function([x, y], z, givens={y: 0})

        with single_threaded_session():
            initialize()

            assert lin(2) == 6
            assert lin(x=3) == 9
            assert lin(2, 2) == 10
            assert lin(x=2, y=3) == 12

    Parameters
    ----------
    inputs: [tf.placeholder, tf.constant, or object with make_feed_dict method]
        list of input arguments
    outputs: [tf.Variable] or tf.Variable
        list of outputs or a single output to be returned from function. Returned
        value will also have the same shape.
    """
    if isinstance(outputs, list):
        return _Function(inputs, outputs, updates, givens=givens)
    elif isinstance(outputs, (dict, collections.OrderedDict)):
        f = _Function(inputs, outputs.values(), updates, givens=givens)
        return lambda *args, **kwargs: type(outputs)(zip(outputs.keys(), f(*args, **kwargs)))
    else:
        f = _Function(inputs, [outputs], updates, givens=givens)
        return lambda *args, **kwargs: f(*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 (type(inpt) is 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

    def _feed_input(self, 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):
        assert len(args) <= len(self.inputs), "Too many arguments provided"
        feed_dict = {}
        # Update the args
        for inpt, value in zip(self.inputs, args):
            if value is not None:
                self._feed_input(feed_dict, inpt, value)
        # Update feed dict with givens.
        for inpt in self.givens:
            feed_dict[inpt] = feed_dict.get(inpt, self.givens[inpt])
        results = tf.get_default_session().run(self.outputs_update, feed_dict=feed_dict)[:-1]
        return results

# ================================================================
# Flat vectors
# ================================================================

def var_shape(x):
    out = x.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(x):
    return intprod(var_shape(x))

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

def flatgrad(loss, var_list, clip_norm=None):
    grads = tf.gradients(loss, var_list)
    if clip_norm is not None:
        grads, _ = tf.clip_by_global_norm(grads, clip_norm=clip_norm)
    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):
        assigns = []
        shapes = list(map(var_shape, var_list))
        total_size = np.sum([intprod(shape) for shape in shapes])

        self.theta = theta = tf.placeholder(dtype, [total_size])
        start = 0
        assigns = []
        for (shape, v) in zip(shapes, var_list):
            size = intprod(shape)
            assigns.append(tf.assign(v, tf.reshape(theta[start:start + size], shape)))
            start += size
        self.op = tf.group(*assigns)

    def __call__(self, theta):
        tf.get_default_session().run(self.op, feed_dict={self.theta: theta})

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

    def __call__(self):
        return tf.get_default_session().run(self.op)

def flattenallbut0(x):
    return tf.reshape(x, [-1, intprod(x.get_shape().as_list()[1:])])

def reset():
    global ALREADY_INITIALIZED
    ALREADY_INITIALIZED = set()
    tf.reset_default_graph()

"""
Random Seeds
"""
def set_global_seeds(i):
    try:
        import tensorflow as tf
    except ImportError:
        pass
    else:
        tf.set_random_seed(i)
    np.random.seed(i)
    random.seed(i)

[(60, 'tensorflow.ConfigProto', 'tf.ConfigProto', True, 'import tensorflow as tf\n'), (74, 'functools.wraps', 'functools.wraps', False, 'import functools\n'), (93, 'tensorflow.train.Saver', 'tf.train.Saver', True, 'import tensorflow as tf\n'), (98, 'tensorflow.train.Saver', 'tf.train.Saver', True, 'import tensorflow as tf\n'), (209, 'tensorflow.split', 'tf.split', True, 'import tensorflow as tf\n'), (222, 'tensorflow.concat', 'tf.concat', True, 'import tensorflow as tf\n'), (318, 'tensorflow.gradients', 'tf.gradients', True, 'import tensorflow as tf\n'), (357, 'tensorflow.reset_default_graph', 'tf.reset_default_graph', True, 'import tensorflow as tf\n'), (369, 'numpy.random.seed', 'np.random.seed', True, 'import numpy as np\n'), (370, 'random.seed', 'random.seed', False, 'import random\n'), (25, 'tensorflow.cast', 'tf.cast', True, 'import tensorflow as tf\n'), (65, 'tensorflow.InteractiveSession', 'tf.InteractiveSession', True, 'import tensorflow as tf\n'), (67, 'tensorflow.Session', 'tf.Session', True, 'import tensorflow as tf\n'), (85, 'tensorflow.variables_initializer', 'tf.variables_initializer', True, 'import tensorflow as tf\n'), (94, 'tensorflow.get_default_session', 'tf.get_default_session', True, 'import tensorflow as tf\n'), (97, 'os.path.dirname', 'os.path.dirname', False, 'import os\n'), (99, 'tensorflow.get_default_session', 'tf.get_default_session', True, 'import tensorflow as tf\n'), (104, 'cloudpickle.load', 'cloudpickle.load', False, 'import cloudpickle\n'), (108, 'os.path.dirname', 'os.path.dirname', False, 'import os\n'), (110, 'cloudpickle.dump', 'cloudpickle.dump', False, 'import cloudpickle\n'), (138, 'os.path.join', 'os.path.join', False, 'import os\n'), (139, 'os.path.exists', 'os.path.exists', False, 'import os\n'), (143, 'os.path.join', 'os.path.join', False, 'import os\n'), (158, 'tensorflow.reshape', 'tf.reshape', True, 'import tensorflow as tf\n'), (160, 'tensorflow.reshape', 'tf.reshape', True, 'import tensorflow as tf\n'), (161, 'tensorflow.squeeze', 'tf.squeeze', True, 'import tensorflow as tf\n'), (187, 'numpy.random.normal', 'np.random.normal', True, 'import numpy as np\n'), (188, 'numpy.linalg.svd', 'np.linalg.svd', True, 'import numpy as np\n'), (198, 'tensorflow.constant', 'tf.constant', True, 'import tensorflow as tf\n'), (204, 'tensorflow.variable_scope', 'tf.variable_scope', True, 'import tensorflow as tf\n'), (214, 'tensorflow.split', 'tf.split', True, 'import tensorflow as tf\n'), (215, 'tensorflow.nn.sigmoid', 'tf.nn.sigmoid', True, 'import tensorflow as tf\n'), (216, 'tensorflow.nn.sigmoid', 'tf.nn.sigmoid', True, 'import tensorflow as tf\n'), (217, 'tensorflow.nn.sigmoid', 'tf.nn.sigmoid', True, 'import tensorflow as tf\n'), (218, 'tensorflow.tanh', 'tf.tanh', True, 'import tensorflow as tf\n'), (278, 'tensorflow.group', 'tf.group', True, 'import tensorflow as tf\n'), (315, 'numpy.prod', 'np.prod', True, 'import numpy as np\n'), (320, 'tensorflow.clip_by_global_norm', 'tf.clip_by_global_norm', True, 'import tensorflow as tf\n'), (332, 'tensorflow.placeholder', 'tf.placeholder', True, 'import tensorflow as tf\n'), (339, 'tensorflow.group', 'tf.group', True, 'import tensorflow as tf\n'), (368, 'tensorflow.set_random_seed', 'tf.set_random_seed', True, 'import tensorflow as tf\n'), (47, 'tensorflow.abs', 'tf.abs', True, 'import tensorflow as tf\n'), (48, 'tensorflow.square', 'tf.square', True, 'import tensorflow as tf\n'), (76, 'tensorflow.Session', 'tf.Session', True, 'import tensorflow as tf\n'), (84, 'tensorflow.global_variables', 'tf.global_variables', True, 'import tensorflow as tf\n'), (85, 'tensorflow.get_default_session', 'tf.get_default_session', True, 'import tensorflow as tf\n'), (115, 'os.path.join', 'os.path.join', False, 'import os\n'), (161, 'tensorflow.split', 'tf.split', True, 'import tensorflow as tf\n'), (173, 'tensorflow.concat', 'tf.concat', True, 'import tensorflow as tf\n'), (175, 'tensorflow.stack', 'tf.stack', True, 'import tensorflow as tf\n'), (220, 'tensorflow.tanh', 'tf.tanh', True, 'import tensorflow as tf\n'), (49, 'tensorflow.abs', 'tf.abs', True, 'import tensorflow as tf\n'), (59, 'multiprocessing.cpu_count', 'multiprocessing.cpu_count', False, 'import multiprocessing\n'), (121, 'os.path.join', 'os.path.join', False, 'import os\n'), (196, 'numpy.random.randn', 'np.random.randn', True, 'import numpy as np\n'), (207, 'tensorflow.constant_initializer', 'tf.constant_initializer', True, 'import tensorflow as tf\n'), (213, 'tensorflow.matmul', 'tf.matmul', True, 'import tensorflow as tf\n'), (213, 'tensorflow.matmul', 'tf.matmul', True, 'import tensorflow as tf\n'), (342, 'tensorflow.get_default_session', 'tf.get_default_session', True, 'import tensorflow as tf\n'), (349, 'tensorflow.get_default_session', 'tf.get_default_session', True, 'import tensorflow as tf\n'), (184, 'numpy.prod', 'np.prod', True, 'import numpy as np\n'), (298, 'tensorflow.get_default_session', 'tf.get_default_session', True, 'import tensorflow as tf\n'), (337, 'tensorflow.reshape', 'tf.reshape', True, 'import tensorflow as tf\n'), (197, 'numpy.square', 'np.square', True, 'import numpy as np\n'), (322, 'tensorflow.zeros_like', 'tf.zeros_like', True, 'import tensorflow as tf\n'), (124, 'os.path.basename', 'os.path.basename', False, 'import os\n')]