#!/usr/bin/env python

"""
This script makes per-event histograms/figures
showing the energy depositions + digizer responce

Inspired by similar plots from PAMELA experiment
(see for example 'Spatial Resolution of double-sided
silicon microstrip detector for the PAMELA apparatus')

Kyrre N. Sjobak (k.n.sjobak@fys.uio.no)
21/2-2010
"""


#ALL units in MeV, mm, us
MeV = 1
GeV = 1000

import sys, os

#Find working folder, set maxEvents
skipToEvent = -1
if not (len(sys.argv) == 2 or len(sys.argv) == 3):
    print "Usage: python", sys.argv[0], " <workdir> { <skipToEvent> }"
    sys.exit(1)
if len(sys.argv) == 3:
    skipToEvent = int(sys.argv[2])


# Setup the modules for reading
from TbModules import *
from TbDigitizers import *
from TbDigitizerConfig import *
mls = TbModuleList(sys.argv[1])
#mls.modList['performance'] = TbModulePerformance("TBsim_performance.dat")
#trigger1 = mls.modList['trigger1']    = TbModuleScint("TBsim_moduleScint_trigger1.dat")
#trigger2 = mls.modList['trigger2']    = TbModuleScint("TBsim_moduleScint_trigger2.dat")
#veto     = mls.modList['veto']        = TbModuleScint("TBsim_moduleScint_veto.dat")
BAT1     = mls.modList['BAT1']        = TbModulePixel("TBsim_moduleBAT_BatMod1.dat")
BAT3     = mls.modList['BAT3']        = TbModulePixel("TBsim_moduleBAT_BatMod3.dat")
BAT6     = mls.modList['BAT6']        = TbModulePixel("TBsim_moduleBAT_BatMod6.dat")
#DUT1     = mls.modList['DUT1']        = TbModulePixel("TBsim_moduleDUT_DUT1.dat")
#DUT2     = mls.modList['DUT2']        = TbModulePixel("TBsim_moduleDUT_DUT2.dat")
#DUT3     = mls.modList['DUT3']        = TbModulePixel("TBsim_moduleDUT_DUT3.dat")
#DUT4     = mls.modList['DUT4']        = TbModulePixel("TBsim_moduleDUT_DUT4.dat")
truth    = mls.modList['truth']       = TbModuleTruth("TBsim_truth.dat")

posIndex = TbDigitizerConfigBAT.posIndex

#ROOT
import ROOT
from ROOT import gStyle, TCanvas, TGraph #, TMarker, TPad, TVirtualPadPainter
from array import array
import numpy

gStyle.SetOptStat("")
gStyle.SetCanvasColor(ROOT.kWhite)
gStyle.SetCanvasBorderMode(0)
gStyle.SetLabelColor(ROOT.kWhite)
#gStyle.SetFillColor(0)

canvasArray = (TCanvas("c_N","N-side graphs",640,900),TCanvas("c_P","P-side graphs",640,900))
[c.Divide(1,2) for c in canvasArray]
#canvas = TCanvas()
#canvas.Divide(2,1)

if not os.path.isdir("clusterBarchartPlots"): #Creating mls => implicit change of directory
    os.mkdir("clusterBarchartPlots")
print "Hit \"S\" to save canvases to file in subfolder \"clusterBarchartPlots\""

