# easy and fast reverse mode autodiff
#
# author:  justin domke
# history:
#   May 26-June  2009, First version 
#   
#   Aug 20-25 2009 next hacking spurt (converted to outputting
#   bytecode with a c interpreter, along with lots of other
#   efficiencies, then converted back to outputting c code after
#   discovering that tcc exists and can compile source codes 100x
#   larger than gcc can.  Now have about a 2x-10x speed advantage
#   over pycppad, depending on the problem and the compiler used.
#   (If we could somehow compile 600mb source code at O3 we would
#   be in business!)
# 
#  Nov 29, 2009 released
#
# THIS CODE IS SUBJECT TO THE MIT LICENSE:
# http://www.opensource.org/licenses/mit-license.php
#
# todo:
# convert code to output integers in hexadecimal to save space!?

import numpy
import array as a
from numpy import *
import copy
from scipy.weave import inline
from scipy.weave import converters
import os
import time

NOOP      = 0
MULT      = 1
MULT_N1   = 2
MULT_N1N2 = 3
MULT_N2N1 = 4
MULT___N1 = 5
MULT_REP  = 6
ADD       = 7
ADD_N1    = 8
ADD_N1N2  = 9
ADD_N2N1  = 10
ADD___N1  = 11
DIV       = 12
DIV_N1    = 13
DIV_N1N2  = 14
DIV_N2N1  = 15
DIV___N1  = 16
SUB       = 17
SUB_N1    = 18
SUB_N1N2  = 19
SUB_N2N1  = 20
SUB___N1  = 21
EXP       = 22
LOG       = 23

NOID = 987654321

# todo
# fix the insanity with the id numbers (not an issue!?)
# full id renumbering (prevent unnecessary allocations
# and allow user to declare arrays in any order.)

# todo
# why need to init D to zero!?

import sys
# this is a bad, bad, bad idea
sys.setrecursionlimit(1000000)

def independent(*args):
    eargs = []
    Atype = type(zeros(0))
    dtype = type(0.123)
    for n in range(len(args)):
        A = args[n]
        if type(A)==Atype:
            eargs.append(getmat(A.shape))
        elif type(A)==dtype:
            eargs.append(etree())
        else:
            eargs.append(independent(*A))
    return eargs    

def compile_fun(fun,*args):
    # first, make all the inputs into etrees
    reset()
    
    # coerce inputs to be floats
    args = copy.deepcopy(list(args))
    testmat  = numpy.zeros(1,dtype=float)
    mattype  = type(testmat)
    matdtype = testmat.dtype
    for i in range(len(args)):
        if type(args[i]) == mattype:
            if args[i].dtype != matdtype:
                args[i] = numpy.array(args[i],dtype=float)
        else:
            args[i] = float(args[i])

    outargs = copy.deepcopy(list(args))
    eargs = independent(*args)
    # now, run the function on the etree types
    z = fun(*eargs)
    f = z.compile(outargs,*eargs)
    # call the function once to force weave to actually compile it
    f(*args)
    return f


def getmat(N):    
    A = numpy.zeros(N,dtype=object)
    for position,element in numpy.ndenumerate(A):
        A[position] = etree()
    return A

def reset():
    etree.id = 0
    etree.ops  = a.array('B')
    etree.arg1 = a.array('l')
    etree.arg2 = a.array('l')

