import sys,os,math,time,random,string
#pandapath=os.environ.get('PANDAPATH')
pandapath='/home/walbrecht/fopiroot/fopiroot_dev/'
user=os.environ.get('USER')
if len(user)>8:
    user=user[0:7]
sys.path.append(pandapath+'/macro/tpc/FOPI/python/argparse-1.2.1')
import argparse

def id_generator(size=6, chars=string.ascii_uppercase + string.digits):
    return ''.join(random.choice(chars) for _ in range(size))

def checkBatch(getnum=False):
    if os.path.isfile('stopDevmap'):
        exit()

    tmpname='/tmp/batch_tmp_'+id_generator()
    if args.e18:
        os.system('qstat > '+tmpname)
    else:
        os.system('squeue -o "%.18i %.9P %.8j %.8u %.2t %.10M %.6D %R %p" > '+tmpname)
    statfile=open(tmpname,'r')
    njobs=0
    nrjobs=0
    nqjobs=0
    ntjobs=0
    running='R'
    queued='PD'
    if args.e18:
        running='r'
        queued='qw'

    for line in statfile:
        if line.find(user)==-1:
            ntjobs+=1
            continue
        ntjobs+=1
        words=line.split()
#        print words
        if words[4]==running:
            nrjobs+=1
        if words[4]==queued:
            nqjobs+=1
        njobs+=1
    statfile.close()
    os.system('rm '+tmpname)
    
    print 'there are',ntjobs-3,'jobs on the batchfarm'
    print njobs,'are my jobs'
    print nrjobs,'jobs are running'
    print nqjobs,'jobs are in the queue'
    print args.ofile
    print ''

    if njobs<args.minjobs or getnum:
        return njobs
    else:
        return -1

def readFromFile():
    for line in open(args.settings,'r'):
        if line[0]=='#':
            continue
        words=line.split()
        args.minjobs=int(words[0])
        args.maxjobs=int(words[1])
        args.checkintervall=int(words[2])

def runCheck():
    global jobcounter
    print 'min:',args.minjobs,'max:',args.maxjobs,'wait:',args.checkintervall,'jobs:',jobcounter
    if jobcounter>=args.maxjobs:
        print 'jobs:',jobcounter
        check=-1
        while check==-1:
            timer=0
            while timer<args.checkintervall:
                print args.checkintervall-timer,'seconds till next check'
                time.sleep(10)
                timer+=10
            if args.settings!='NONE':
                readFromFile()
            check=checkBatch()

        jobcounter=check
        print 'jobs now:',jobcounter

