"""Convenience functions built on top of `grad`."""
from __future__ import absolute_import
import autograd.numpy as np
from autograd.core import grad, getval, jacobian
from collections import OrderedDict


def multigrad(fun, argnums=[0]):
    """Takes gradients wrt multiple arguments simultaneously."""
    def combined_arg_fun(multi_arg, *args, **kwargs):
        extra_args_list = list(args)
        for argnum_ix, arg_ix in enumerate(argnums):
            extra_args_list[arg_ix] = multi_arg[argnum_ix]
        return fun(*extra_args_list, **kwargs)
    gradfun = grad(combined_arg_fun, argnum=0)
    def gradfun_rearranged(*args, **kwargs):
        multi_arg = tuple([args[i] for i in argnums])
        return gradfun(multi_arg, *args, **kwargs)
    return gradfun_rearranged

def multigrad_dict(fun):
    "Takes gradients wrt all arguments simultaneously,"
    "returns a dict mapping 'argname' to 'gradval'"

    import funcsigs
    sig = funcsigs.signature(fun)

    def gradfun(*args, **kwargs):
        argdict = dict(sig.bind(*args, **kwargs).arguments)
        newfun = lambda dct: fun(**{argname:dct[argname] for argname in argdict})
        grad_dict = grad(newfun)(argdict)
        return OrderedDict((argname, grad_dict[argname]) for argname in argdict)

    return gradfun

def grad_and_aux(fun, argnum=0):
    """Builds a function that returns the gradient of the first output and the
    (unmodified) second output of a function that returns two outputs."""
    def grad_and_aux_fun(*args, **kwargs):
        saved_aux = []
        def return_val_save_aux(*args, **kwargs):
            val, aux = fun(*args, **kwargs)
            saved_aux.append(aux)
            return val
        gradval = grad(return_val_save_aux, argnum)(*args, **kwargs)
        return gradval, saved_aux[0]

    return grad_and_aux_fun

def value_and_grad(fun, argnum=0):
    """Returns a function that returns both value and gradient. Suitable for use
    in scipy.optimize"""
    def double_val_fun(*args, **kwargs):
        val = fun(*args, **kwargs)
        return val, getval(val)
    gradval_and_val = grad_and_aux(double_val_fun, argnum)

    def value_and_grad_fun(*args, **kwargs):
        gradval, val = gradval_and_val(*args, **kwargs)
        return val, gradval

    return value_and_grad_fun

def elementwise_grad(fun, argnum=0):
    """Like `jacobian`, but produces a function which computes just the diagonal
    of the Jacobian, and does the computation in one pass rather than in a loop.
    Note: this is only valid if the Jacobian is diagonal. Only arrays are
    currently supported."""
    def sum_output(*args, **kwargs):
        return np.sum(fun(*args, **kwargs))
    return grad(sum_output, argnum=argnum)

def hessian_vector_product(fun, argnum=0):
    """Builds a function that returns the exact Hessian-vector product.
    The returned function has arguments (*args, vector, **kwargs), and takes
    roughly 4x as long to evaluate as the original function."""
    fun_grad = grad(fun, argnum)
    def vector_dot_grad(*args, **kwargs):
        args, vector = args[:-1], args[-1]
        return np.dot(vector, fun_grad(*args, **kwargs))
    return grad(vector_dot_grad, argnum)  # Grad wrt original input.

def hessian(fun, argnum=0):
    "Returns a function that computes the exact Hessian."
    return jacobian(jacobian(fun, argnum), argnum)

@@ -197,16 +197,51 @@
 redDict%0A
+from future.utils import iteritems%0A
 %0A%0Adef mu
@@ -1045,23 +1045,19 @@
     
-argdict = dict(
+bindings = 
 sig.
@@ -1081,53 +1081,555 @@
 rgs)
-.arguments)%0A        newfun = lambda dct: fun(
+%0A        var_positional = next((name for name, parameter in sig.parameters.items()%0A                               if parameter.kind == parameter.VAR_POSITIONAL), None)%0A        var_keyword = next((name for name, parameter in sig.parameters.items()%0A                            if parameter.kind == parameter.VAR_KEYWORD), None)%0A        args = lambda dct: dct%5Bvar_positional%5D if var_positional else ()%0A        kwargs = lambda dct: %5C%0A            dict(((key, dct%5Bvar_keyword%5D%5Bkey%5D) for key in dct%5Bvar_keyword%5D) if var_keyword else (),%0A                 
 **%7Ba
@@ -1635,16 +1635,17 @@
 argname:
+ 
 dct%5Bargn
@@ -1664,20 +1664,173 @@
 name
- in argdict%7D
+, parameter in sig.parameters.items()%0A                    if parameter.kind != parameter.VAR_POSITIONAL%7D)%0A%0A        newfun = lambda dct: fun(*args(dct), **kwargs(dct)
 )%0A  
@@ -1860,23 +1860,40 @@
 newfun)(
-arg
 dict
+(bindings.arguments)
 )%0A      
@@ -1962,15 +1962,26 @@
  in 
-argdict
+bindings.arguments
 )%0A%0A 

make multigrad_dict handle functions with *args and **kwargs