import ROOT
from math import *
import inspect, ctypes, array




def maketree( **kwargs ) :

    """Make a branch for every name=type mentioned in the args.
       Pythonizes the tree so you can set the values with T.name = object 

	   e.g: 

	   maketree( time = TTimeStamp(), name = "x" )
    """

    name     = kwargs.pop( 'name'    , 'T')
    friendof = kwargs.pop( 'friendof',  None)

    T = ROOT.TTree( name , name ) 
    T.objs = []

    for name,typ in kwargs.items() :
        obj = typ()
        T.objs.append( obj )
        T.Branch( name,  obj  , 32000, 4 ) 
  
    if friendof : friendof.AddFriend(T)


    # pythonize the tree object

    def treeset( tree, **kwargs ) :
        for name,value in kwargs.items() :
            print ("asigning to", name )
            getattr( tree, name ).__assign__( value )

    T.set = types.MethodType( treeset, T ) # make it a (bound) method 

    # Change setattr so that we can do Tree.object = otherobject to do assignement
    # this has to be done at class-level.

    if not hasattr( ROOT.TTree, "__old_setattr" ) :

        ROOT.TTree.__old_setattr = ROOT.TTree.__setattr__

        def treesetattr(self,name,value):
            
            try : 
                self.set( **{name : value} )
                return value
            except AttributeError :
                pass    

            return self.__old_setattr( name, value )

        ROOT.TTree.__setattr__ = treesetattr

    return T



class Bin ( object ) :

	def __init__( self, hist, binnrs ) :

		self.hist = hist
		self.bins = binnrs[:]
		self.b    = hist.GetBin( *binnrs )

	@property
	def axes(self) :
		return ( self.hist.GetXaxis(), self.hist.GetYaxis(), self.hist.GetZaxis() )

	@property
	def content(self):
		return self.hist.GetBinContent( self.b )
	
	@content.setter
	def content( self, value ) :
		self.hist.SetBinContent( self.b, value )

def d(n,i) :
	setattr( Bin, f"i{n}"    , property( lambda self : self.bins[i] ) )
	setattr( Bin, f"{n}"     , property( lambda self : self.axes[i].GetBinCenter( self.bins[i] )))
	setattr( Bin, f"{n}low"  , property( lambda self : self.axes[i].GetBinLowEdge( self.bins[i])))
	setattr( Bin, f"{n}up"   , property( lambda self : self.axes[i].GetBinUpEdge( self.bins[i] )))

for i,n in enumerate( "xyz") :
	d(n,i)
		

def bins ( hist ) :

	"""
	Generator whos returned objects represent a bin. 
	Works for 1,2,3 D histograms .

	example
	for bin in bins( histogram ) :
		bin.content = some_func ( bin.x, bin.y )

	"""
	ok = False

	if type(hist) in ( ROOT.TH1F, ROOT.TH1D) :
		ok = True
		for ix in range( 1, hist.GetXaxis().GetNbins() + 1 ) :
			yield Bin ( hist, (ix,) )

	if type(hist) in ( ROOT.TH2F, ROOT.TH2D ) :
		ok = True
		for ix in range( 1, hist.GetXaxis().GetNbins() + 1 ) :
			for iy in range ( 1, hist.GetYaxis().GetNbins() + 1 ) :
				yield Bin( hist, (ix, iy) )

	if not ok : 
		print ("bins failes for", hist)
		raise TypeError
	