def gencode(ops,arg1s,arg2s,froz):
    nops = len(arg1s)

    # run some inline c code that will generate the workhorse c code
    writercode =  "const int NOOP = %d;\n" % NOOP
    writercode += "const int MULT = %d;\n" % MULT
    writercode += "const int ADD = %d;\n" % ADD
    writercode += "const int SUB = %d;\n" % SUB
    writercode += "const int DIV = %d;\n" % DIV
    writercode += "const int EXP = %d;\n" % EXP
    writercode += "const int LOG = %d;\n" % LOG
    writercode += "const int nops = %d;\n" % nops
    writercode += "int i;\n"

    # theoretically shouldn't need to cast to (int*)
    writercode += "int *touched = (int*)malloc(nops * sizeof(int));\n"
    writercode += "for(i=0; i<nops; i++)\n  touched[i]=0;\n"

    writercode += '''
    FILE *file;
    file = fopen("mycode.c","w");
    fprintf(file,"#include \\"math.h\\"\\nvoid rock(double* A, double* D){\\n");
    char arg1str[50];
    char arg2str[50];
    char arg1sym[50];
    char arg2sym[50];
    for(i=0; i<nops; i++){
      sprintf(arg1str,"A[%d]", arg1s[i]);
      sprintf(arg2str,"A[%d]", arg2s[i]);
      if(arg1s[i]<0)
        sprintf(arg1str,"%f",froz[-arg1s[i]-1]);
      if(arg2s[i]<0)
        sprintf(arg2str,"%f",froz[-arg2s[i]-1]);
      //if(ops[i])==NOOP){
        // do nothing
      //}
      if(ops[i]==MULT){
        fprintf(file,"A[%d]=%s * %s;\\n",i,arg1str,arg2str);
      } else if(ops[i]==ADD){
        fprintf(file,"A[%d]=%s + %s;\\n",i,arg1str,arg2str);
      } else if(ops[i]==SUB){
        fprintf(file,"A[%d]=%s - %s;\\n",i,arg1str,arg2str);
      } else if(ops[i]==DIV){
        fprintf(file,"A[%d]=%s / %s;\\n",i,arg1str,arg2str);
      } else if(ops[i]==EXP){
        fprintf(file,"A[%d]=exp(%s);\\n",i,arg1str);
      } else if(ops[i]==LOG){
        fprintf(file,"A[%d]=log(%s);\\n",i,arg1str);
      }
    }
    fprintf(file,"D[%d] = 1.0;\\n",nops-1);
    touched[nops-1]=1;
    for(i=nops-1; i>=0; i--){
      // if no-one has updated you, you have no impact
      if(touched[i]==0){
        //fprintf(file,"D[%d] = 0.0;\\n",i);
        continue;      
      }
    
      sprintf(arg1str,"A[%d]", arg1s[i]);
      sprintf(arg2str,"A[%d]", arg2s[i]);
      if(arg1s[i]<0)
        sprintf(arg1str,"%f",froz[-arg1s[i]-1]);
      if(arg2s[i]<0)
        sprintf(arg2str,"%f",froz[-arg2s[i]-1]);
    
      sprintf(arg1sym,"=");
      if(ops[i] > NOOP && arg1s[i] >= 0){
        if(touched[arg1s[i]]>0)
          sprintf(arg1sym,"+=");
        else
          touched[arg1s[i]] = 1;
      }
      sprintf(arg2sym,"=");
      if(ops[i] > NOOP && ops[i] < EXP && arg2s[i] >= 0){
        //printf("t: %d \\n", touched[arg2s[i]]);
        if(touched[arg2s[i]]>0)
          sprintf(arg2sym,"+=");
        else
          touched[arg2s[i]] = 1;
      }

      if(ops[i]==MULT){
            if(arg1s[i]>= 0)
                fprintf(file,"D[%d] %s D[%d] * %s;\\n", arg1s[i],arg1sym,i,arg2str); // HERE!
            if(arg2s[i] >= 0)
                fprintf(file,"D[%d] %s D[%d] * %s;\\n", arg2s[i],arg2sym,i,arg1str);
      } else if(ops[i]==ADD){
            if(arg1s[i] >= 0)
                fprintf(file,"D[%d] %s D[%d];\\n", arg1s[i],arg1sym,i);
            if(arg2s[i] >= 0)
                fprintf(file,"D[%d] %s D[%d];\\n", arg2s[i],arg2sym,i);
      } else if(ops[i]==DIV){
            if(arg1s[i] >= 0)
                fprintf(file,"D[%d] %s D[%d] / %s;\\n", arg1s[i],arg1sym,i,arg2str);
            if(arg2s[i] >= 0)
                fprintf(file,"D[%d] %s -D[%d] * %s / (%s*%s);\\n", arg2s[i],arg2sym,i,arg1str,arg2str,arg2str);
      } else if(ops[i]==SUB){
            if(arg1s[i] >= 0)
                fprintf(file,"D[%d] %s D[%d];\\n", arg1s[i],arg1sym,i);
            if(arg2s[i] >= 0)
                fprintf(file,"D[%d] %s -D[%d];\\n", arg2s[i],arg2sym,i);
      } else if(ops[i]==EXP){
            if(arg1s[i] >= 0)
                fprintf(file,"D[%d] %s D[%d] * A[%d];\\n", arg1s[i],arg1sym,i,i);
      } else if(ops[i]==LOG){
            if(arg1s[i] >= 0)
                fprintf(file,"D[%d] %s D[%d] / A[%d];\\n", arg1s[i],arg1sym,i,arg1s[i]);
      }

    }
    free(touched);
    fprintf(file,"}\\n");
    fclose(file);

    '''
    #print writercode
    inline(writercode,['ops','arg1s','arg2s','froz'], extra_compile_args=['-w'])

