#!/usr/bin/python
#
#
# tomwrap.cfg is a flat text file with one line for every input required by 
# the tomography program, usually empty and containing a linefeed. The only 
# option in this file is the keyword ROTATION which is replaced with the
# rotation specified at the commandline
#
# tomwrap-idl.cfg is a two line text file that calls a specific IDL routine. 
# Keywords in the file are SOURCE, DESTINATION, VELOCITY where SOURCE is 
# replaced automatically with the location of the nv3h files, DESTINATION is
# replaced with the images directory and VELOCITY is replaced with the 
# /velocity keyword if good velocities are detected during the tomography run
import sys, os, time, glob, math
from subprocess import Popen, PIPE, STDOUT
from optparse import OptionParser

usage = "usage: %prog START_ROTATION NUMBER_OF_ROTATIONS [options]"
parser = OptionParser(usage)
parser.add_option("-b", "--basedir",
                  dest="outbase",
                  default='/smeitom/tom',
                  help="Base directory for output files [default: %default]")
parser.add_option("-c", "--cfg",
                  dest="configfiles",
                  default='/smeitom/tom/wrapper',
                  help="Directory containing Tomography and IDL wrapper configuration files [default: %default]")
parser.add_option("-P", "--pro",
                  dest="idlscript",
                  default='/smeitom/tom/pro',
                  help="Directory containing script called from tomwrap.cfg [default: %default]")
parser.add_option("-t", "--timedir",
                  dest="timedir",
                  default=time.strftime('%d-%b-%Y',time.localtime()),
                  help="Subdirectory of base where events and images are written, only use this to overwrite or resume an existing run [default: %default]")
parser.add_option("-T", "--tomo",
                  dest="tomo",
                  action="store_true",
                  default=False,
                  help="Run only tomography, skip IDL, usually used with -t")
parser.add_option("-I", "--idl", 
                  dest="idl", 
                  action="store_true",
                  default=False,
                  help="Run only IDL, skip tomography, usually used with -t")
parser.add_option("-S", "--start", 
                  dest="start", 
                  default=None, 
                  help="Used in conjunction with --idl, specify a carrington rotation (with fractional rotation) or the file index number (if 20 png files exist for a given type of image, then enter 21)")
parser.add_option("-e", "--exe", 
                  dest="exe", 
                  default='/smeitom/tom/exe/smeiipstd0n_intel', 
                  help="Full path to tomography executable [default: %default]")
parser.add_option("-s", "--smei",
                  dest="smei", 
                  default='/smeitom/int/inter_data/base/*.*', 
                  help="smei path for tomography [default: %default]")
parser.add_option("-n", "--nagoya", 
                  dest="nagoya", 
                  default='nagoya,/zshare/dat/nagoya/yearly', 
                  help="nagoya path for tomography [default: %default]")
parser.add_option("-v", "--verbose", 
                  dest="verbose", 
                  action="store_true", 
                  default=False,
                  help="Write tomography/IDL output to the screen [default: %default]")
(options, args) = parser.parse_args()

if len(args) != 2:
    parser.error("Incorrect number of arguments")
if options.idl and options.tomo:
    parser.error("Cannot use -I and -T, pick one, or none")
if not os.path.exists(options.configfiles):
    parser.error("Configuration directory not found: %s" %options.configfiles)
if not os.path.exists(options.outbase):
    parser.error("Base directory not found: %s" % options.outbase)
if not os.path.exists(options.exe):
    parser.error("Tomography executable not found: %s" %options.exe)
if not os.path.exists(options.idlscript):
    parser.error("IDL script directory not found: %s" %options.idlscript)
if options.idl == False and options.start:
    parser.error("Don't use -S/--start without -I/--idl and -t")

rotation = str(args[0])
rotations = int(args[1])
cmd = options.exe + ' nagoya='+options.nagoya+' smei='+options.smei
outbase = options.outbase
configfile = os.path.join(options.configfiles,'tomwrap.cfg')
idlconfigfile = os.path.join(options.configfiles,'tomwrap-idl.cfg')
rotationdir = str(int(math.floor(float(rotation))))
timedir = options.timedir
idlscriptdir = options.idlscript
if not options.tomo:
    if not os.path.exists(os.path.join(outbase,rotationdir,timedir)):
        print glob.glob(os.path.join(outbase,rotationdir)+'/*')
        timedir = filter(os.path.isdir, glob.glob(os.path.join(outbase,rotationdir)+'/*'))[0]
        timedir = os.path.split(timedir)[-1]
