import sys,os,time,math,array,copy
pandapath=os.environ.get('PANDAPATH')
sys.path.append(pandapath+'/macro/tpc/FOPI/python/argparse-1.2.1')
sys.path.append(pandapath+'/macro/tpc/FOPI/mberger')
import glob,math,argparse,ROOT
from functions import thisIsTheEnd

def writeDir(dirkey,file1,file2):
    if file2.GetDirectory(dirkey.GetName()):
        return
    dir=file1.Get(dirkey.GetName())
    file2.mkdir(dirkey.GetName())
    keys=dir.GetListOfKeys()
    for k in keys:
        file2.cd(dirkey.GetName())
        #print k,k.GetName(),file1.Get(dirkey.GetName()+'/'+k.GetName())
        file1.Get(dirkey.GetName()+'/'+k.GetName()).Write()

parser=argparse.ArgumentParser(description='plot chi2 values of parameter scan')
parser.add_argument('file',help='the file with all plots',type=str)
parser.add_argument('--isTree',help='file contains a tree created from this script',action='store_true')
parser.add_argument('--fixDiff',help='only plot for fixed diff parameter',type=float,default=-1)
parser.add_argument('--fixXY',help='fix parameter offsetXY',type=float,default=-1)
parser.add_argument('--fixZ',help='fix parameter offsetZ',type=float,default=-1)
parser.add_argument('--fixAmp',help='fix parameter offsetAmp',type=float,default=-1)
parser.add_argument('--maxnum',help='only run up to file number maxnum',type=int,default=-1)


args=parser.parse_args()

c1=ROOT.TCanvas('c1',args.file)

infiles=[]
if os.path.isdir(args.file):
    infiles=glob.glob(args.file+'/*')
    treepath=args.file
    if not os.path.isdir(args.file+'/done'):
        os.system('mkdir '+args.file+'/done')
else:
    infiles=[ROOT.TFile(args.file,'read')]
    treepath=os.path.dirname(args.file)

counter=-1    
graph=ROOT.TGraph()
mingraph=ROOT.TGraph()
minchi2=99999999999
minsetting={}
minname=''
minfile=None
olddone=0
fcounter=-1
precision_scale=100
offxy_scale=1
offZ_scale=10*precision_scale
diffXY_scale=100*(precision_scale)**2
diffZ_scale=1000*(precision_scale)**3
amp_scale=10000*(precision_scale)**4


if args.isTree:
    tuplefile=ROOT.TFile(args.file,'read')
    tree=tuplefile.Get('bat_param')
else:
    tuplefile=ROOT.TFile("{0}/bat_chi2Tuple.root".format(treepath),"recreate")
    resultTuple=ROOT.TNtuple("bat_param",'','offXY:offZ:diffXY:diffZ:amp:chi2')
    tree=None
    mergefile=ROOT.TFile("{0}/bat_parscan_merged.root".format(treepath),"update")