class etree:
    id = 0
    ops  = a.array('B')
    arg1 = a.array('l')
    arg2 = a.array('l')
    froz = a.array('f')
    #mark = a.array('b')

    def __init__(self,doinits=True):
        self.id   = etree.id
        etree.id += 1
        #self.mark = 0
        #self.op   = NOOP
        if doinits:
            etree.ops .append(NOOP)
            etree.arg1.append(NOID)
            etree.arg2.append(NOID)

    def __str__(self):
        return "etree"

    def __add__(self,other):
        new_etree = etree(False)
        etree.ops.append(ADD)
        etree.arg1.append(self.id)
        if hasattr(other,'id'):
            etree.arg2.append(other.id)
        else:
            etree.froz.append(other)
            etree.arg2.append(-len(etree.froz))
        return new_etree

    __radd__ = __add__

    def __mul__(self,other):
        new_etree = etree(False)        
        etree.ops .append(MULT)
        etree.arg1.append(self.id)
        if hasattr(other,'id'):
            etree.arg2.append(other.id)
        else:
            etree.froz.append(other)
            etree.arg2.append(-len(etree.froz))
        return new_etree

    __rmul__ = __mul__

    def __sub__(self,other):
        new_etree = etree(False)
        etree.ops .append(SUB)
        etree.arg1.append(self.id)
        if hasattr(other,'id'):
            etree.arg2.append(other.id)
        else:
            etree.froz.append(other)
            etree.arg2.append(-len(etree.froz))
        return new_etree

    def __rsub__(self,other):
        new_etree = etree(False)
        etree.ops.append(SUB)
        if hasattr(other,'id'):
            etree.arg1.append(other.id)
        else:
            etree.froz.append(float(other))
            etree.arg1.append(-len(etree.froz))
        etree.arg2.append(self.id)
        return new_etree        

    def __div__(self,other):
        new_etree = etree(False)
        etree.ops .append(DIV)
        etree.arg1.append(self.id)
        if hasattr(other,'id'):
            etree.arg2.append(other.id)
        else:
            etree.froz.append(other)
            etree.arg2.append(-len(etree.froz))
        return new_etree

    def __rdiv__(self,other):
        new_etree = etree(False)
        etree.ops.append(DIV)
        if hasattr(other,'id'):
            etree.arg1.append(other.id)
        else:
            etree.froz.append(other)
            etree.arg1.append(-len(etree.froz))
        etree.arg2.append(self.id)
        return new_etree

    def __neg__(self):
        return self*(-1)
    
    def exp(self):
        new_etree = etree(False)
        etree.ops .append(EXP)
        etree.arg1.append(self.id)
        etree.arg2.append(NOID)
        return new_etree

    def log(self):
        new_etree = etree(False)
        etree.ops .append(LOG)
        etree.arg1.append(self.id)
        etree.arg2.append(NOID)
        return new_etree

    def compile(self,outargs,*args):
        mattype = type(numpy.zeros(1))
        N = self.id+1

        graphsize = self.id+1
        # this will get increased to include frozen-in data
        # at the end

        #frozens = array(frozendata,dtype=float)

        print "generating init code"

        initcode  = '//printf("how are you my children\\n");\n'
        initcode += "const int graphsize =%d;\n" % graphsize
        initcode += "const int nops = graphsize\n;\n"
        initcode += "int c2d = sizeof(double)/sizeof(char);\n"
        initcode += "double* A = new double[graphsize];\n"
        initcode += "double* D = new double[graphsize];\n"
        #initcode += "for(int i=0; i<graphsize; i++){ A[i]=0.0; D[i]=0.001;}\n"
        initcode += "PyObject * cargs[%d], *coutargs[%d];\n" % (len(args), len(args))        
        initcode += "PyArrayObject * arr[%d], *oarr[%d];\n" % (len(args), len(args))
        initcode += "double* adata[%d], *odata[%d];\n" % (len(args), len(args))
        initcode += "int *astr[%d], *ostr[%d];\n" % (len(args),len(args))
        initcode += "PyArrayObject *array;\n"
        initcode += "int where;\n"

        for n in range(len(args)):
            initcode += "cargs[%d]    = PySequence_GetItem(args2,%d);\n" % (n, n)
            # this prevents memory leaks, but I have no idea what I am doing
            initcode += "Py_XDECREF(cargs[%d]);\n" % n # kill reference count!?
            initcode += "coutargs[%d] = PySequence_GetItem(outargs,%d);\n" % (n, n)
            # this prevents memory leaks, but I have no idea what I am doing
            initcode += "Py_XDECREF(coutargs[%d]);\n" % n # kill reference count!?

        for n in range(len(args)):
            if type(args[n])==mattype:
                initcode += "arr[%d]   = (PyArrayObject*)cargs[%d];\n" % (n,n)
                initcode += "oarr[%d]  = (PyArrayObject*)coutargs[%d];\n" % (n,n)
                # will access data as doubles, not chars
                initcode += "adata[%d] = (double*)arr[%d]->data;\n" % (n,n)
                initcode += "odata[%d] = (double*)oarr[%d]->data;\n" % (n,n)

                # copy strides into a new array, and compensate for being doubles instead of chars
                ndim = args[n].ndim
                initcode += "int astr%d[%d];\n" % (n,ndim)
                initcode += "astr[%d] = astr%d;\n" % (n,n)
                for dim in range(ndim):
                    initcode += "astr[%d][%d] =  arr[%d]->strides[%d]/c2d;\n" % (n,dim,n,dim)

        print "generating copy code"

        copycode = ""
        for i in range(len(args)):
            A = args[i]
            # this dynamically generates a for loop with the appropriate number of fors.
            # this was difficult to write!
            if type(A)==mattype:
                where0 = A.flat[0].id
                copycode += "where=%d;\n" % where0
                for dim in range(A.ndim):
                     copycode += "for(int i%d=0; i%d<%d; i%d++){\n" % (dim, dim, A.shape[dim], dim)
                copycode += "A[where++] = *(adata[%d]" % i
                for dim in range(A.ndim):
                    copycode += "+i%d*astr[%d][%d]" % (dim, i, dim)
                copycode += ");\n"
                for dim in range(A.ndim):
                    copycode += "}\n"
            else:
                copycode += "A[%d] = PyFloat_AS_DOUBLE(PyNumber_Float(cargs[%d]));\n" % (A.id,i)

        print "generating backcopy code"
                
        backcopycode = ""
        for i in range(len(args)):
            A = args[i]
            if type(A)==mattype:
                where0 = A.flat[0].id
                backcopycode += "where=%d;\n" % where0
                for dim in range(A.ndim):
                    backcopycode += "for(int i%d=0; i%d<%d; i%d++){\n" % (dim, dim, A.shape[dim], dim)
                backcopycode += "*(odata[%d]" % i
                for dim in range(A.ndim):
                    backcopycode += "+i%d*astr[%d][%d]" % (dim, i, dim)
                backcopycode += ")= D[where++];\n"
                for dim in range(A.ndim):
                    backcopycode += "}\n"
            else:
                backcopycode += "PySequence_SetItem(outargs,%d,PyFloat_FromDouble(D[%d]));\n" % (i,A.id)


        backcopycode += "return_val = A[graphsize-1];\n"

        backcopycode += "delete[] A;\n"
        backcopycode += "delete[] D;\n"

        initcode += copycode

        print "converting data to numpy format"

        # just skip this altogether??
        ops   = numpy.array(etree.ops ,dtype=numpy.int16)
        arg1s = numpy.array(etree.arg1,dtype=numpy.int32)
        arg2s = numpy.array(etree.arg2,dtype=numpy.int32)
        froz  = numpy.array(etree.froz,dtype=numpy.double)

        print "generating ccode"
        t1 = time.time()

        gencode(ops,arg1s,arg2s,froz)

        t2 = time.time()
        print "%f seconds to generate." % (t2-t1)
        print "compiling workhorse function."
        t1 = time.time()
        os.system("gcc -c mycode.c -ffast-math -O2") # 
        #os.system("tcc -c mycode.c")
        t2 = time.time()
        print "%f seconds to compile." % (t2-t1)

        allcode = initcode + "rock(A,D);\n" + backcopycode        
        #print allcode

        # copy the inputs once, then re-use them
        # (saves time!)

        return lambda *args2: evaluator(allcode,outargs,*args2)
        

