#!/usr/bin/python


import os
import sys
import shutil
import re
import optparse
import os.path
import time
import glob
import commands
import subprocess
import getpass

from bcv import git_version

# =============================================================================
#
#                               PREAMBLE
#
# =============================================================================
def parse_commandline():
        """
        Parse the options given on the command-line.
        """
        parser = optparse.OptionParser(usage = __doc__,version=git_version.verbose_msg)

	parser.add_option("-G", "--source",help="Path to source directory.",default='source')
	parser.add_option("-A", "--sourceA",help="Path to source A directory.",default='0')
	parser.add_option("-B", "--sourceB",help="Path to source B directory.",default='0')
        parser.add_option("-r", "--fs",help="Expected sampling rate of (primary) input time series.",default = 16384)
        parser.add_option("-T", "--threshMan",help="Manual threshold, on or off.",default = 1)
        parser.add_option("-c", "--HbicohThresh",help="Manual threshold for function 1 or 2.",default = 0.996)
	# updating HbicohThresh default from 0.9606 to 0.996 (3-19-18)
        parser.add_option("-H", "--HbicohBspcThresh",help="Manual threshold for function 3.",default = 0.7)
        parser.add_option("-F", "--fnct",help="Function selection.")
        parser.add_option("-m", "--mode",help="Mode selection.",default = 0)
        parser.add_option("-X", "--ChannelA",help="A channel name.",default = 'A')
        parser.add_option("-Y", "--ChannelB",help="B channel name.",default = 'B')
        parser.add_option("-x", "--filt",help="Filter, on/off.",default = 0)
        parser.add_option("-s", "--segslice",help="Slice time series, on/off.",default = 0)
        parser.add_option("-S", "--segStart",help="Segment slice start time (s).",default = 0.0)
        parser.add_option("-E", "--segEnd",help="Segment slice end time (s).",default = 1.0)
        parser.add_option("-N", "--full_SNR",help="Process SNR on full time segment before segment slice, on/off.",default = 1)
        parser.add_option("-a", "--autom",help="Automatic overlap parameter calculating, on/off .",default = 1)
        parser.add_option("-t", "--time_window",help="Desired length of time for analysis (s).",default = 4.0)
        parser.add_option("-o", "--offset_multiplier",help="(See manual.  Normally this should be calculated automatically by setting autom to 1.)",default = 1.0)
        parser.add_option("-O", "--ol",help="Desired segment overlap (%).",default = 98.0)
        parser.add_option("-g", "--seg",help="Desired number of experiments.",default = 30)
        parser.add_option("-l", "--tile",help="Tiled option, on/off.",default = 0)
        parser.add_option("-C", "--GPS_central",help="Expected GPS central time.",default = 0.0)
        parser.add_option("-k", "--check",help="Check parameters, on/off.",default = 1)
        parser.add_option("-D", "--dnsmple",help="Downsample, on/off.",default = 0)
        parser.add_option("-d", "--decm",help="Decimate, on/off.",default = 1)
        parser.add_option("-p", "--plot_img",help="Plot, on/off.",default = 0)
        parser.add_option("-R", "--randomize",help="Use phase randomization, on/off.",default = 1)
        parser.add_option("-q", "--preSeq",help="Use user-provided random number sequence, on/off.",default = 0)
        parser.add_option("-w", "--sequenceN",help="User-provided random number list (typically a text file), full path (string).",default = '/home/bernard.hall/bispectral_suite/random_seq/seq_552015_2.txt')
        parser.add_option("-f", "--fsd",help="Target downsampling or decimation sampling rate (Hz).",default = 4096)
        parser.add_option("-u", "--uSnrV",help="User-input SNR value.",default = 1.0)
        parser.add_option("-U", "--upr",help="Upper bispectral value for bispectrogram calculation.",default = 1.0)
        parser.add_option("-P", "--path",help="Output path.",default = '')
        parser.add_option("-i", "--ifoChannel",help="IFO and channel information.",default = '')
	parser.add_option("-M", "--Memory",help="Amount of memory to request.",default = 8192)
	parser.add_option("-K", "--checkpow2",help="Ensure cross bicoherence channels have a power of 2 sampling rate calculated.",default = 1)	

	opts, args = parser.parse_args()
        return opts
