#===============================================================================
# Full Numpy integration for aanet. 
#===============================================================================

import numpy as np
import collections
import ROOT

def get_array( tree, varexp, condition = "" , neve=1000000000):

    """ This function is like TTree:Draw, except that it returns numpy
        arrays. in this example, the numpy array holds a pointer to the
        data, which means it is only valid until the next TTree:Draw 
        commond. (of course it's trivial to copy them).
          
        If you 'draw' more than one variable, you get a list of numpy
        arrays.
    """

    # It would be interesting to see if we can get root to give up
    # ownership and let numpy assume ownership...

    # The number of rows may be (much) more than the number of events
    # in the tree. E.g. when plotting hits. If the arrays are not
    # large enough we only get the last events. We cannot predict this,
    # but we can detect it if it happens, and then dimension the arrays
    # correctly and (unfortunately) Draw one more time.
    
    tree.SetEstimate( tree.GetEntries() )

    while True : # sometimes, I really want a goto statement!
        tree.Draw(varexp,condition,"goff",neve  )
        n = tree.GetSelectedRows()
        if n > tree.GetEstimate() : 
            tree.SetEstimate(n) 
        else : break
    
    m = tree.GetPlayer().GetDimension()

    r = []
    
    for im in range(m):
        p = tree.GetVal(im)
        a = np.ndarray( (n,) ,dtype= np.float64, buffer=p )
        r.append(a)
        
    if len(r) == 0 : return None
    if len(r) == 1 : return r[0]
    return r

# Monkey-patch root TTree
ROOT.TTree.DrawNumpy = get_array



#===============================================================================
# Full Numpy integration for aanet. 
#===============================================================================


nptypes = { 
"int"          : np.int32  ,
"long int"     : np.int64  ,
"unsigned int" : np.uint32 ,
"unsigned long": np.uint64 ,
"double"       : np.float64,
"size_t"       : np.uint32 }

skiptypes = ["TClass","atomic<TClass*>", "string" ]
    
def get_dtype( rootclass , verb = False ):

    D = collections.defaultdict( list )

    try :
        L = rootclass.GetListOfDataMembers()
    except AttributeError:

        rootclass = rootclass.Class()
        L = rootclass.GetListOfDataMembers()
        # maybe this works...
        
        
    for m in L :

        title = None
        
        if verb : print " datamember : ", m.GetTypeName(), m.GetName(), m.GetUnitSize()

        
        if m.IsBasic() :
            dtype = nptypes[ m.GetTypeName() ] # todo: check m.GetUnitSize()

            
        elif m.GetTypeName().startswith("vector<") :

            value_type_str = m.GetTypeName().split("<",1)[1].rsplit(">",1)[0]

            if value_type_str not in nptypes :
                print " cannot handle vector of ", value_type_str
                continue

            value_type_np = nptypes[ value_type_str ]
                
            # record a void dtype and store the type in the title
            # for later rebuilding.
            
            dtype  = np.dtype( (np.uint8, m.GetUnitSize() ) )
            title  = value_type_np

            
        else :
            n = m.GetTypeName()

            if n in skiptypes : continue
            
            rootclass_of_member = ROOT.TClass.GetClass( n )
            dtype  = get_dtype ( rootclass_of_member )

        if verb : print "-->",m.GetName(), m.GetTypeName(),  dtype, title
            
        D['names'].  append( m.GetName()   ) 
        D['titles']. append( title         )
        D['formats'].append( dtype         )
        D['offsets'].append( m.GetOffset() )
  
    D['aligned'] = False 
    D['itemsize']= rootclass.Size()
    
    # print D
    
    return np.dtype( D )


class AArray(np.recarray):

    # https://docs.scipy.org/doc/numpy-1.13.0/user/basics.subclassing.html

    # I actualy forgot why I made my own class for this.
    
    def __new__(subtype, shape, dtype, buf=None, offset=0,
                strides=None, order=None, info=None):

        obj = super(AArray, subtype).__new__(subtype, shape, dtype,
                                             buf, offset, strides,
                                             order)
        return obj
    
    def __array_finalize__(self, obj):

        self.info = "whathere?"


        
def npfy( V ): # V should be root-dictionaried stl vector

    """ This function takes a std::vector (which must have root 
        type-information), and produces a numpy array """
    
    if V.size() == 0 : return None
    
    dtyp = get_dtype ( V[0].Class() )
    buf  = ROOT._avoid( V[0] ) # void ptr becomes python buffer
    buf.SetSize( dtyp.itemsize * len(V) )
    
    # Offsets are baked into the dtype, so no offsets here.
    # Note that the buffer arugment is called buffer in ndarray, but
    # buf in recarray (!)

    return np.recarray( (len(V),), dtype=dtyp , buf=buf ) 
    #return AArray(      (len(V),), dtype=dtyp , buf=buf )