#!/usr/bin/python
#
# Werner Lustermann
# ETH Zurich
#

from ROOT import TH1F, TObjArray

import numpy as np


class histogramList( object ):
    """
    """
    def __init__( self, name ):
        """ set the name and create empty lists """
        self.name  = name         # name of the list
        self.list  = []           # list of the histograms
        self.dict  = {}           # dictionary of histograms
        self.hList = TObjArray()  # list of histograms for ROOT

        self.book()

    def __del__(self):
        for h in self.list:
            h.Delete()
        
    def add( self, tag, h ):
        self.list.append( h )
        self.dict[tag] = h
        self.hList.Add( h )
        
    def h1d( self, name, title, Nbin, first, last, xtitle, ytitle ):

        h = TH1F( name, title, Nbin, first, last );
        h.GetXaxis().SetTitle( xtitle );
        h.GetYaxis().SetTitle( ytitle );

        # print 'self.add( ',name, ',', h, ' )'
        self.add( name, h )
        
    def book( self ):
        # print 'histogramList::book'
        pass
    

class hist( object ):
    """
    Use numpy histogram function for histogram filling
    Allows multiple calls for a  single histogram
    Create ROOT histo 
    """
    def __init__( self, nbins, first, last, name = 'name', title = 'title',
                  xtitle = 'x', ytitle = 'y' ):

        self.nbins  = nbins
        self.first  = first
        self.last   = last
        self.name   = name
        self.title  = title
        self.xtitle = xtitle
        self.ytitle = ytitle
        self.y = np.zeros( self.nbins ) # initialize all bins to ZERO

        self.rh = 0
        
    def fill( self, x ):
        h, b = np.histogram( x, self.nbins, self.first, self.last,
                             new = True )
        self.y += h

    def zero( self ):
        self.y[:] = 0.
        
    def mkroot( self ):
        
        if self.rh == 0:
            self.rh = TH1F( self.name, self.title, self.nbins,
                   self.first, self.last )
            self.rh.GetXaxis().SetTitle( xtitle );
            self.rh.GetYaxis().SetTitle( ytitle );

        for i in range( self.nbins ):
            self.rh.SetBinContent( i+1, self.y[i] )
        
        
class hist_array( object ):
    """
    """
    def __init__( self, nhist, nbins, first, last, name = 'name',
                  title = 'title', xtitle = 'x', ytitle = 'y', root = True):
        """
        - store the histogram info in the object
        - initialize data array
        """
        self.nhist  = nhist 
        self.nbins  = nbins
        self.first  = first
        self.last   = last
        self.name   = name
        self.title  = title
        self.xtitle = xtitle
        self.ytitle = ytitle
        
        t = np.zeros( self.nhist * self.nbins ) # initialize array ZERO
        self.y = t.reshape( nhist, nbins ) # shape array of histograms

        if root == True:
            self.bookRootHistos()
        else:
            self.list = 0

    def __del__(self):
        for h in self.list:
            h.Delete()
            
    def fill( self, hist, x ):
        h, b = np.histogram( x, self.nbins, ( self.first, self.last ),
                             new = True )
        self.y[ hist ] += h

    def fillall( self, x ):
        for hist in range( self.nhist ):
            self.fill( hist, x[hist, :] )

    def zero( self ):
        self.y.flat[:] = 0.
        
    def bookRootHistos( self ):

        self.list  = [ x for x in range( self.nhist ) ]
        self.hList = TObjArray()
        
        for hist in range( self.nhist ):

            hname  = self.name + ' ' + str( hist )
            htitle = self.title + ' ' + str( hist )
            self.list[hist] = TH1F( hname, htitle,
                                    self.nbins, self.first, self.last )
            
            self.list[hist].GetXaxis().SetTitle( self.xtitle )
            self.list[hist].GetYaxis().SetTitle( self.ytitle )
            self.hList.Add( self.list[hist] )

    def SetBinContent( self ):

        if self.list == 0:
            print 'no root histograms booked'
        else:
            for hist in range( self.nhist ):        
                for i in range( self.nbins ):
                    self.list[hist].SetBinContent( i+1, self.y[hist, i] )
		    
		    
class hist_list( object ):

	def __init__( self, booking, nhist, hname ):
	
		self.list = []
	
		for i in range(nhist):
			self.list.append( booking( str( hname )+"_"+str(i+1)))