import ROOT,glob,os,time,sys
pandapath=os.environ.get('PANDAPATH')
sys.path.append(pandapath+'/macro/tpc/FOPI/python/argparse-1.2.1')
import argparse

parser=argparse.ArgumentParser(description="meger root files of given run number(s)")
parser.add_argument("--run",help="the run numbers to merge",type=str,default="NOT")
parser.add_argument("path",help="path to the place where the files are to find.",type=str,default="")
parser.add_argument("--files",help="the files to merge",type=str,default=[""],nargs="+")
parser.add_argument('--filelist',help='a file with a list of the files to use',type=str,default="")
parser.add_argument('--move',help='just move the files into a folder and create a text file with all files',type=str,default='')
parser.add_argument("--out",help="the filename of the output file",type=str,default="")
parser.add_argument("--pattern",help='a certain pattern for the filenames',type=str,default="")
parser.add_argument("--check",help="only check the files",action="store_true")
parser.add_argument("--finecheck",help="only check the files but do this eventwise",action='store_true')
parser.add_argument("--flagcheck",help='check statusd flag of tracks',action='store_true')
parser.add_argument("--rm",help='remove bad files',action='store_true')
parser.add_argument("--delIfEntry",help='delete part files if total event number matches entries in chain',type=int,default=-1)
parser.add_argument("--list",help='just list the files to merge',action='store_true')
parser.add_argument("--wait",help='wait until a defined nuber of files is available for merging',type=int,default=-1)
parser.add_argument('--runwise',help='merge each run into a file',action='store_true')
parser.add_argument("--hlp",help="show help",action="store_true")

args=parser.parse_args()

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

ROOT.gROOT.ProcessLine(".x rootlogon.C") 
ROOT.gROOT.ProcessLine('gSystem->Load("libPhysics")')
ROOT.gROOT.ProcessLine('gStyle->SetPalette(1)')
ROOT.gROOT.ProcessLine('gROOT->SetStyle("Plain")')
ROOT.gROOT.ProcessLine('gSystem->Load("libGeom")')
ROOT.gROOT.LoadMacro("stlPYROOT.h+")

chain=ROOT.TChain("cbmsim")

files=glob.glob(args.path + "/*"+args.pattern+"*reco.root")
#files=glob.glob(args.path + "/Cosmics*"+args.pattern+"*.root")
#print files

if args.wait!=-1:
    files=glob.glob(args.path + "/*"+args.pattern+"*reco.root")
    #files=glob.glob(args.path + "/Cosmics*"+args.pattern+"*.root")
    while len(files)!=args.wait:
        time.sleep(10)
files.sort()

if len(files)==0:
    print 'no files to merge in',args.path
    exit()

runList=[]

if args.run != "NOT":
    colindex = args.run.find(",")
    if colindex > 0:
        runs = args.run.split(",")
        for i in range(len(runs)) :
            dashIndex = runs[i].find("-")
            if dashIndex > 0 : 
                runs[i] = runs[i].split("-")
                for j in range(int(runs[i][0]), int(runs[i][1])+1) :
                    runList.append(j)
            if dashIndex < 0 and len(runs) > 1 :
                runList.append(int(runs[i]))
    else :
        dashIndex = args.run.find("-")
        if dashIndex > 0 :
            runs = args.run.split("-")
            runList = range(int(runs[0]), int(runs[1])+1)
        if dashIndex < 0 and len(args.run) > 1 :
            runList = [args.run]
else:
    runList=files

if args.filelist!="":
    runList=[]
    for line in open(args.filelist):
        runList.append(line.replace("\n",""))

if args.list:
    for f in files:
        print f
    exit()

damfilename=args.path.replace("/parts","")+"/"+args.pattern+"Errorfile.txt"
damfile=open(damfilename,"w")
if args.runwise:
    chain={}
else:
    chain=ROOT.TChain("cbmsim")
found=False
damaged_file=False

moved_files=0
if args.move!='':
    os.system('mkdir {0}'.format(args.move))
    filelist=open('{0}/files.txt'.format(args.move),'w')

