#!/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, TPolyLine, 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)

canvas = TCanvas()
canvas.Divide(2,2)

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
    keepInMemory = [] #Avoid premature garbage collection and overwriting

    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)
            
            #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 )
            keepInMemory.append(barChart)
            if layer == 0:
                barChart.SetFillColor(ROOT.kBlue)
                barChart.GetXaxis().SetTitle("X [strips]")
                sideName = "(N-side)"
            elif layer == 1:
                barChart.SetFillColor(ROOT.kRed)
                barChart.GetXaxis().SetTitle("Y [strips]")
                sideName = "(P-side)"
            barChart.GetYaxis().SetTitle("ADU count")
            barChart.SetTitle("Digitized BAT event, Event#=" + str(event.eventNum) + ", " + name +  " " + sideName)
            barChart.GetYaxis().SetRangeUser(-5, 2000)
            barChart.Draw('AB')

            #Draw detector
            pad = canvas.cd(3+layer)
            pad.Clear()

            #pad.GetCanvasPainter().SetFillColor(38)
            boxStart = 0.1; boxStop = 0.9; boxWidth = boxStop-boxStart
            boxX = array('d', (boxStart, boxStop, boxStop, boxStart, boxStart))
            boxY = array('d', (0.1, 0.1, 0.9, 0.9, 0.1))
            box = TPolyLine(5,boxX, boxY)
            box.SetFillColor(38)
            box.Draw("f")
            keepInMemory.append(box)
            
            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.PaintTextNDC(smid, 0.875, str(sn))
                boxX = array('d', (smid-swidth/2, smid+swidth/2, smid+swidth/2, smid-swidth/2, smid-swidth/2))
                boxY = array('d', (0.85, 0.85, 0.9, 0.9, 0.85))
                box = TPolyLine(5,boxX, boxY)
                if layer == 0:
                    box.SetFillColor(ROOT.kBlue)
                elif layer == 1:
                    box.SetFillColor(ROOT.kRed)
                box.Draw("f")
                keepInMemory.append(box)

            #Draw energy deposits
            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

                size = hit.edep/MeV*250
                if x > boxStart and x < boxStop:
                    mark = TMarker(x,y, 20)
                    mark.SetMarkerSize(size)
                    mark.Draw()
                    keepInMemory.append(mark)
                elif x < boxStart:
                    pass
#                     boxX = array('d', (boxStart-size, boxStart, boxStart, boxStart-size, boxStart-size))
#                     boxY = array('d', (y-hs, y-hs, y+hs, y+hs, y-hs))
#                     box = TPolyLine(5,boxX, boxY)
#                     box.SetFillColor(ROOT.kBlack)
#                     box.Draw("f")
#                     keepInMemory.append(box)
                elif x > boxStop:
                    pass
#                     boxX = array('d', (boxStop, boxStop+size, boxStop+size, boxStop, boxStop))
#                     boxY = array('d', (y-hs, y-hs, y+hs, y+hs, y-hs))
#                     box = TPolyLine(5,boxX, boxY)
#                     box.SetFillColor(ROOT.kBlack)
#                     box.Draw("f")
#                     keepInMemory.append(box)
            
        while True:
            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")
            else:
                break
        keepInMemory = [] #Drop unneeded stuff