def minimize( function,
              start_values = None,
              ranges       = None,
              verbose      = False,
              maxcalls     = 1000000,
              tollerance   = 1e-8 ):
    
    """ 
    Minimize a python function 'function'. The function should have
    any number of (float) arguments. Python instrospection is used
    to know the number (npar) and names of the arguments. 

    ROOT's TMinuit is used to minimize the function with respect
    to all its arguments. The return value is a list with the parameters
    for which 'function' is minimial.

    function     : the function to be minimized
    
    start_values : list of values to use for the initial parameters,
                   default [0.0, 0.0, ... , 0.0]
    
    ranges       : list of lists, defining the min and max values of the 
                   parameters: [ [par1min, par1max], .... , [parnmin, parnmax] ].
                   (Important when you take e.g. logarithms in your 'function'!)
                   
    maxcalls     : maximum number of iterations 

    tollerance   : tollerance, see ROOT.TMinuit().mnhelp("MIGrad")
    """

    args = inspect.getargspec( function )[0]  # list of functions arguments
    npar = len(args)
    
    if not start_values : start_values = [0.0] * npar
    if not ranges       : ranges       = [[0,0]] * npar

    step_sizes   = [1] * npar
    
    minuit = ROOT.TMinuit( npar )


    # set parameter names, start_values, ...
    print(" ----------> npar=",npar)
    for i in range(npar) :
        ierflg = ctypes.c_int()
        step_sizes[i] = (ranges[i][1]-ranges[i][0])*0.0001
        minuit.mnparm( i, args[i] , start_values[i] , step_sizes[i], ranges[i][0], ranges[i][1], ierflg )
        print(" >>>>>>>>>>>>>>>>>",args[i],start_values[i] , step_sizes[i], ranges[i][0], ranges[i][1])

    def fcn( _npar, gin, f, par, iflag ):
        "Wrapper around function, to provide signature expected by minuit."

        # the buffer objects we get passed in behave strangly/buggy ->
        # dont use _npar (but npar instead), and unpack par by hand:
        parlist = [ par[i] for i in range (npar) ] 

        f.value = function( *parlist )

        if verbose :
            strpars = ",".join( str(x) for x in parlist )
            print ("evaluating %s(%s) = %g " % (function.__name__, strpars, f[0] ))

            
    # call migrad
    print(tollerance)
    minuit.SetFCN( fcn ) # tell minuit to use our wrapper
    print(tollerance)

    arglist = array.array( 'd', [maxcalls, tollerance] )
    print(arglist,ierflg)
    
    minuit.mnexcm( "SIMPLEX", arglist , len(arglist) , ierflg )
    minuit.mnexcm( "MIGRAD", arglist , len(arglist) , ierflg )


    # return a list of fitted values
    
    r, errs = [],[]
    for i in range( npar) :
        r.append( ctypes.c_double() )
        errs.append( ctypes.c_double() )
        minuit.GetParameter( i, r[-1], errs[-1] )
        
    return [ e.value for e in r ]





def set_attributes( rootobject , **args ) :

	""" shortcut to set properties of root-objects

	    example : 
	    set_attributes( histogram, LineColor = 3, FillStyle = 1001 )
	    will call
	    histogram.SetLineColor(3) and histogram.SetFillStyle(1001)
	"""

	for k,v in args.items() :
		try :
			getattr( rootobject, "Set"+k )( v )
		except AttributeError:
			print ("cant set attribute", k )


def fill_hist( func ,
			   nbins  = 100,
			   xrange = [0,1],
			   hist   = None,
			   clone  = False,
			   name   = None ,
			   verbose= False,
			   **args) :

	"""
	Fill a histogram with the result of a python funciton, func.
	"""

	name = name or "hist"+func.__name__
	hist = hist or ROOT.TH1D( name, name, nbins, *xrange )

	if clone :
		hist = hist.Clone( name )

	for bin in bins( hist ) :
		bin.content = func ( bin.x )

		if verbose :
			print ( bin.x, bin.content )


	set_attributes( hist, **args )

	return hist



def splitcanvas( N , name = "cc", size = [900,900] , **args ) :

	""" 
	Return a canvas with at least N sub-canvasses.
	"""

	ny = int(sqrt(N))
	nx = int(0.9999 + N / ny )
	print ("need",N,"sub-canvases, doing ", nx ,"x", ny )
	c = ROOT.TCanvas( name, name, *size )
	c.Divide( nx, ny )
	ROOT.SetOwnership( c, False )
	set_attributes(c, **args )
	return c


def equalize( histos , value = 0 ) :

	"set the max of each histgram to value"
	for h in histos :
		m = max( b.content for b in bins(h) )
		for b in bins(h) : b.content -= m + value


def draw( obj , opt ) :
	obj.DrawCopy(opt)


def draw_into( canv, objs , opt ="") :
	for i, o in enumerate( objs ) :
		canv.cd(i+1)
		draw( o, opt )

def canvas( objs , name = "cc", size = [900,900], opt = "") :

	c = splitcanvas( len(objs), name, size )
	draw_into( c, objs , opt )
	return c


def plot_funcs( template, *funcs ) :

	L = [ fill_hist( f, hist = template, clone= True ) for f in funcs ]
	return canvas( L )



#------------------------------------------------------------------------------
# 
#  Functional calculus 
#
#------------------------------------------------------------------------------


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

derivative_h = 1e-4




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

	h = derivative_h

	def subem( v1, v2 , a =1  ) :

		if type(v1) == float :
			return a*(v2-v1)
		if type(v1) in (list, tuple) :
			return [ subem( x1, x2, a ) for x1,x2 in zip( v1,v2 ) ]


	def d(x) :
		f2, f1 = f(x+h), f( x-h)
		return subem( f2,f1, 1.0/(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)