for ifile in infiles:
    if ifile.find("done")!=-1:
        continue

    
    fcounter+=1
    
    if args.isTree:        
        keys=tree
    else:
        if ifile.find('Tuple')!=-1:
            continue
    
        if ifile.find('merged')==-1:
            tmp=ifile.split('_')
            runnum=int(tmp[-1].replace('.root',''))
            if runnum>args.maxnum and args.maxnum!=-1:
                continue
    
        infile=ROOT.TFile(ifile)
        keys=infile.GetListOfKeys()
        

    done=int(float(fcounter*100./len(infiles)))
    if done!=olddone:
        print "{0}%".format(done)
        olddone=done

    mergedFile=False
    if ifile.find('merged')!=-1:
        mergedFile=True
    
    #minchi2=0
    olddonek=0
    kcounter=-1
    for k in keys:
        counter+=1
        kcounter+=1
        if mergedFile:
            donek=int(float(kcounter*100./len(keys)))
            if donek!=olddonek:
                print "->{0}%".format(donek)
                olddonek=donek

        params={}
        if args.isTree:
            chi2=k.chi2
            params['offXY']=k.offXY
            params['offZ']=k.offZ
            params['diffXY']=k.diffXY
            params['diffZ']=k.diffZ
            params['amp']=k.amp
        else:
            if ifile.find('Tuple')!=-1 or ifile.find("done")!=-1:
                continue
            chi2=(float(infile.Get("{0}/chi2".format(k.GetName())).GetTitle()))
            if chi2==0:
                continue
            if math.isnan(chi2) or math.isinf(chi2):
                continue
            
            tmp_params=k.GetName()
            tmp_params=tmp_params.split('_')
            if k.GetName().find('diffZ')==-1:
                params[tmp_params[0][0:5]]=float(tmp_params[0][5:len(tmp_params[0])])
                params[tmp_params[1][0:4]]=float(tmp_params[1][4:len(tmp_params[1])])
                params[tmp_params[2][0:4]]=float(tmp_params[2][4:len(tmp_params[2])])
                params[tmp_params[3][0:3]]=float(tmp_params[3][3:len(tmp_params[3])])
                params['diffZ']=params['diff']
                params['diffXY']=params['diff']
            else:
                params[tmp_params[0][0:5]]=float(tmp_params[0][5:len(tmp_params[0])])
                params[tmp_params[1][0:4]]=float(tmp_params[1][4:len(tmp_params[1])])
                params[tmp_params[2][0:6]]=float(tmp_params[2][6:len(tmp_params[2])])
                params[tmp_params[3][0:5]]=float(tmp_params[3][5:len(tmp_params[3])])        
                params[tmp_params[4][0:3]]=float(tmp_params[4][3:len(tmp_params[3])])        
            resultTuple.Fill(params['offXY'],params['offZ'],params['diffXY'],params['diffZ'],params['amp'],chi2)
            
        
        if (args.fixDiff!=-1 and args.fixDiff!=params['diff']):
            continue
        if (args.fixXY!=-1 and args.fixXY!=params['offXY']):
            continue
        if (args.fixZ!=-1 and args.fixZ!=params['offZ']):
            continue
        if (args.fixAmp!=-1 and args.fixAmp!=params['amp']):
            continue
        

        
        if chi2<minchi2:
            minchi2=chi2
            minsetting=params
            if args.isTree:
                minnam='offXY{0}_offZ{1}_diffXY{2}_diffZ{3}_amp{4}'.format(params['offXY'],params['offZ'],params['diffXY'],params['diffZ'],params['amp'])
            else:
                minname=k.GetName()
                #minfile=infile
            mingraph.SetPoint(0,params['offXY']+params['offZ']*offZ_scale+params['diffXY']*diffXY_scale+params['diffZ']*diffZ_scale+params['amp']*amp_scale,chi2)
        
        
        if ifile.find("merged")==-1 and not args.isTree:
                writeDir(k,infile,mergefile)
        graph.SetPoint(counter+1,params['offXY']+params['offZ']*offZ_scale+params['diffXY']*diffXY_scale+params['diffZ']*diffZ_scale+params['amp']*amp_scale,chi2)
        
    if ifile.find("merged")==-1 and not args.isTree:
        infile.Close()
        os.system('mv '+ifile+' '+args.file+'/done')

        #print params
        #print chi2
        

graph.SetMarkerStyle(20)
graph.SetMarkerSize(0.5)
graph.SetTitle('offXY+offZ*{0}+diffXY*{1}+diffZ*{2}+amp*{3}'.format(offZ_scale,diffXY_scale,diffZ_scale,amp_scale))
graph.GetXaxis().SetTitle('setting')
graph.GetYaxis().SetTitle('chi2')

mingraph.SetMarkerColor(2)
mingraph.SetMarkerStyle(20)
mingraph.SetMarkerSize(0.8)
print 'minchi2:',minchi2,' min setting=',minsetting    

c1.cd()
graph.Draw('AP')
mingraph.Draw('P')
c1.Update()

print minname
if not args.isTree:
    c2=mergefile.Get(minname+'/tmppulls')
    c2.SetTitle(args.file)
    c2.Draw()
    c2.Update()

if not args.isTree:
    tuplefile.cd()
    resultTuple.Write()
            
tuplefile.Close()

thisIsTheEnd()

if not args.isTree:
    mergefile.Close()