import ROOT,argparse
parser=argparse.ArgumentParser(description='plot samples of a digi and some other stuff')
parser.add_argument('recofile',help='the recofile with cluster tree and digi and sample persistence',type=str)
parser.add_argument('--mc',help='also do some stuff with monte carlo info',action='store_true')
parser.add_argument('--mcfile',help='the mc file',type=str,default='')
parser.add_argument('--digifile',help='the raw file',type=str,default='')
parser.add_argument('--title',help='add a title to everything',type=str,default='')
parser.add_argument('--hlp',help='show help',action='store_true')
parser.add_argument('--ev',help='go to event',type=int, default=-1)
parser.add_argument('--cal',help='do calib only',action='store_true')
parser.add_argument('--nevts',help='number of events to process',type=int,default=-1)
parser.add_argument('--compare',help='do comparison between 2 digi trees',action='store_true')
parser.add_argument('--sampleFreq',help='sampling frequency',type=float, default=15.55)
parser.add_argument('--sizecut',help='only plot digis bigger than size',type=int,default=[-1,-1],nargs=2)
parser.add_argument('--caltimecut',help='cut on caltime',type=int,default=[-1,-1],nargs=2)
parser.add_argument('--timecut',help='cut on sample time',type=int,default=[-1,-1],nargs=2)
parser.add_argument('--padcut',help='cut on a particular pad',default=-1,type=int)
parser.add_argument('--ampcut',help='add an digi amplitude cut',default=[-1,-1],type=int,nargs=2)
parser.add_argument('--trk',help='use olny digis in track',action='store_true')
parser.add_argument('--outf',help='name of the rootfile to write out',type=str,default='Null')
parser.add_argument('--timebits',help='number of timebits',type=int,default=9)
parser.add_argument('--timebins',help='number of bins in time',type=int,default=512)
parser.add_argument('--signumcut',help='cut on the signal number',type=int,default=[-1,-1],nargs=2)
parser.add_argument('--startampcut',help='cut on the first sample amp',type=int,default=-1)
parser.add_argument('--verbose',help='verbose cut information',action='store_true')
args=parser.parse_args()

if args.hlp:
    parser.show_help()
    exit(0)

ROOT.gROOT.ProcessLine(".x rootlogon.C") 
ROOT.gROOT.ProcessLine("gStyle->SetPalette(1)") 
ROOT.gROOT.LoadMacro("stlPYROOT.h+")

Infile=ROOT.TFile(args.recofile,"read")
tree=Infile.Get("cbmsim")
if tree==None:
    print "no cbmsim"
    Infile.Close()
    exit()
if args.mc:
    if args.mcfile=="":
        mcfile=args.recofile.replace("reco","mc")
        mcfile=mcfile.replace("_recl","")
    else:
        mcfile=args.mcfile
    if args.digifile=="":
        digifile=args.recofile.replace("reco","raw")
        digifile=digifile.replace("_recl","")
    else:
        digifile=args.digifile

    print 'adding',mcfile,'and',digifile,'as friends'
    tree.AddFriend("cbmsim",mcfile)
    tree.AddFriend("cbmsim",digifile)

tree.SetBranchStatus("*",1)
tree.SetBranchStatus("TpcSample.*",1)
tree.SetBranchStatus("TpcDigi.*",1)
tree.SetBranchStatus("TpcSignal.*",1)
if args.compare:
    tree.SetBranchStatus("TpcSample2.*",1)
    tree.SetBranchStatus("TpcDigi2.*",1)
tree.SetBranchStatus("TpcCluster.*",1)
tree.SetBranchStatus("TrackPostFit.*",1)


evcounter=-1
hsignal=ROOT.TH1D('hsignal','signal',30000,0,30000)
hsample=ROOT.TH1D('hsamples','Samples of Digi',200,0,200)
hdigi=ROOT.TH1D('hdigi','Digi',200,0,200)

hsample2=ROOT.TH1D('hsamples2','Samples of Digi2',50,0,50)
hdigi2=ROOT.TH1D('hdigi2','Digi2',50,0,50)