# =============================================================================
#
#                                    Main
#
# =============================================================================
def main():
	#######################################################################
        ## parse the command line
        opts = parse_commandline()
	#sourceDir = 'gaussian_1000_source'
	#sourceDir = 'gaussian_source_full'
	sourceDir = opts.source
	sourceDirA = opts.sourceA
	sourceDirB = opts.sourceB
	#funct = 2
	funct = int(opts.fnct)
	filt = int(opts.filt)
	#outDir = 'gaussian_output'
	#outDir = 'gaussian_T_0_8_50p'
	outDir = opts.path 
	#outDir = 'gaussian_1000_output'
	#overlap = 50.0 #98.0 (default)
	overlap = float(opts.ol)
	#decimate = 0
	decimate = int(opts.decm)	
	fs = float(opts.fs)
	threshMan = int(opts.threshMan)
	HbicohThresh = float(opts.HbicohThresh)
        HbicohBspcThresh = float(opts.HbicohBspcThresh)
	mode = int(opts.mode)
	ChannelA = opts.ChannelA
	ChannelB = opts.ChannelB	
	segslice = int(opts.segslice)
	segStart = float(opts.segStart)
	segEnd = float(opts.segEnd)
	full_SNR = int(opts.full_SNR)
	autom = int(opts.autom)
	time_window = float(opts.time_window)
	offset_multiplier = float(opts.offset_multiplier)
	seg = int(opts.seg)
	tile = int(opts.tile)
	GPS_central = float(opts.GPS_central)
	checkM = int(opts.check)
	dnsmple = int(opts.dnsmple)
	plot_img = int(opts.plot_img)
	randomize = int(opts.randomize)
	preSeq = int(opts.preSeq)
	sequenceN = opts.sequenceN
	fsd = float(opts.fsd)
	uSnrV = float(opts.uSnrV)
	upr = float(opts.upr)
	ifoChannel = opts.ifoChannel
	mem = str(opts.Memory)
	checkpow2 = int(opts.checkpow2)

	if ((funct == 0) or (funct== 1)) and ((sourceDirA == '0') or (sourceDirB == '0')):
		print 'Error: cross-bispectrum/cross-bicoherence selected, but no channel A or channel B information provided!  Exiting...'
		sys.exit()
	###################################################################
	accGroup = 'ligo.dev.o2.detchar.nonlin_coup.twochanveto'
	#accGroup = 'ligo.dev.o1.detchar.nonlin_coup.bcv'
	check = 1 # check whether a job has already been submitted to Condor
	ext = 'txt' # string against which to compare the "tag" variable to verify type
	#option = 2
	#wd = os.getcwd()
	# Get full path of current directory --------------------------------
	wd = commands.getstatusoutput('pwd')
	wd = wd[1]
	#--------------------------------------------------------------------
	#print wd
	lim = commands.getstatusoutput('ulimit -Hn')
	#currProcessCount = commands.getstatusoutput('lsof -u %s 2>/dev/null | wc -l'%getpass.getuser())
	lim = int(lim[1])
	spaceleft = commands.getstatusoutput('df -hT /home/%s/'%getpass.getuser())
	spaceleft = spaceleft[1]
	print spaceleft
	#currProcessCount = int(currProcessCount[1]) 
	#print lim
	#print currProcessCount
	###################################################################
	#runList = os.listdir(sourceDir)
	runList =  glob.glob(sourceDir + '/*.' + ext)
	runList.sort()
	if sourceDirA != '0':
		runListA =  glob.glob(sourceDirA + '/*.' + ext)
		if len(runList) != len(runListA): # check if the number of time-series files for channel A is the same as for GW channel; they should be equal
			print 'Error: Channel A file list is not the same length as GW file list!'
			sys.exit()
		runListA.sort()
	if sourceDirB != '0':
		runListB =  glob.glob(sourceDirB + '/*.' + ext)
		if len(runList) != len(runListB): # see above
                        print 'Error: Channel B file list is not the same length as GW file list!'
                        sys.exit()
		runListB.sort()
	#print runList
		
	if check:
		doneList = os.listdir(outDir)
	#limit = 100
	limit = 2

	k = 1
	j = 1
	#flag = 1

	if not os.path.isfile('batch.ini'):
        	num = open('batch.ini', 'w')
	        num.write('1')
        	num.close()

		f = open('list_'+ str(k) + '.txt', 'w')
		flag = 1
		for line in runList: # GW channel lists
			if not (j % limit):
				f.write(line + '\n')
				f.close()
				k = k + 1
				f = open('list_'+ str(k) +'.txt', 'w')
				flag = 0
			if flag:
				f.write(line + '\n')
			else:
				flag = 1
			j = j + 1	
		f.close()
		
		if sourceDirA != '0': # Channel A lists
			k = 1
			j = 1
			f = open('list_A_'+ str(k) + '.txt', 'w')
        	        flag = 1
                	for line in runListA:
                        	if not (j % limit):
                                	f.write(line + '\n')
	                                f.close()
        	                        k = k + 1
                	                f = open('list_A_'+ str(k) +'.txt', 'w')
                        	        flag = 0
	                        if flag:
        	                        f.write(line + '\n')
                	        else:
                        	        flag = 1
	                        j = j + 1
		        f.close()
		if sourceDirB != '0': # Channel B lists
                        k = 1
                        j = 1
                        f = open('list_B_'+ str(k) + '.txt', 'w')
                        flag = 1
                        for line in runListB:
                                if not (j % limit):
                                        f.write(line + '\n')
                                        f.close()
                                        k = k + 1
                                        f = open('list_B_'+ str(k) +'.txt', 'w')
                                        flag = 0
                                if flag:
                                        f.write(line + '\n')
                                else:
                                        flag = 1
                                j = j + 1
                        f.close()
	#print runList

	num = open('batch.ini', 'r')
	number = num.readline()
	number = number.strip('\n')
	num.close()	

	#print number
	if os.path.isfile('list_'+ number + '.txt'):
		print 'Opening ' + '"list_'+ number + '.txt"'
		lt = open('list_'+ number + '.txt', 'r')
		listThr = list(lt)
		lt.close()
		if sourceDirA != '0':
			if os.path.isfile('list_A_'+ number + '.txt'):
		                print 'Opening ' + '"list_A_'+ number + '.txt"'
                		ltA = open('list_A_'+ number + '.txt', 'r')
                		listAThr = list(ltA)
                		ltA.close()
			else:
				print 'Missing channel A file(s)! (or incorrect list A file name)...exiting...'
		                sys.exit()
		if sourceDirB != '0':
                        if os.path.isfile('list_B_'+ number + '.txt'):
                                print 'Opening ' + '"list_B_'+ number + '.txt"'
                                ltB = open('list_B_'+ number + '.txt', 'r')
                                listBThr = list(ltB)
                                ltB.close()
                        else:
                                print 'Missing channel B file(s)! (or incorrect list A file name)...exiting...'
                                sys.exit()
	else:
		print 'Reached end of list (or incorrect list file name)...exiting...'
		sys.exit()
	#print listThr

	#num = open('batch.ini', 'w')
	#num.write(str(int(number) + 1))
	#num.close()

	#print doneList
	#print ('bicoh_out_300.txt' in doneList)
	#sys.exit()

	i = 1 + ((int(number) - 1) * 100)
	
	if not  outDir == '.':
		if not os.path.isfile(outDir + '/links.txt'):  # create a text file to record file linkage between source and output directories
			match = open(outDir + '/links.txt', 'w')
		else:
			match = open(outDir + '/links.txt', 'a')
	else:
		if not os.path.isfile('links.txt'):
                        match = open('links.txt', 'w')
                else:
                        match = open('links.txt', 'a')
	#setcmd = 'export CONDOR_REQUIREMENTS="(Memory * 1024) >= ImageSize"'
	#print "Setting environment variables..."
	#os.system(setcmd)
	if outDir == '.':
		outDirC = ''
	else:
		outDirC = outDir + '/'
	print len(listThr)
	currProcessCount = 0
	success = 1
	#cntA = 0
	#cntB = 0
	print "Submitting jobs..."
	#for line in listThr:
	for cntAB, line in enumerate(listThr):
		fl =  1
		cPC = currProcessCount
		#print cPC
		try:
			currProcessCount = commands.getstatusoutput('lsof -u %s 2>/dev/null | wc -l'%getpass.getuser())
			currProcessCount = int(currProcessCount[1])
			if (cPC > currProcessCount) and (currProcessCount < 5):
				fl = 0
			else:
				print 'Current process count ' + str(currProcessCount) + ' of ' + str(lim) + ' limit...'
		except:
			fl = 0
			success = 0
			print 'Current process count ' + str(currProcessCount)	
		if (currProcessCount < lim) and fl:
			line = line.split('/')
	                line = line[len(line) - 1]
			#print line
			prefix = line.split('.')
			tag = prefix[len(prefix)-1]  # get extension for filename to verify type
			tag = tag.strip('\n')
			tag = tag.strip('\r')
			line = line.strip('\n')
			line = line.strip('\r')
			prefix = prefix[0:len(prefix)-1]
			pjoin = '.'
			prefix = pjoin.join(prefix)
			chG = line.split('+')
			chG = float(chG[0])
			#print prefix
			#print tag
			#prefix = prefix[0]
			if sourceDirA != '0': # check if source A directory parameter has been entered, or is blank/zero
				lineA = listAThr[cntAB]
				lineA = lineA.split('/')
	                        lineA = lineA[len(lineA) - 1]
        	                #print line
                	        prefixA = lineA.split('.')
                        	tagA = prefixA[len(prefixA)-1]
	                        tagA = tagA.strip('\n')
        	                tagA = tagA.strip('\r')
                	        lineA = lineA.strip('\n')
                        	lineA = lineA.strip('\r')
	                        prefixA = prefixA[0:len(prefixA)-1]
        	                pjoinA = '.'
                	        prefixA = pjoinA.join(prefixA)
				#cntA = cntA + 1
				chA = lineA.split('+')
                        	chA = float(chA[0])
				if chA != chG:  # if channels are not aligned in time, then there is a problem -> abort
					print 'Error: GW/Channel A central time mismatch!  Exiting...'
					sys.exit()
			if sourceDirB != '0': # check if source B directory parameter has been entered, or is blank/zero
                                lineB = listBThr[cntAB]
                                lineB = lineB.split('/')
                                lineB = lineB[len(lineB) - 1]
                                #print line
                                prefixB = lineB.split('.')
                                tagB = prefixB[len(prefixB)-1]
                                tagB = tagB.strip('\n')
                                tagB = tagB.strip('\r')
                                lineB = lineB.strip('\n')
                                lineB = lineB.strip('\r')
                                prefixB = prefixB[0:len(prefixB)-1]
                                pjoinB = '.'
                                prefixB = pjoinB.join(prefixB)
                                #cntB = cntB + 1
				chB = lineB.split('+')
                                chB = float(chB[0])
                                if chB != chG:
                                        print 'Error: GW/Channel B central time mismatch!  Exiting...'
                                        sys.exit()
			if tag == ext: # may not need this now that this is accounted for with glob
				if (funct == 0) or (funct== 1):
                                        #shRun = '"bicoherence_analysis -G ' + sourceDir + '/' + line + ' -A ' + sourceDirA + '/' + lineA + ' -B ' + sourceDirB + '/' + lineB + ' -F ' + str(funct) +' -m ' + str(mode) + ' -u ' + str(uSnrV) + ' -P ' + wd + '/' + outDirC + ' -f ' + str(fsd) + ' -q ' + str(preSeq) + ' -w ' + sequenceN + ' -T ' + str(threshMan) + ' -d ' + str(decimate) + ' -i ' + prefix  + ' -O ' + str(overlap) + ' -x ' + str(filt) + '"'

					shRun = '"bicoherence_analysis -G ' + sourceDir + '/' + line + ' -A ' + sourceDirA + '/' + lineA + ' -B ' + sourceDirB + '/' + lineB + ' -F ' + str(funct) +' -m ' + str(mode) + ' -u ' + str(uSnrV) + ' -P ' + wd + '/' + outDirC + ' -f ' + str(fsd) + ' -q ' + str(preSeq) + ' -w ' + sequenceN + ' -T ' + str(threshMan) + ' -d ' + str(decimate) + ' -i ' + prefix  + ' -O ' + str(overlap) + ' -x ' + str(filt) + ' -r ' + str(fs) + ' -c ' + str(HbicohThresh) + ' -s ' + str(segslice) + ' -S ' + str(segStart) + ' -E ' + str(segEnd) + ' -N ' + str(full_SNR) + ' -a ' + str(autom) + ' -t ' + str(time_window) + ' -o ' + str(offset_multiplier) + ' -g ' + str(seg) + ' -l ' + str(tile) + ' -C ' + str(GPS_central) + ' -k ' + str(checkM) + ' -D ' + str(dnsmple) + ' -R ' + str(randomize) + ' -U ' + str(upr) + ' -p ' + str(plot_img) + ' -X ' + str(checkpow2) + '"'
                                        cmd = 'nohup condor_run -a RequestMemory=' + mem + ' -a accounting_group=' + accGroup + ' ' + shRun + ' > ' + outDirC + 'bicoh_out_'+ str(i) +'.txt &'
                                        #i = i + 1
                                        if check == 1:
                                                if not ('bicoh_out_'+ str(i) +'.txt' in doneList):
                                                        print cmd + '\n'
                                                        match.write(sourceDir + ' ' + line + ' ' + outDir + ' bicoh_out_' + str(i) + '\n')
                                                        os.system(cmd)
                                                        if not (i % 10):
                                                                time.sleep(5)
                                                else:
                                                        print 'Job already processed...skipping...(bicoh_out_'+ str(i) +'.txt)'
                                        else:
                                                print cmd + '\n'
                                                match.write(sourceDir + ' ' + line + ' ' + outDir + ' bicoh_out_' + str(i) + '\n')
                                                os.system(cmd)
                                                if not (i % 10):
                                                        time.sleep(5)
                                        i = i + 1

				else:	
					#if option == 1:
						#shRun =	'"bicoherence_analysis -G source/' + line + ' -F 3 -w random_seq/seq_552015_2.txt -m 1 -u 14.026610259768070 -P /home/bernard.hall/bispectral_suite/output/ -f 4096 -q 1 -w /home/bernard.hall/bispectral_suite/random_seq/seq_552015_2.txt -l 1 -a 0 -i ' + prefix + '"'
					#	shRun = '"bicoherence_analysis -G ' + sourceDir + '/' + line + ' -F ' + str(funct) +' -m 1 -u 14.026610259768070 -P /home/bernard.hall/bispectral_suite/' + outDirC + ' -f 4096 -q 1 -w /home/bernard.hall/bispectral_suite/random_seq/seq_552015_2.txt -T 1 -c 0.8 -d 1 -i ' + prefix  + '"'
					#elif option == 2:
					#	shRun = '"bicoherence_analysis -G ' + sourceDir + '/' + line + ' -F ' + str(funct) +' -m 1 -u 14.026610259768070 -P ' + wd + '/' + outDirC + ' -f 4096 -q 1 -w /home/bernard.hall/bispectral_suite/random_seq/seq_552015_2.txt -T 1 -d ' + str(decimate) + ' -i ' + prefix  + ' -O ' + str(overlap) + '"'
					shRun = '"bicoherence_analysis -G ' + sourceDir + '/' + line + ' -F ' + str(funct) +' -m ' + str(mode) + ' -u ' + str(uSnrV) + ' -P ' + wd + '/' + outDirC + ' -f ' + str(fsd) + ' -q ' + str(preSeq) + ' -w ' + sequenceN + ' -T ' + str(threshMan) + ' -d ' + str(decimate) + ' -i ' + prefix  + ' -O ' + str(overlap) + ' -x ' + str(filt) + ' -r ' + str(fs) + ' -c ' + str(HbicohThresh) + ' -s ' + str(segslice) + ' -S ' + str(segStart) + ' -E ' + str(segEnd) + ' -N ' + str(full_SNR) + ' -a ' + str(autom) + ' -t ' + str(time_window) + ' -o ' + str(offset_multiplier) + ' -g ' + str(seg) + ' -l ' + str(tile) + ' -C ' + str(GPS_central) + ' -k ' + str(checkM) + ' -D ' + str(dnsmple) + ' -R ' + str(randomize) + ' -U ' + str(upr) + ' -p ' + str(plot_img) + '"'
					cmd = 'nohup condor_run -a RequestMemory=' + mem + ' -a accounting_group=' + accGroup + ' ' + shRun + ' > ' + outDirC + 'bicoh_out_'+ str(i) +'.txt &'
					#i = i + 1
					if check == 1:
						if not ('bicoh_out_'+ str(i) +'.txt' in doneList):
							print cmd + '\n'
							match.write(sourceDir + ' ' + line + ' ' + outDir + ' bicoh_out_' + str(i) + '\n')
	                        			os.system(cmd)
			                        	if not (i % 10):
        			                        	time.sleep(5)
						else:
							print 'Job already processed...skipping...(bicoh_out_'+ str(i) +'.txt)'
					else:
						print cmd + '\n'
						match.write(sourceDir + ' ' + line + ' ' + outDir + ' bicoh_out_' + str(i) + '\n')
						os.system(cmd)
						if not (i % 10):
							time.sleep(5)
					i = i + 1
		else: # process limit has been reached; no further submissions possible -> exit
			print 'Process limit reached.  Exiting...'
			print '(Current process count ' + str(currProcessCount) + ')'
			#print fl
			match.close()
			sys.exit()
	match.close()
	#success = 0
	if success:
		num = open('batch.ini', 'w')
        	num.write(str(int(number) + 1))
	        num.close()

	spaceleft = commands.getstatusoutput('df -hT /home/%s/'%getpass.getuser()) # check and report space left on user's account
	spaceleft = spaceleft[1]
        print spaceleft
if __name__ == "__main__":
    main()
