from math import *

def npars ( f ) :
	"Return the number of parameters a function has."
	if hasattr(f,'npars') : return f.npars # home-brewed, for *args
	return f.__code__.co_argcount


def deco( f , np = None ) :

	"Add argument checking to f and allow a list/tuple of the appropriate size"

	if np : f.npars = np
	nreq = npars(f)

	def ff( *args ):

		if len(args) == nreq : 
			return f(*args)

		if len(args)==1 and type(args[0]) in (list,tuple) and len(args[0]) == nreq :
			return f( *args[0] )

		print ("wrong arguments to function", f, "expecting", nreq , np )
		print ("args=", args )
		raise TypeError

	ff.npars = nreq
	return ff


def derivative( f  ) :
	"Returns derivative of f. f should have exactly one float arugment, but may be multi-valued"

	h = 1e-12

	def d(x) :
		f2, f1 = f(x+h), f( x-h)
		if type(f1) in (list, tuple) :
			return [ (v2-v1) / (2*h) for v1,v2 in zip(f1,f2) ]
		return (f2-f1) / (2*h)

	return d;


def partial_derivative( f, i ) :
	"If f(X), return derivative wrt X_i, which is a function of X"

	def r(*X) :
		def fi( xi ) :
			return f(*(X[:i] +(xi,)+ X[i+1:]) )
		return derivative( fi ) ( X[i] ) 

	return deco( r , npars(f) )


def gradient( f ) :  
	"Returns the gradient function of f"
	n = npars(f)
	r = lambda *X : [ partial_derivative(f,i)(*X) for i in range(n) ]
	return deco( r, n )


def hessian( f ) :
	"returns the hessian function of f"	
	r = lambda *X : gradient( gradient( f ) )(*X)
	return deco( r, npars(f) )


def test_functional() :
	def f(x,y) : return x**2 + 3* y**2
	h = hessian( f )(0,0) # should be [ [ 2,0 ] , [ 0, 6 ] ]
	print (h)