tbin=args.timebins
tmax=2**args.timebits
hsampleint=ROOT.TH1D('hsampleint','sample in amp norm (digi)',220,0,220)
hsampleintvst=ROOT.TH2D('hsampleintvst','sample int amp norm (digi) vs t',tbin,0,tmax,220,0,220)
htimecalib=ROOT.TH1D('htimecalib','time calib',120,-20,100)
htimecalibcalib=ROOT.TH1D('htimecalibcalib','time calib calib',480,-20,100)
htimecalibvsdigisize=ROOT.TH2D('htimecalibvsdigisize','htimecalibvsdigisize',120,-20,100,220,0,220)
htimecalibvst=ROOT.TH2D('htimecalibvsT','time calib vs t',tbin,0,tmax,120,-20,100)
htimecalibvsamp=ROOT.TH2D('htimecalibvsamp','time calib vs digi amp',120,-20,100,3500,0,70000)
htimecalibvsnsig=ROOT.TH2D('htimecalibvsnsig','time calib vs number of signals',120,-20,100,200,0,200)
#hdigisize

if args.cal:
    csample=ROOT.TCanvas('csample',args.title)
    csample2=ROOT.TCanvas("csample2","csaple2",700,1100)
    csample3=ROOT.TCanvas("csample3","csaple3",700,1100)
    csample4=ROOT.TCanvas("csample4","csample4",700,1100)
    csample5=ROOT.TCanvas("csample5","csample5",700,1100)
else:
    csignal=ROOT.TCanvas('csignal',args.title)
    csample=ROOT.TCanvas('csample',args.title)

for e in tree:
    evcounter+=1
    print "at event:",evcounter
    if args.ev!=-1 and args.ev!=evcounter:
        continue
    if args.nevts!=-1 and evcounter==args.nevts:
        break
#    for sig in e.TpcSignal:
#        print sig.index()
    if args.trk:
        trackarray=e.TrackPostFit
    else:
        trackarray=[0]

    for trk in trackarray:
#        print "next track"
        if args.trk:
            rep=trk.getCardinalRep()
            flag=int(rep.getStatusFlag())
            cand=trk.getCand()
            detid=cand.getDetIDs()[0]
            hits=cand.getHitIDs(detid)
        else:
            hits=[0]

        for hit in hits:
            if args.trk:
                if not args.cal:
                    print "next cluster"
                sphit=e.TpcSPHit.At(hit)
                cl=e.TpcCluster.At(sphit.getClusterID())
                digimap=cl.getDigiMap() 
            else:
                digimap=e.TpcDigi
            for dig in digimap:
                if args.trk:
                    digi=e.TpcDigi.At(dig.first)
                else:
                    digi=dig
                samplemap=digi.getSampleMap()
                samcounter=-1
                t0=0
                t1=0
                maxsama=0
                maxsamt=-99
