'''
Automatic differentiation.
(C) R. Brette 2008
This file is released under the terms of the BSD licence
--------------------------------------------------------
Method: forward accumulation by operator overloading.

Use::
  df=differentiate(f)                  # where f is a single-variable function
  y=differentiate(f,x0)                # differentiate at x=x0
  y=differentiate(f,x0,order=2)        # second-order derivative at x=x0
  y=differentiate(f,(a,b),order=(1,2)) # d2f/dxdy(a,b)
  y=differentiate(f,order=(1,2))       # d2f/dxdy
  
  a0,a1,a2=taylor_series(f,x0,order=2) # 2nd order Taylor series of f at x0
  J=gradient(f,(a,b))                  # gradient of f at (a,b) (returns a list)
  H=hessian(f,(a,b))                   # hessian of f at (a,b) (returns a numpy array)
  
TODO:
* iadd, etc. (more operators)
* NonDifferentiableException and maybe use this on conditions (when cmp=0)
* Differential operators: Laplacian...
* gradient, taylor_series, hessian: x0 should be optional
* test with units, functions Rp->something, alternative differential operators,
  differential (f:Rp->Rn), complex numbers, arithmetic derivative
* remove most dependency with numpy (except Hessian)
* simplify the code

If several partial derivatives of the same function need to be
calculated, it is faster to use the lower level object Differentiable.
Example::
  y=f(Differentiable('x',3),Differentiable('y',2))
will return a Differentiable object with y.val=f(3,2)
and y.diff={'x': df/dx(3,2), 'y': df/dy(3,2)}
'''

from numpy import exp,sin,cos,log,zeros # todo: look in frame instead
import types

__all__=['differentiate','Differentiable','taylor_series','gradient','hessian']

def differentiate(f,x=None,order=1):
    '''
    Calculate the derivative of f at point x.
    Higher-order derivatives are possible.
    Ex.: differentiate(lambda x:3*x*x+2,x=2,order=2)
    (Returns: 6.0)
    
    Partial derivatives:
    1) Give a tuple with the order.
    Ex.: differentiate(lambda x,y:x*y+2*y+1,x=(1,2),order=(0,1))
    (Returns: 3.0 (partial derivative d/dy) )
    2) Give a dictionnary with the order:
    Ex.: differentiate(lambda x,y:x*y+2*y+1,x=(1,2),order={'y':1})
    '''
    if x is None:
        return lambda x:differentiate(f,x,order)
    if (type(order)==types.ListType) or (type(order)==types.TupleType): # several variables
        # Build the list of arguments
        n=sum(order) # total order
        args=[HigherDifferentiable(i,x[i],n) for i in range(len(x))]
        y=f(*args)
        for i in range(len(x)):
            for j in range(order[i]):
                if isinstance(y,Differentiable) and (i in y.diff):
                    y=y.diff[i]
                else: # constant
                    return 0.
        return y
    elif (type(order)==types.DictType): # named variables
        # Build the list of arguments
        n=sum(order.itervalues())
        vars=list(f.func_code.co_varnames)
        args=[]
        args.extend(x)
        for name in order.iterkeys():
            i=vars.index(name)
            args[i]=HigherDifferentiable(name,x[i],n)
        y=f(*args)
        for name,n in order.iteritems():
            for j in range(n):
                if isinstance(y,Differentiable) and (name in y.diff):
                    y=y.diff[name]
                else: # constant
                    return 0.
        return y
    elif order==0:
        return f(x)
    elif order==1:
        y=f(Differentiable('x',x))
        if isinstance(y,Differentiable) and ('x' in y.diff):
            return y.diff['x']
        else: # constant
            return 0.
    else:
        return differentiate(lambda y:differentiate(f,y,order=order-1),x)
    
def taylor_series(f,x,order=1):
    '''
    Returns the list of coefficients of the Taylor Series of f
    at x (scalar).
    '''
    y=f(HigherDifferentiable('x',x,order))
    series=[]
    n=1.
    for i in range(order+1):
        # Value
        val=y
        while isinstance(val,Differentiable):
            val=val.val
        series.append(val/n)
        n*=(i+1)
        # Next order
        if isinstance(y,Differentiable) and ('x' in y.diff):
            y=y.diff['x']
        else:
            y=0.
    return series

