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

from ROOT import TH1F, TFile, TObjArray
import numpy as np


class histogramList( object ):
    """ maintain a list of ROOT histograms
    """
    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 ):
        """ booking of histograms

        to be overwritten by the user with hist own definitions
        """
        # 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 ):
    """ array of 1d histograms data

    y : np.array( nhist, nbins ), each row is representing a 1d histogram

    all histograms have identical binning
    histogramming of data vectors is done using np.histogram

    creating of ROOT histograms is optional
    """
    def __init__( self, nhist, nbins, first, last, name = 'name',
                  title = 'title', xtitle = 'x', ytitle = 'y', root = True):
        """ store histogram info and initialize data array
        nhist  : number of histograms
        nbins  : number of bins in a histogram
        first  : first bin
        last   : last bin
        name   : name of the histogram
        title  : title of the histogram
        xtitle : label for the x axe
        ytitle : label for the y axe

        y : np.array( nhist, nbins ), each row is representing a 1d histogram
        """
        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

        # creating of ROOT histograms
        if root == True:
            self.bookRootHistos()
        else:
            self.list = 0

    def __del__(self):
        """ delete the ROOT histogram objects """
        
        for h in self.list:
            h.Delete()
            
    def fill( self, hist, x ):
        """ compute histogram for new data and add them to existing histogram
        
        hist : index of the histogram { 0:nhist }
        x    : data vector
        """
        
        h, b = np.histogram( x, self.nbins, ( self.first, self.last ),
                             new = True )
        self.y[ hist ] += h

    def fillall( self, x ):
        """ compute all histograms for now data array and add them to existing histograms

        x : np.array( nhist, nbins ), data array in the shape of the histogram array
        """
        for hist in range( self.nhist ):
            self.fill( hist, x[hist, :] )

    def zero( self ):
        """ set histogram array to ZERO
        """
        self.y.flat[:] = 0.
        
    def bookRootHistos( self ):
        """ book ROOT histograms equivalent to the histogram array

        self:
            list  : list containing the ROOT histogram objects
            hList : ROOT TObjArray containing all histograms
        """
        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 ):
        """ set the ROOT contents
        """
        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] )
        
        
def SaveHistograms( histogramLists, fname = 'histo.root', opt = 'RECREATE' ):
    """
    Saves all histograms in all given histogram lists to a root file
    Each histogram list is saved to a separate directory
    """
    rf = TFile( fname, opt)
    
    for list in histogramLists:
        rf.mkdir( list.name )
        rf.cd( list.name )
        list.hList.Write()

    rf.Close()