#                print "next digi"
                if args.sizecut[0]!=-1 and (digi.tlength()>args.sizecut[1] or digi.tlength()<args.sizecut[0]):
                    if args.verbose:
                        print "sizecut!",digi.tlength(),args.sizecut[0],args.sizecut[1]
                    continue
                elif args.sizecut[0]!=-1:
                    if args.verbose:
                        print "padid:",digi.padId(),evcounter

                if args.padcut!=-1 and digi.padId()!=args.padcut:
                    if args.verbose:
                        print "padcut!",digi.padId(),args.padcut
                    continue

                if args.ampcut[0]!=-1 and (digi.amp()<args.ampcut[0] or digi.amp()>args.ampcut[1]):
                    if args.verbose:
                        print "ampcut!",digi.amp(),args.ampcut[0],args.ampcut[1]
                    continue
                if args.timecut[0]!=-1 and (digi.t()<args.timecut[0] or digi.t()>args.timecut[1]):
                    if args.verbose:
                        print 'timecut',digi.t(),args.timecut[0],args.timecut[1]
                    continue
                signals=[]
                discard=False
                for sam in samplemap :
                    samcounter+=1
                    sample = e.TpcSample.At(sam.first)
                    t1=sample.t()
                    if samcounter==0:
                        t0=sample.t()
                    if discard:
                        continue
                    if samcounter<2:
                        if (args.startampcut!=-1 and sample.amp()>args.startampcut):
                            discard=True
                            continue
                    stime=sample.t()
                    samp=sample.amp()
                    if samp>maxsama:
                        maxsama=samp
                        maxsamt=stime
                    if signals.count(sample.signalId())==0:
                        signals.append(sample.signalId())
                    if not args.cal:
                        signal = e.TpcSignal.At(sample.signalId())
                        if type(signal)==ROOT.TpcSignal:
                            sigtime=signal.t()
                            sigamp=signal.amp()
                        else:
                            sigtime=-1
                            sigamp=-1
                        hsample.Fill(stime-t0,samp)
                        hsignal.Fill(sigtime,sigamp)
                    if args.cal:
                        hsampleint.Fill(stime-t0,samp/digi.amp())
                        hsampleintvst.Fill(digi.t(),stime-t0,samp/digi.amp())
                if discard:
                    continue
                if args.caltimecut[0]!=-1 and (maxsamt-t0>args.caltimecut[1] or maxsamt-t0<args.caltimecut[0]):
                    hsampleint.Reset()
                    hsampleintvst.Reset()
                    hsample.Reset()
                    hsignal.Reset()
                    if args.verbose:
                        print "timecalcut!",maxsamt-t0,args.caltimecut[0],args.caltimecut[1]
                    continue
                elif args.caltimecut[0]!=-1:
                    if args.verbose:
                        print "padid:",digi.padId(),evcounter
                if args.signumcut[0]!=-1 and (len(signals)<args.signumcut[0] or len(signals)>args.signumcut[1]):
                    if args.verbose:
                        print "signumcut",len(signals),args.signumcut[0],args.signumcut[1]
                    continue
                if args.cal:
                    htimecalib.Fill(maxsamt-t0)
                    htimecalibcalib.Fill(digi.t()-t0)
                    htimecalibvst.Fill(maxsamt,maxsamt-t0)
                    htimecalibvsdigisize.Fill(maxsamt-t0,digi.tlength())
                    htimecalibvsamp.Fill(maxsamt-t0,digi.amp())
                    htimecalibvsnsig.Fill(maxsamt-t0,len(signals))
                    
                if not args.cal:
                    if args.compare:
                        digi2=e.TpcDigi2.At(dig.first)
                        hdigi2.Fill(digi2.t()-t0+10,digi2.amp())
                    csample.cd()
                    hdigi.Fill(maxsamt-t0,digi.amp())
                    hsample.Draw()
                    hdigi.SetLineColor(2)
                    hdigi.Draw('same')
                    hsignal.SetLineColor(3)
                    csignal.cd()
                    hsignal.Draw()
                    csignal.Update()
                    text=ROOT.TLatex()
                    text.SetTextSize(2)
                    if args.compare:
                        hdigi2.SetLineColor(3)
                        hdigi2.Draw('same')
                    csample.Update()
                    u=raw_input('next digi?')
                    if u=='q':
                        exit()
                    hsample.Reset()
                    hdigi.Reset()
                    hsignal.Reset()
                    if args.compare:
                        hdigi2.Reset()
if args.cal:
    csample.cd()
    htimecalib.Draw()
    htimecalibcalib.SetLineColor(2)
    htimecalibcalib.Draw("same")

    csample2.Divide(1,2)
    csample2.cd(1)
    htimecalibvst.Draw("colz")
    csample2.cd(2)
    htimecalibvst.FitSlicesY(0,0,-1,0,"QNRG5")
    ROOT.gDirectory.Get(htimecalibvst.GetName()+"_1").Draw()
   
    csample3.Divide(1,2)
    csample3.cd(1)
    htimecalibvsdigisize.Draw("colz")
    csample3.cd(2)
    htimecalibvsamp.Draw("colz")

    csample4.Divide(1,2)
    csample4.cd(1)
    hsampleint.Draw()
    csample4.cd(2)
    hsampleintvst.Draw("colz")

    csample5.Divide(1,2)
    csample5.cd(1)
    htimecalibvsnsig.Draw("colz")
    csample5.cd(2)

    csample.Update()
    csample2.Update()
    csample3.Update()
    csample4.Update()
    csample5.Update()

    if args.outf!='Null':
        outf=ROOT.TFile(args.outf,'RECREATE')
        htimecalib.Write()
        htimecalibcalib.Write()
        htimecalibvst.Write()
        htimecalibvsdigisize.Write()
        htimecalibvsamp.Write()
        hsampleint.Write()
        hsampleintvst.Write()
        htimecalibvsnsig.Write()
        outf.Close()
    while 1:
        u=raw_input("done?")
        if u=='q':
            exit()