for run in runList:
    if args.runwise:
        chain[str(run)]=ROOT.TChain("cbmsim")

    for f in files:
        if f.find('FairRunInfo')!=-1:
            continue
        good=True
        if f.find(str(run))!=-1:
            found=True
            RFile=ROOT.TFile.Open(f)
            tree=RFile.Get("cbmsim")
            if tree==None:
                val=f.rfind("/")
                fil=f[val+1:]
                damfile.write(str(fil)+"\n")
                damaged_file=True
                print "no tree found in file ",f
                RFile.Close()
                if args.rm:
                    print 'removing file',f
                    os.system('rm '+str(f))
                continue

            if args.runwise:
                chain[str(run)].AddFile(f)
            elif args.move=='':
                chain.AddFile(f)

            nofEvts=tree.GetEntriesFast()
            if args.finecheck:
                print 'checking events:'
                ecounter=0
                tree.SetBranchStatus("*",0)
                tree.SetBranchStatus("TrackFitStat_0.*",1)
                flags=0
                tracks=0
                for e in tree:
                    ecounter+=1
                    if args.flagcheck:
                        for tr in e.TrackFitStat_0:
                            flags+=tr.GetStatFlag()
                            tracks+=1
                if args.flagcheck:
                    if tracks>0:
                        val=float(flags)/float(tracks)
                        print 'flagvaldinges=',val,'track=',tracks,'flags=',flags
                        if val>0.95:
                            print 'bad',val
                            good=False
                    elif tracks<3:
                        good=False
                        print 'bad, no tracks inside'
                print ecounter,'events, all good'

            RFile.Close()
            print f,nofEvts
            if not good and args.rm:
                print 'removing file',f
                os.system('rm '+str(f))

            if f.find("_000")>0:
                outfname=f.replace("_000","")

            if args.finecheck:
                print ""
            
            if args.move!='':
                print 'moving files into folder:',args.move
                os.system('mv {0} {1}/'.format(f,args.move))
                moved_files+=1
                fname=f.split('/')[-1]
                filelist.write('{0}/{1}\n'.format(args.move,fname))

                
if args.move!='':
    filelist.close()
    number_of_files = len([item for item in os.listdir(args.move) if os.path.isfile(os.path.join(args.move, item))])
    #print 'test',moved_files,number_of_files-1
    if(number_of_files-1==moved_files):
        print 'moved {0} files, found {1} files'.format(moved_files,number_of_files-1)

damfile.close()
if not damaged_file:
    os.system("rm {0}".format(damfilename))

if args.out!="":
    outfname=args.out

totentries=0
merged=False

if args.runwise:
    for run in chain:
        print "Number of events in chain (",run,")",chain[run].GetEntries()
        totentries+=chain[run].GetEntries()
        if not args.check and found:    
            if args.runwise:
                runOutfname=args.path.replace('/parts','')+"/run_"+run+outfname
                paramfile=args.path+"/runC_"+run+outfname
                paramfile=paramfile.replace('.reco.root','part_000.param.root')
            else:
                runOutfname=outfname
                paramfile=outfname
                paramfile=paramfile.replace('.reco.root','.param.root')
            print "merging all above files to:",runOutfname
            chain[run].Merge(runOutfname,"fast")
            merged=True
            thedir=args.path.replace('/parts','')+'/run_'+run
            os.system('mkdir '+thedir)
            print 'moving outfile ({0}) to {1}'.format(runOutfname,thedir)
            os.system('mv -v '+runOutfname+' '+thedir)
            print 'the parameter file:',paramfile
            if os.path.isfile(paramfile):
                newparamfile=paramfile.split('/')[-1]
                newparamfile=newparamfile.replace('part_000','')
                newparamfile=newparamfile.replace('runC','run')
                os.system('cp -v '+paramfile+' '+thedir+'/'+newparamfile)

    print 'Total number of events:',totentries
elif args.move=='':
    print "Total number of events:",chain.GetEntries()
    if not args.check and found:
        print "merging all above files to:",args.path.replace('/parts','')+outfname
        chain.Merge(args.path.replace('/parts','')+outfname,"fast")
        totentries=chain.GetEntries()
        merged=True


print args.delIfEntry,totentries,merged
if args.delIfEntry!=-1 and args.delIfEntry==totentries and merged:
    for f in files:
        os.system('rm -v '+f)

if args.files[0] != "":
    args.files.sort()
    for f in args.files:
        print "adding file:",f
        chain.AddFile(f)

    outfname=args.path
    if args.out!="":
        outfname=args.out
    print "Merging all files to:",outfname
    chain.Merge(outfname,"fast")