#!/usr/bin/python2.6
#
# Werner Lustermann
# ETH Zurich
#
from ctypes import *

# get the ROOT stuff + my shared libs
from ROOT import gSystem
gSystem.Load('../troot/fitslib.so')
from ROOT import *

import numpy as np


class rawdata( object ):
    """
    raw data access and calibration
    """
    def __init__( self, dfname,  calfname ):
        """
        open data file and calibration data file
        get basic information about the data inf dfname
        allocate buffers for data access

        dfname   - fits or fits.gz file containing the data including the path
        calfname - fits or fits.gz file containing DRS calibration data
        """
        self.dfname   = dfname
        self.calfname = calfname
        
        # access data file
        try:
            df = fits( self.dfname )
        except IOError:
            print 'problem accessing data file: ', dfname
            raise # stop ! no data
        self.df = df
        
        # get basic information about the data file
        self.NROI    = df.GetUInt( 'NROI' ) # region of interest (length of DRS pipeline read out)
        self.NPIX    = df.GetUInt( 'NPIX' ) # number of pixels (should be 1440)
        self.NEvents = df.GetNumRows()      # find number of events
        # allocate the data memories
        self.evNum = c_ulong()
        self.Data  = np.zeros( self.NPIX * self.NROI, np.int16 ) 
        self.startCells = np.zeros( self.NPIX, np.int16 )
        # set the pointers to the data++
        df.SetPtrAddress( 'EventNum', self.evNum )
        df.SetPtrAddress( 'StartCellData', self.startCells ) # DRS readout start cell
        df.SetPtrAddress( 'Data', self.Data ) # this is what you would expect
        # df.GetNextRow() # access the first event
        
        # access calibration file
        try:
            calf = fits( self.calfname )
        except IOError:
            print 'problem accessing calibration file: ', calfname
            raise
        self.calf = calf
        #
        BaselineMean      = calf.GetN('BaselineMean')
        GainMean          = calf.GetN('GainMean')
        TriggerOffsetMean = calf.GetN('TriggerOffsetMean')

        self.blm = np.zeros( BaselineMean, np.float32 )
        self.gm  = np.zeros( GainMean, np.float32 )
        self.tom = np.zeros( TriggerOffsetMean, np.float32 )

        self.Nblm = BaselineMean / self.NPIX
        self.Ngm  = GainMean / self.NPIX
        self.Ntom  = TriggerOffsetMean / self.NPIX

        calf.SetPtrAddress( 'BaselineMean', self.blm )
        calf.SetPtrAddress( 'GainMean', self.gm )
        calf.SetPtrAddress( 'TriggerOffsetMean', self.tom )
        calf.GetRow(0)

        self.v_bsl = np.zeros( self.NPIX ) # array with baseline values (all ZERO)

    def next( self ):
        """
        load the next event from disk and calibrate it
        """
        self.df.GetNextRow()
        self.calibrate_drsAmplitude()

        
    def calibrate_drsAmplitude( self ):
        """
        perform amplitude calibration for the event 
        """
        tomV = 2000./4096.
        acalData = self.Data * tomV # convert into mV

        # reshape arrays: row = pixel, col = drs_slice
        acalData = np.reshape( acalData, (self.NPIX, self.NROI) )
        blm = np.reshape( self.blm, (self.NPIX, self.NROI) )
        tom = np.reshape( self.tom, (self.NPIX, self.NROI) )
        gm  = np.reshape( self.gm,  (self.NPIX, self.NROI) )
        
        # print 'acal Data ', acalData.shape
        # print 'blm shape ', blm.shape
        # print 'gm shape  ', gm.shape
        
        for pixel in range( self.NPIX ):
            # rotate the pixel baseline mean to the Data startCell
            blm_pixel = np.roll( blm[pixel,:], -self.startCells[pixel] )
            acalData[pixel,:] -= blm_pixel[0:self.NROI]
            acalData[pixel,:] -= tom[pixel, 0:self.NROI]
            acalData[pixel,:] /= gm[pixel,  0:self.NROI]
            
	self.acalData = acalData * 1907.35
    
        # print 'acalData ', self.acalData[0:2,0:20]

    def ReadBaseline( self, file, bsl_hist = 'bsl_sum/hplt_mean' ):
        """
        open ROOT file with baseline histogram and read baseline values
        file       name of the root file
        bsl_hist   path to the histogram containing the basline values
        """
        try:
            f = TFile( file )
        except:
            print 'Baseline data file could not be read: ', file
            return
        
        h = f.Get( bsl_hist )

        for i in range( self.NPIX ):
            self.v_bsl[i] = h.GetBinContent( i+1 )

        f.Close()

        
    def CorrectBaseline( self ):
        """
        apply baseline correction
        """
        for pixel in range( self.NPIX ):
            self.acalData[pixel,:] -= self.v_bsl[pixel]
            
        
    def info( self ):
        """
        print information
        """
        print 'data file:  ', dfname
        print 'calib file: ', calfname
        print '\ncalibration file'
        print 'N BaselineMean: ', self.Nblm
        print 'N GainMean: ', self.Ngm
        print 'N TriggeroffsetMean: ', self.Ntom
        

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 a la ROOT of the histograms

    def add( self, tag, h ):
        self.list.append( h )
        self.dict[tag] = h
        self.hList.Add( h )


class pixelHisto1d ( object ):

    def __init__( self, name, title, Nbin, first, last, xtitle, ytitle, NPIX ):
        """
        book one dimensional histograms for each pixel
        """
        self.name = name

        self.list = [ x for x in range( NPIX ) ]
        self.hList = TObjArray()

        for pixel in range( NPIX ):

            hname  = name + ' ' + str( pixel )
            htitle = title + ' ' + str( pixel )
            self.list[pixel] = TH1F( hname, htitle, Nbin, first, last )

            self.list[pixel].GetXaxis().SetTitle( xtitle )
            self.list[pixel].GetYaxis().SetTitle( ytitle )
            self.hList.Add( self.list[pixel] )


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()

# simple test method
if __name__ == '__main__':
    """
    create an instance
    """
    dfname = '/data03/fact-construction/raw/2011/11/24/20111124_121.fits'
    calfname = '/data03/fact-construction/raw/2011/11/24/20111124_111.drs.fits'
    rd = rawdata( dfname, calfname )
    rd.info()
    rd.next()
    
# for i in range(10):
#    df.GetNextRow() 

#    print 'evNum: ', evNum.value
#    print 'startCells[0:9]: ', startCells[0:9]
#    print 'evData[0:9]: ', evData[0:9]