def evaluator(allcode,outargs,*args2):
    # get types for testing later
    testmat  = numpy.zeros(1,dtype=float)
    mattype  = type(testmat)
    matdtype = testmat.dtype

    # outargs has to be a mutable type, otherwise you get a runtime error in the c code
    # also need to make a copy
    #print args2

    #outargs = copy.deepcopy(list(args2))
    #print outargs
    #outargs = copy(args2)
    #print outargs
    
    for i in range(len(args2)):
        if type(args2[i])==mattype:
            if args2[i].dtype != matdtype:
                raise Exception("etree: arg %d is non-float numpy array. Sage? do RealNumber=float" % i)
                                
    #print allcode
            
    import os
    ret_val = inline(allcode,['args2','outargs'],
                     extra_compile_args=['-O3 -w'],
                     type_converters=converters.default,
                     compiler='gcc',
                     extra_objects=["mycode.o"],
                     export_symbols = ["rock"],
                     headers=['"' + os.path.abspath('.') + '/mycode.h' + '"'],
                     auto_downcast=0)

    return [ret_val,outargs]

def runtests():
    "generate a bunch of etrees and compare the generated derivs \
     with finite difference derivs"
    import time
    random.seed(1)
    e = 1e-6

#     print "comparing convolution speed"
#     siz = [20,20]
#     A = random.random(siz)
#     B = random.random(siz)
#     C = random.randn(3,3)
#     D = random.randn(3,3)
#     def fun(X,Y):
#         rez = jlib.my_convolve2d(A,X+Y,mode='same')
#         return sum( (rez.flatten()-B.flatten())**2)
#     f_etree = compile_fun(fun,C,D)
#     rez = f_etree(C,D)