def sub(r,z):
    if args.settings!='NONE':
        readFromFile()
    global jobcounter
    if args.zStep!=-1:
        zStep=args.zStep
    else:
        zStep=10

    for rBin in range(0,args.nr):
        if (r!=-1 and r!=rBin):
            continue

        countBIN=0
        if args.BIN != -1:
            if rBin<args.BIN:
                countBIN=1
        if countBIN==1:
            continue

        os.environ['efile']=args.efile
        os.environ['parfile']=args.parfile
        os.environ['ofile']=outfile
        os.environ['nr']=str(args.nr)
        os.environ['nz']=str(args.nz)
        os.environ['rBin']=str(rBin)
        os.environ['specZ']=args.specZ
        os.environ['zStep']=str(zStep)
        arrcomm="-t "
        if not args.array:
            for zBin in range(0,int(math.ceil( float(args.nz) / float(zStep) ) ) ):

                jobname="Dev"+str(rBin)+"_"
                if (z.count(-1)==0 and z.count(zBin)==0):

                    continue
                if args.specZ!="-1":
                    os.environ['zBin']=str(-1)
                    if not args.array:
                        print "submitting rBin",rBin, "zBin", args.specZ
                    jobname+=args.specZ
                else:
                    print zBin
                    os.environ['zBin']=str(zBin)
                    if not args.array:
                        print "submitting rBin",rBin, "zBin", zBin
                    jobname+=str(zBin)    
   
                  


                if args.inter:
                    print 'starting interactive mode'
                    if args.align:
                        if args.e18:
                            os.system('macro/tpc/FOPI/BatchTools/batchDevMapAlignE18.sh')
                        else:
                            os.system('macro/tpc/FOPI/BatchTools/batchDevMapAlign.sh')
                    else:
                        if args.e18:
                            os.system('macro/tpc/FOPI/BatchTools/batchDevMapE18.sh')
                        else:
                            os.system('macro/tpc/FOPI/BatchTools/batchDevMap.sh')
                            
                elif not args.array:
                    runCheck()
                    if args.e18:
                        os.system('qsub -N '+jobname+' -cwd -v "efile","parfile","ofile","nr","nz","rBin","zBin","zStep","specZ" -o /nfs/mnemosyne/user/ne86qok/private/sgeout/ -j y -l h_vmem=1000M,short=TRUE macro/tpc/FOPI/BatchTools/batchDevMapE18.sh')
                        jobcounter+=1
                    else:
                        if args.align:
                            os.system('sbatch -J '+jobname+' -o ~/sgeout/'+args.ofile+'_'+jobname+' -e ~/sgeout/'+args.ofile+'_'+jobname+' --time=120 --mem-per-cpu=500 macro/tpc/FOPI/BatchTools/batchDevMapAlign.sh')
                        else:
                            os.system('sbatch -J '+jobname+' -o ~/sgeout/'+args.ofile+'_'+jobname+' -e ~/sgeout/'+args.ofile+'_'+jobname+' --time=120 --mem-per-cpu=500 macro/tpc/FOPI/BatchTools/batchDevMap.sh')
                        jobcounter+=1                       

                if args.array:
                    arrcomm+=str(zBin)+","
                if args.specZ!="-1":
                    break
        if args.array:
            runCheck()
            if arrcomm[-1]==',':
                arrcomm=arrcomm[:-1]
            arrcomm="-t 1-"+str(int(math.ceil(args.nz/zStep))+1)
            jobname="Dev"+str(rBin)
            print "submitting array job, rBin:",rBin,"zbins:",arrcomm
            if args.align:
                os.system('qsub '+arrcomm+" -N "+jobname+' -cwd -v "efile","parfile","ofile","nr","nz","rBin","zStep" -o /nfs/mnemosyne/user/ne86qok/private/sgeout/ -j y -l h_vmem=1000M,short=TRUE macro/tpc/FOPI/BatchTools/batchDevMapArrayAlign.sh')
            else:
                os.system('qsub '+arrcomm+" -N "+jobname+' -cwd -v "efile","parfile","ofile","nr","nz","rBin","zStep" -o /nfs/mnemosyne/user/ne86qok/private/sgeout/ -j y -l h_vmem=1000M,short=TRUE macro/tpc/FOPI/BatchTools/batchDevMapArray.sh')
            jobcounter+=1





#ROOT.gROOT.ProcessLine(".x rootlogon.C") 

parser=argparse.ArgumentParser(description='submits the creation of the deviation file onto the batchfarm in radial steps')
parser.add_argument("efile",help='file with the e-field')
parser.add_argument("parfile",help='parameter file')
parser.add_argument("ofile",help='prefix for the outputfile')
parser.add_argument("--nr",help="number of bins in R to calculate",default=-1,type=int)
parser.add_argument("--nz",help="number of bins in Z to calculate",default=-1,type=int)
parser.add_argument("--BIN",help="calculates only from BIN",default=-1,type=int)
parser.add_argument("--rBin",help="if a r single bin should be submitted",default=[-1],type=int,nargs='+')
parser.add_argument("--zBin",help="if a z single bin should be submitted",default=[-1],type=int,nargs='+')
parser.add_argument("--rRange",help='submit all r bins inside a range',default=[-1,-1],type=int,nargs=2)
parser.add_argument("--zRange",help='submit alr z bins inside a range',default=[-1,-1],type=int,nargs=2)
parser.add_argument("--ask",help="ask which bins should be used",action="store_true")
parser.add_argument("--file",help='read bins to submit from file',type=str,default="NOT")
parser.add_argument("--zStep",help="the size of one z-step (number of bins in one step)",type=int,default=-1)
parser.add_argument("--inter",help='use interactive session instead of batch farm',action='store_true')
parser.add_argument("--specZ",help='calculate specific given z-bins, komma separated list',default="-1",type=str)
parser.add_argument("--specfile",help='give a devmap errorfile in format of merge2',default="",type=str)
parser.add_argument("--array",help='submit job as an array job',action='store_true')
parser.add_argument("--e18",help='run at e12',action='store_true')
parser.add_argument("--align",help='turn alignment on',action='store_true')
parser.add_argument('--minjobs',help='minimum number of jobs for new submisson',type=int,default=50)
parser.add_argument('--maxjobs',help='maximum number of jobs at a time',type=int,default=100)
parser.add_argument('--checkintervall',help='the time intervall (in sec) in which the batch farm is checked for jobs',type=int,default=60)
parser.add_argument('--settings',help='a file with the settings for submission (one line: minjobs  maxjobs  checkintervall)',type=str,default='NONE')
args = parser.parse_args()

