import ROOT, argparse,sys
sys.path.append('$PANDAPATH/macro/tpc/FOPI/mberger')
from functions import sort_inner, sort_inner2, Drawh, Divideh, fillcutred, plotfullevent,set_palette

parser=argparse.ArgumentParser(description='count digis before and after reclustering')
parser.add_argument("filename",help="the file to analyze")
parser.add_argument("--mc",help='enable mc stuff',action="store_true")
parser.add_argument("--nevts",help='number of the events to analyze',type=int,default=-1)
parser.add_argument("--hlp",help="display help",action='store_true')
args=parser.parse_args()

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

ROOT.gROOT.ProcessLine(".x rootlogon.C") 
ROOT.gROOT.ProcessLine('gROOT->SetStyle("Plain")')
ROOT.gROOT.LoadMacro("stlPYROOT.h+")
set_palette()


Rfile=ROOT.TFile.Open(args.filename)
Dfile=args.filename.replace("reco","raw")
Dfile=Dfile.replace("_recl","")
Mfile=args.filename.replace("reco","mc")
Mfile=Mfile.replace("_recl","")

tree=Rfile.Get("cbmsim")
if tree == None :
    print "No cbmsim in file", file
    if type(Rfile)==ROOT.TFile:
        Rfile.Close()
        exit()

if args.mc:
    tree.AddFriend("cbmsim",Dfile)
    tree.AddFriend("cbmsim",Mfile)

tree.SetBranchStatus("*",1)

ev=-1
for e in tree:
    ev+=1
    predigis=0
    recldigis=0
    if args.nevts!=-1 and ev>args.nevts:
        break
    preDigiIds=[]
    digiIds=[]

    for tr in e.TrackPrePostFit:
        cand=tr.getCand()
        hitIds=cand.getHitIDs(cand.getDetIDs()[0])
        for hitId in hitIds:
            hit=e.TpcPreSPHit.At(hitId)
            cl=e.TpcPreCluster.At(hit.getClusterID())
            digimap=cl.getDigiMap()
            for dig in digimap:
                digi=e.TpcDigi.At(dig.first)
                if preDigiIds.count(dig.first)>0:
                    continue
                preDigiIds.append(dig.first)

    for tr in e.TrackPostFit:
        cand=tr.getCand()
        hitIds=cand.getHitIDs(cand.getDetIDs()[0])
        digiIds=[]
        for hitId in hitIds:
            hit=e.TpcSPHit.At(hitId)
            cl=e.TpcCluster.At(hit.getClusterID())
            digimap=cl.getDigiMap()
            for dig in digimap:
                digi=e.TpcDigi.At(dig.first)
#                if digiIds.count(dig.first)>0:
#                    continue
                digiIds.append(dig.first)
                if preDigiIds.count(dig.first)==0:
                    print "Missing digi:",ev," ",dig.first
                if digiIds.count(dig.first)>1:
                    print "Double counted digi:",ev," ",dig.first
    if len(preDigiIds)!=len(digiIds):
        print "event:",ev,"pre:",len(preDigiIds),"recl:",len(digiIds)