#     t1_etree  = time.time()
#     for i in range(100):
#         rez_etree = f_etree(C,D)
#     t2_etree  = time.time()

#     t1_numpy  = time.time()
#     for i in range(100):
#         rez_numpy = fun(C,D)
#     t2_numpy  = time.time()

#     print 'etree time vs. numpy time'
#     print (t2_etree-t1_etree), (t2_numpy-t1_numpy)
#     print 'etree results vs. numpy results'
#     print rez_etree[0], rez_numpy

#     print 'etree derivs'
#     print rez_etree[1][0]

#     raw_input('press enter to continue')

    siz = [2000]
    C = random.random(siz)
    def fun(X,Y):
        return sum(exp(X*Y).flatten())
        
    A = random.random(siz)
    B = random.random(siz)

    f_etree = compile_fun(fun,A,B)
    rez = f_etree(A,B)

    t1_etree  = time.time()
    for i in range(100):
        rez_etree = f_etree(A,B)
    t2_etree  = time.time()

    t1_numpy  = time.time()
    for i in range(100):
        rez_numpy = fun(A,B)
    t2_numpy  = time.time()

    print 'etree time vs. numpy time'
    print (t2_etree-t1_etree), (t2_numpy-t1_numpy)
    print 'etree results vs. numpy results'
    print rez_etree[0], rez_numpy

    raw_input('press enter to continue')

    print "now, test a bunch of derivative computations against finite differences."
    print "These tests print a lot of pairs."
    print "Each number should be very close to the other."
    e = 1e-6

    print "--basic operations, scalar arguments"
    a = .123
    b = .592
    def fun(x,y):
        return x+y*x-y/x
    f = compile_fun(fun,a,b)
    o ,[da ,db ] = f(a  ,b  )
    o2,[da2,db2] = f(a+e,b  )
    rez1 = [da, (1/e)*(o2-o)]
    o2,[da2,db2] = f(a,b+e  )
    rez2 =  [db, (1/e)*(o2-o)]
    print rez1
    print rez2

    reset()

    print "--some more advanced operations"
    a = .123
    b = .592
    def fun(x,y):
        return exp(x)+log(y*x)-y/x
    f = compile_fun(fun,a,b)
    o ,[da ,db ] = f(a  ,b  )
    o2,[da2,db2] = f(a+e,b  )
    rez3 = [da, (1/e)*(o2-o)]
    o2,[da2,db2] = f(a,b+e  )
    rez4 = [db, (1/e)*(o2-o)]
    print rez3
    print rez4

    reset()

    print "--vector test"
    A = numpy.random.rand(5)
    B = numpy.random.rand(5)
    def fun(X,Y):
        return (X*Y).sum() + (X/Y).sum() + (X+Y).prod()
    f = compile_fun(fun,A,B)
    o, [dA, dB] = f(A ,B )
    A2 = A.copy()
    A2[2] += e
    o2,[dA2,dB2]= f(A2,B)
    rez5 = [dA[2], (1/e)*(o2-o)]
    B2    = B.copy()
    B2[4] += e
    o2,[dA2,dB2]= f(A,B2)
    rez6 = [dB[4], (1/e)*(o2-o)]
    print rez5
    print rez6

    reset()

    print "--multi-dim test"
    A = numpy.random.rand(5,4,2)
    B = numpy.random.rand(5,4,2)
    def fun(X,Y):
        return (X*Y).sum()
    f = compile_fun(fun,A,B)
    o,[dA, dB] = f(A ,B )
    A2 = A.copy()
    A2[2,2,1] += e
    o2,[dA2,dB2]= f(A2,B)
    rez7 = [dA[2,2,1], (1/e)*(o2-o)]
    B2    = B.copy()
    B2[4,1,0] += e
    o2,[dA2,dB2]= f(A,B2)
    rez8 = [dB[4,1,0], (1/e)*(o2-o)]

    reset()

    print "--taking in floats simple"
    a = .123
    b = -.345
    def fun(x,y):
        return x*.2 + y*.5
    f = compile_fun(fun,a,b)
    o,[da, db] = f(a,b)
    a2  = a
    a2 += e
    o2,[da2,db2]= f(a2,b )
    rez9 = [da, (1/e)*(o2-o)]
    b2  = b
    b2 += e
    o2,[da2,db2]= f(a ,b2)
    rez10 = [db, (1/e)*(o2-o)]
    print rez9
    print rez10
    
    reset()

    print "--taking in floats complex"
    A = numpy.random.rand(5)
    Y = numpy.random.rand(5)
    def fun(X):
        return (X*Y).sum()
    f = compile_fun(fun,A)

    o,[dA] = f(A)
    A2 = A.copy()
    A2[2] += e
    o2,[dA2]= f(A2)
    rez11 = [dA[2], (1/e)*(o2-o)]

    print "--basic operations, scalar arguments"
    print rez1
    print rez2
    print "--some more advanced operations"
    print rez3
    print rez4
    print "--vector test"
    print rez5
    print rez6
    print "--multi-dim test"
    print rez7
    print rez8
    print "--taking in floats simple"
    print rez9
    print rez10
    print "--taking in floats complex"
    print rez11

def example():
    def fun(A,B):
        rez = 0
        for i in range(A.shape[0]):
            for j in range(B.shape[1]):
                rez += dot(A[i,:],B[:,j])
        return rez

    A = numpy.random.rand(60,60)
    B = numpy.random.rand(60,60)
    dfun = compile_fun(fun,A,B)
    F,[dFdA,dFdB] = dfun(A,B)

    print F, dFdA, dFdB

    return fun, dfun