if os.path.isfile('stopDevmap'):
    os.system('rm stopDevmap')

basepath='/nfs/hicran/project/panda/SIM/bergerm/fopiroot/'
scriptpath='macro/tpc/FOPI/mberger/last_strip/createDevMapMulti_mb.py'
global jobcounter
#jobcounter=args.maxjobs
jobcounter=0

#rBin=1
if args.ofile.find("/")!=-1:
    print "are you sure you put in the ofile thingy correct?"
    exit()
outfile=args.ofile+"_"    
if (args.nr==-1 or args.nz==-1):
    efile=open(args.efile,'r')
    line=efile.readline()
    words=line.split()
    if args.nr==-1:
        args.nr=int(words[3])
        print args.nr
    if args.nz==-1:
        args.nz=int(words[6])
        print args.nz

if args.zRange.count(-1)==0:    
    args.zBin=range(args.zRange[0],args.zRange[1]+1)


if args.ask:
    args.rBin=int(raw_input("Which r-Bin?"))
    elem=(raw_input("which z-Bins?"))
    elem=elem.split(" ")
    args.zBin=[]
    for e in elem:
        if e!='':
            args.zBin.append(int(e))
    print "not implementet at the moment!"
    exit()

if not args.inter:
    print 'setting jobcounter'
    if args.settings!='NONE':
        readFromFile()
    jobcounter=checkBatch(True)
#    jobcounter=checkBatch()
    print 'jobcounter now is:',jobcounter

if args.file!="NOT":
    r=-1
    inf=open(args.file,"r")
    for line in inf:
        zBin=[]
        words=line.split(" ")
        if words[0]=="nothing":
            continue
        for word in words:
#            print word
            if word=="missing" or word=="for":
                continue
            if word.find("Bin")>0:
                r=int(word.split(":")[1])
#                print "r=",r,"zbins=",zBin
                sub(r,zBin)
                continue
            zBin.append(int(word))
    exit()

if args.specfile!="":
    inf=open(args.specfile,"r")
    zbins=[]
    for line in inf:
        words=line.split()
        if words[0]=="r-Bin:":
            r=int(words[1])
        if words[2]=='z-Bin:':
            zbins=words[3].split(",")
        if words[2]=='complete':
            args.specZ='-1'
            sub(r,[-1])
            continue
        if words[2]=='is':
            sub(r,[-1])
            continue
        counter=0
        args.specZ=""
        for zbin in zbins:
            if counter==0:
                args.specZ+=zbin
            else:
                args.specZ+=","+zbin
            counter+=1
            if counter==args.zStep:
                sub(r,[-1])
                args.specZ=""
                counter=0


        if args.specZ!="":
            sub(r,[-1])
    print 'all done, used specfile:',args.specfile
    exit()

if args.rRange[0]!=-1:
    args.rBin=[]
    for i in range(args.rRange[0],args.rRange[1]+1):
        args.rBin.append(i)



for rb in args.rBin:
    print 'submitting',rb,args.zBin
    sub(rb,args.zBin)