print "Tomography/IDL wrapper beginning with the following parameters:"
print "Rotation: %s" % rotation
print "Number of rotations: %s" % rotations
print "Base directory: %s" % outbase
print "Subdirectory for this run: %s" % os.path.join(outbase, rotationdir,timedir)
print "Tomography command: %s" % cmd
print "Tomography configuration: %s" % configfile
print "IDL configuration: %s" % idlconfigfile
for num in range(rotations):
    starttime = time.time()
    rotrun = str(float(rotation)+num)
    rotationdir = str(int(math.floor(float(rotrun))))
    if not options.idl:
        config = open(configfile, 'r')
        writedir = os.path.join(outbase,rotationdir,timedir,'events_'+rotrun+'_HR')
        outputfile = os.path.join(outbase,rotationdir,timedir,'tomoutput-'+timedir+'-'+rotrun+'-tomo.log')
        try:
            os.makedirs(writedir)
        except OSError:
            print 'Events directory exists, continuing anyways.'

        output = open(outputfile, 'a')
        smeiproc = Popen(cmd,stdin=PIPE,stdout=PIPE, cwd=writedir,shell=True)

        print 'Piping tomography config file: ',
        for line in config.readlines():
            smeiproc.stdin.write(line.replace('ROTATION',rotrun))
            output.write(line.replace('ROTATION',rotrun))

        print 'done'

        config.close()
        print 'Tomography starting for rotation: ', rotrun
        settime = 0
        hasvelocity = None
        while not smeiproc.poll():
            timenow = time.time()
            if options.verbose:
                if int(timenow-starttime)%60 == 1 and settime == 0:
                    print '------------------------------'
                    print time.strftime('%d-%b-%Y %H:%M',time.localtime()), 'Still Alive'
                    print '------------------------------'
                    settime = 1
                if int(timenow-starttime)%60 == 50 and settime == 1:
                    settime = 0
            outputval = smeiproc.stdout.readline()
            smeiproc.stdout.flush()
            errorcheck = outputval.find('******* Error ********')
            if errorcheck != -1:
                print "Error: check that all points files exist"
                sys.exit()
            if hasvelocity == None:
                vobs = outputval.find('V observations ')
                if vobs != -1:
                    velocities = outputval[vobs+len('V observations '):]
                    output.write('VELOCITIES FOUND: '+velocities)
                    if int(velocities) != 0:
                        hasvelocity = True
                    else:
                        hasvelocity = False
            if options.verbose:
                print outputval.rstrip()
            output.write(outputval)
            output.flush()

        output.close()
        if hasvelocity:
            print 'Velocities found, renaming directory'
            os.rename(writedir,writedir+'_v')
            writedir = writedir+'_v'
        print 'Finished tomography for rotation:', rotrun
    if not options.tomo:
        if os.path.exists(os.path.join(outbase,rotationdir,timedir,'events_'+rotrun+'_HR_v')):
            hasvelocity = True
            writedir = os.path.join(outbase,rotationdir,timedir,'events_'+rotrun+'_HR_v')
            imagedir = os.path.join(outbase,rotationdir,timedir,'images_'+rotrun+'_HR_v')
        elif os.path.exists(os.path.join(outbase,rotationdir,timedir,'events_'+rotrun+'_HR')):
            hasvelocity = False
            writedir = os.path.join(outbase,rotationdir,timedir,'events_'+rotrun+'_HR')
            imagedir = os.path.join(outbase,rotationdir,timedir,'images_'+rotrun+'_HR')
        else:
            hasvelocity = None
            print 'No Tomography directory found, cannot run IDL from this path'
            sys.exit(1)
        print 'IDL Parameters'
        print 'Velocities  : ', hasvelocity
        print 'Source      : ', writedir
        print 'Destination : ', imagedir
        if options.start and (num == 0):
            print 'Restarting : ', options.start
            filelist = glob.glob(os.path.join(writedir, 'nv3h*'))
            # I hate that globs aren't sorted
            filelist.sort()
            try:
                if len(options.start) > 6:
                    indexval = [str('nv3h'+options.start) in i for i in filelist].index(True)
                else:
                    indexval = options.start
            except ValueError:
                print 'Cannot resume, something wrong with --start %s', options.start
                sys.exit(1)

            filterlist = filelist[indexval:]
            # changes the filtered array into something palatable by IDL's filter='' arg in file_search by
            # replacing the [ brackets with { brackets and removing the single quotes from around the filenames
            # ['somefile', 'anotherfile'] => {somefile,anotherfile}
            # this is very picky, the expansion is basically a regular expression so spaces between comma and element
            # name are not ignored.
            filterlist = repr([os.path.basename(i) for i in filterlist]).replace('[','{').replace(']','}').replace("'", "").replace(' ', '')


        
        idlconfig = open(idlconfigfile, 'r')
        try:
            os.makedirs(imagedir)
        except OSError:
            print 'failed to make images directory: already exists or IO error, continue anyways'

        idlproc = Popen('idl',stdin=PIPE,stdout=PIPE, cwd=idlscriptdir,shell=True)
        idloutputfile = os.path.join(outbase,rotationdir,timedir,'tomoutput-'+timedir+'-'+rotrun+'-idl.log')
        idloutput = open(idloutputfile, 'a')
        time.sleep(5)

        print 'Parsing IDL configuration file: ',
        for line in idlconfig.readlines():
            line = line.replace('SOURCE',writedir)
            line = line.replace('DESTINATION',imagedir)
            if hasvelocity:
                line = line.replace('SPEED', ', /speed')
            else:
                line = line.replace('SPEED', '')
            if options.start and (num == 0):
                line = line.replace('FILTER', ", filter='" + filterlist + "'")
                print 'FILTER: ', line
            else:
                line = line.replace('FILTER', '')
            print line
            idlproc.stdin.write(line)
            idloutput.write(line)

        print 'done'
        idlconfig.close()
        print 'Starting IDL'

        settime = 0
        while not idlproc.poll():
            timenow = time.time()
            if options.verbose:
                if int(timenow-starttime)%60 == 1 and settime == 0:
                    print '------------------------------'
                    print time.strftime('%d-%b-%Y %H:%M',time.localtime()), 'IDL Still Alive'
                    print '------------------------------'
                    settime = 1
                if int(timenow-starttime)%60 == 50 and settime == 1:
                    settime = 0
            outputval = idlproc.stdout.readline()
            if options.verbose:
                print outputval.rstrip()
            idlproc.stdout.flush()
            if outputval.strip() != '':
                idloutput.write(outputval)
                idloutput.flush()
            else:
                if len(outputval) == 0:
                    print 'Exiting, either IDL finished or broke.'
                    if options.verbose:
                        print "Reason: ",len(outputval),",",outputval,",",repr(outputval)
                    break
                if options.verbose:
                    time.sleep(1)
        idlproc.wait()
        idloutput.close()
        print 'IDL program finished for rotation: ', rotrun