def gradient(f,x):
    '''
    Gradient of f at x (= tuple).
    Result = list.
    '''
    args=[Differentiable(i,val=x[i]) for i in range(len(x))]
    y=f(*args)
    if isinstance(y,Differentiable):
        result=[]
        for i in range(len(x)):
            if i in y.diff:
                result.append(y.diff[i])
            else:
                result.append(0.)
        return result
    else: # constant
        return [0.]*len(x)
    
def hessian(f,x):
    '''
    Hessian of f at x (= tuple).
    Result = array.
    N.B.: not working with vectors.
    '''
    args=[HigherDifferentiable(i,x[i],2) for i in range(len(x))]
    y=f(*args)
    result=zeros((len(x),len(x)))
    if isinstance(y,Differentiable): # non-zero Hessian
        for i in range(len(x)):
            if i in y.diff and isinstance(y.diff[i],Differentiable):
                for j in range(len(x)):
                    if j in y.diff[i].diff:
                        result[i,j]=y.diff[i].diff[j]
    return result
    
class Differentiable(object):
    '''
    A differentiable variable.
    Implemented operations: +,-,*,[],len,exp,sin,cos,/,abs,sqrt,comparisons,**
    '''
    def __init__(self,name=None,val=0.):
        '''
        Initializes a variable and its derivative.
        If the name is None, then it is a constant.
        '''
        self.val=val # value
        if name==None: # constant
            self.diff={} # derivative
        #elif isinstance(val,ndarray): # a vector
            # TODO: different class? (DifferentiableVector)
            #self.diff={name:eye_lil_matrix(len(val))} # Jacobian matrix
        else:
            self.diff={name:1.} # or jacobian? (eye)
            
    def sqrt(self):
        return self**.5
    
    def __cmp__(self,x):
        if isinstance(x,Differentiable):
            return cmp(self.val,x.val)
        else:
            return cmp(self.val,x)
            
    def __abs__(self):
        # NOT WORKING WITH VECTORS
        if self.val==0:
            # TODO: specific exception
            raise ZeroDivisionError,"x:abs(x) is not differentiable at x=0."
        elif self.val>0:
            return self
        else:
            return -self
            
    def __add__(self,y):
        # VECTOR-READY
        if isinstance(y,Differentiable):
            zdict={}
            zdict.update(self.diff)
            zdict.update(y.diff)
            for key in zdict.iterkeys():
                if (key in self.diff) and (key in y.diff):
                    zdict[key]=self.diff[key]+y.diff[key]
            z=Differentiable()
            z.val=self.val+y.val
            z.diff=zdict
            return z
        else: # Adding a constant
            z=Differentiable(val=self.val+y)
            z.diff=self.diff
            return z

    def __radd__(self,x):
        # VECTOR-READY
        if isinstance(x,Differentiable):
            zdict={}
            zdict.update(self.diff)
            zdict.update(x.diff)
            for key in zdict.iterkeys():
                if (key in self.diff) and (key in x.diff):
                    zdict[key]=x.diff[key]+self.diff[key]
            z=Differentiable()
            z.val=x.val+self.val
            z.diff=zdict
            return z
        else: # Adding a constant
            z=Differentiable(val=x+self.val)
            z.diff=self.diff
            return z
        
    def __mul__(self,y):
        # !Not working with vectors!
        if isinstance(y,Differentiable):
            #TODO: with vectors
            zdict={}
            zdict.update(self.diff)
            zdict.update(y.diff)
            for key in zdict.iterkeys():
                if (key in self.diff) and (key in y.diff):
                    zdict[key]=self.diff[key]*y.val+self.val*y.diff[key]
                elif (key in self.diff):
                    zdict[key]=self.diff[key]*y.val
                else:
                    zdict[key]=self.val*y.diff[key]
            z=Differentiable()
            z.val=self.val*y.val
            z.diff=zdict
            return z
        else: # Multiplying by a constant
            z=Differentiable(val=self.val*y)
            for key in self.diff:
                z.diff[key]=self.diff[key]*y
            return z
        
    def __rmul__(self,x):
        if isinstance(x,Differentiable):
            #TODO: with vectors
            zdict={}
            zdict.update(self.diff)
            zdict.update(x.diff)
            for key in zdict.iterkeys():
                if (key in self.diff) and (key in x.diff):
                    zdict[key]=x.val*self.diff[key]+x.diff[key]*self.val
                elif (key in self.diff):
                    zdict[key]=x.val*self.diff[key]
                else:
                    zdict[key]=x.diff[key]*self.val
            z=Differentiable()
            z.val=x.val*self.val
            z.diff=zdict
            return z
        else: # Multiplying by a constant
            z=Differentiable(val=x*self.val)
            for key in self.diff:
                z.diff[key]=x*self.diff[key]
            return z
        
    def __div__(self,x):
        return self*(x**-1)

    def __rdiv__(self,x):
        return x*(self**-1)
        
    def __sub__(self,y):
        # VECTOR-READY
        if isinstance(y,Differentiable):
            zdict={}
            zdict.update(self.diff)
            zdict.update(y.diff)
            for key in zdict.iterkeys():
                if (key in self.diff) and (key in y.diff):
                    zdict[key]=self.diff[key]-y.diff[key]
            z=Differentiable()
            z.val=self.val-y.val
            z.diff=zdict
            return z
        else: # Subtracting a constant
            z=Differentiable(val=self.val-y)
            z.diff=self.diff
            return z
        
    def __rsub__(self,x):
        # VECTOR-READY
        if isinstance(x,Differentiable):
            zdict={}
            zdict.update(self.diff)
            zdict.update(x.diff)
            for key in zdict.iterkeys():
                if (key in self.diff) and (key in x.diff):
                    zdict[key]=x.diff[key]-self.diff[key]
            z=Differentiable()
            z.val=x.val-self.val
            z.diff=zdict
            return z
        else: # Subtracting a constant
            z=Differentiable(val=x-self.val)
            for key in self.diff:
                z.diff[key]=-self.diff[key]
            return z
        
    def __neg__(self):
        # VECTOR-READY
        z=Differentiable(val=-self.val)
        for key in self.diff:
            z.diff[key]=-self.diff[key]
        return z
            
    def __pow__(self,x):
        # NOT WORKING WITH VECTORS
        if isinstance(x,Differentiable):
            return exp(x*log(self))        
        elif x==0:
            z=Differentiable(val=self.val**x)
            for key in self.diff:
                z.diff[key]=0*self.diff[key]
            return z
        else:
            z=Differentiable(val=self.val**x)
            for key in self.diff:
                z.diff[key]=x*self.val**(x-1)*self.diff[key]
            return z
        
    def __rpow__(self,x):
        # NOT WORKING WITH VECTORS
        return exp(self*log(x))        
            
    def __getitem__(self,i):
        z=Differentiable(val=self.val[i])
        z.diff={}
        for key in self.diff.iterkeys():
            z.diff[key]=self.diff[key][i,:]
        return z
    
    def __len__(self):
        return len(self.val)
            
    def __str__(self):
        s=str(self.val)
        for key,value in self.diff.iteritems():
            s+=' ; d/d'+key+'='+str(value)
        return s
    
    def exp(self):
        # !! Not working for vectors !!
        zdict={}
        zdict.update(self.diff)
        for key in zdict.iterkeys():
            zdict[key]=exp(self.val)*self.diff[key]
        z=Differentiable(val=exp(self.val))
        z.diff=zdict
        return z

    def cos(self):
        # !! Not working for vectors !!
        zdict={}
        zdict.update(self.diff)
        for key in zdict.iterkeys():
            zdict[key]=-sin(self.val)*self.diff[key]
        z=Differentiable(val=cos(self.val))
        z.diff=zdict
        return z

    def sin(self):
        # !! Not working for vectors !!
        zdict={}
        zdict.update(self.diff)
        for key in zdict.iterkeys():
            zdict[key]=cos(self.val)*self.diff[key]
        z=Differentiable(val=sin(self.val))
        z.diff=zdict
        return z
    
    def log(self):
        # !! Not working for vectors !!
        zdict={}
        zdict.update(self.diff)
        for key in zdict.iterkeys():
            zdict[key]=(1./(self.val))*self.diff[key]
        z=Differentiable(val=log(self.val))
        z.diff=zdict
        return z        

class HigherDifferentiable(Differentiable):
    '''
    A differentiable variable with order n.
    '''
    def __init__(self,name=None,val=0.,order=1):
        Differentiable.__init__(self,name,val)
        if order>1:
            self.val=HigherDifferentiable(name,val,order-1)