#Loop!
currentEvent = 0
while mls.ReadNextEvent():
    if skipToEvent >= 0 and currentEvent < skipToEvent:
        currentEvent +=1
        continue
    currentEvent += 1

    currentlyOnDisplay = ["",""] #Used to make filename if the user ask for saving

    BATevents = [BAT1.currentEvent, BAT3.currentEvent, BAT6.currentEvent]
    BATnames  = ["BAT1", "BAT3", "BAT6"]
    for event, name in zip(BATevents, BATnames):
        #digitizer = event.digitizers['BAT'] = TbDigitizerBAT(event)
        #digitizer = event.digitizers['BAT'] = TbDigitizerBAT_v2(event)
        digitizer = event.digitizers['BAT'] = TbDigitizerBAT_v3(event)
        for layer in xrange(2):
            #pad = canvas.cd(1+layer)
            canvas = canvasArray[layer]
            canvas.cd(1)
            
            #Check the range of the event
            firstFire = TbDigitizerConfigBAT.numStrips+1
            lastFire  = -1
            for dig in digitizer.digits:
                if isinstance(dig, TbDigitBAT_CMC):
                    continue
                if int(dig.isPside) == layer:
                    if dig.stripNum > lastFire:
                        lastFire = dig.stripNum
                    if dig.stripNum < firstFire:
                        firstFire = dig.stripNum

            
            #Create a graph
            def getStrip(strip):
                for dig in digitizer.digits:
                    if isinstance(dig, TbDigitBAT_CMC):
                        continue
                    if int(dig.isPside) == layer and dig.stripNum == strip:
                        return dig
                    
            strips  = array('i', range(firstFire, lastFire+1))
            try:
                ADU     = array('i', [getStrip(s).ADU for s in strips])
            except AttributeError:
                #Two hits: Skip this event!
                print "MultiHit"
                continue
            
            numHits = len(strips)
            if numHits < 1:
                print "NoHit"
                continue
            
            currentlyOnDisplay[layer] = "event" + str(event.eventNum) + name
            
            barChart = TGraph ( len(strips), strips, ADU )
            #print barChart.GetFillColor(), barChart.GetFillStyle(), barChart.GetLineColor(), barChart.GetLineStyle(), barChart.GetLineWidth(), barChart.GetMarkerColor()
            if layer == 0:
                barChart.SetFillColor(ROOT.kBlue)
                barChart.GetXaxis().SetTitle("X [strips]")
                #barChart.SetTitle("barChart_N")
                sideName = "(N-side)"
            elif layer == 1:
                barChart.SetFillColor(ROOT.kRed)
                barChart.GetXaxis().SetTitle("Y [strips]")
                #barChart.SetTitle("barChart_P")
                sideName = "(P-side)"
            #barChart.SetLineColor(ROOT.kBlack)
            barChart.GetYaxis().SetTitle("ADU count")
            barChart.SetTitle("Digitized BAT event, Event#=" + str(event.eventNum) + ", " + name +  " " + sideName)
            
            barChart.GetYaxis().SetRangeUser(-5, 2000)

            barChart.Draw('AB')
            canvas.Update()
            
            #Draw detector
            #pad = canvas.cd(3+layer)
            #canvas = canvasArray[layer]
            canvas.cd(2)
            canvas.GetCanvasPainter().SetFillColor(38)
            boxStart = 0.1; boxStop = 0.9; boxWidth = boxStop-boxStart
            canvas.PaintBox(boxStart,0.1, boxStop,0.9)

            if layer == 0:
                canvas.GetCanvasPainter().SetFillColor(ROOT.kBlue)
            elif layer == 1:
                canvas.GetCanvasPainter().SetFillColor(ROOT.kRed)

            occup = 0.7                             #Occupancy
            swidth = occup*boxWidth/numHits         #Width of one strip
            sdist  = (1 - occup)*boxWidth/numHits/2 #Remaining space on the side of each strip
            sPosFirst = boxStart + sdist + swidth/2
            sPosLast  = boxStop  - sdist - swidth/2
            
            for smid, sn in zip(numpy.linspace(sPosFirst, sPosLast, numHits, True), strips):
                canvas.PaintBox(smid-swidth/2,0.85, smid+swidth/2,0.9)
                #canvas.PaintTextNDC(smid, 0.875, str(sn))

            #Draw energy deposits
            canvas.GetCanvasPainter().SetFillColor(ROOT.kBlack)
            for hit in event.hits:
                y = 0.0 #Direction of depth in sensor
                x = 0.0 #Direction perpendicular to strips
                if layer == 0:
                    y = (-hit.posLocal[2]+TbDigitizerConfigBAT.d/2)/TbDigitizerConfigBAT.d*0.8+0.1 #Z-direction
                    x = hit.posLocal[posIndex[layer]]
                elif layer == 1:
                    y = ( hit.posLocal[2]+TbDigitizerConfigBAT.d/2)/TbDigitizerConfigBAT.d*0.8+0.1 #Z-direction
                    x = hit.posLocal[posIndex[layer]]
                x = x/0.050 + 640/2 - strips[0]   #Number in units of strips, 0 at edge of first strip in cluster
                x = x/numHits*boxWidth + boxStart #Transformed to NDC

                #mark = TMarker(x,y,20)
                #mark.Draw()
                #canvas.PaintPolyMarker(1, array('d',[x]), array('d',[y]))
                size = hit.edep/MeV*0.6 #1 MeV => 0.35x0.35 box
                hs = size/2
                if x > boxStart and x < boxStop:
                    canvas.PaintBox(x-hs,y-hs, x+hs,y+hs)
                elif x < boxStart:
                    canvas.PaintBox(boxStart-size,y-hs, boxStart,y+hs)
                elif x > boxStop:
                    canvas.PaintBox(boxStop,y-hs, boxStop+size,y+hs)
                #print x,y
                #print canvas.GetUxmin(), canvas.GetUymin(), canvas.GetUxmax(), canvas.GetUymax()
                
            canvas.Update()
        #canvas.Update()
        inp = raw_input()
        if inp == "s" or inp == "S":
            canvasArray[0].SaveAs("clusterBarchartPlots" + os.path.sep + \
                                  currentlyOnDisplay[0] + "N.png")
            canvasArray[1].SaveAs("clusterBarchartPlots" + os.path.sep + \
                                  currentlyOnDisplay[1] + "P.png")
        
        for canvas in canvasArray:
            canvas.Clear()
            canvas.Divide(1,2)

        #canvas.Clear()
        #canvas.Divide(2,2)