#!/usr/bin/env python

import sys,os
import atexit
import string
import random
import argparse
import shutil
import subprocess

import glue.lal
from glue.ligolw import dbtables

def file_remove_try(file):
    try:
        os.remove(file)
    except:
        pass

parser = argparse.ArgumentParser(description='Convert pegasus-proof command line call to an actual lalapps_tmplbankl call. This code takes directly the output file and frame locations and translates these to lalapps_tmplbank.')

parser.add_argument("--output-file", action="store", required=True, 
                    metavar="OUTFILE", help="Path to output file.")
parser.add_argument("--frame-files", nargs="*", action="store", required=True,
                    metavar="FRAMEFILE",
                    help="List of frame files to be used for analysis.")
parser.add_argument("--gps-start-time", action="store", required=True,
                    metavar="TIME", type=int,
                    help="Start time. Sent straight to tmplbank but needed to figure out file naming")
parser.add_argument("--gps-end-time", action="store", required=True,
                    metavar="TIME", type=int,
                    help="End time. Sent straight to tmplbank but needed to figure out file naming")
parser.add_argument("--channel-name", action="store", required=True,
                     metavar="STR", help="Channel name. Sent straight to tmplbank but needed to figure out file naming")

# Need to add signal handlers to deal with condor remove commands
dbtables.install_signal_trap()

# Leftovers are all the other options that will just be sent to lalapps_inspiral
args, leftovers = parser.parse_known_args()

# Choose a random string for file naming

randStr = ''.join([random.choice(string.ascii_uppercase) for x in range(10)])

# Let's start with the frame files, these need to be translated to a frame
# cache file, that will be written to local directory
if (len(args.frame_files) == 1) and (args.frame_files[0].endswith('.lcf')):
    # Assume I was given a frame file as input and use that.
    frameFileName = args.frame_files[0]
else:
    frameCache = glue.lal.Cache.from_urls(args.frame_files)
    frameFileName = '%s.lcf' %(randStr)
    atexit.register(file_remove_try, frameFileName)
    with open(frameFileName, 'w') as fP:
        frameCache.tofile(fP)


# Just get lalapps_tmpltbank to make a random ifotag to guarantee uniqueness
ifoTag = randStr

# Add this to arguments to send to tmpltbank
leftovers.extend(['--ifo-tag', ifoTag])

# Also need to know the temporary output file name
ifo = args.channel_name[0:2]
startTime = args.gps_start_time
duration = args.gps_end_time - startTime
inspOutName = "%s-TMPLTBANK_%s-%d-%d" %(ifo, ifoTag, startTime, duration)
atexit.register(file_remove_try, inspOutName)

if "--write-compress" in leftovers:
    inspOutName += ".xml.gz"
else:
    inspOutName += ".xml"

cmdList = ['lalapps_tmpltbank']
cmdList += ['--gps-start-time', str(args.gps_start_time), '--gps-end-time',
           str(args.gps_end_time), '--channel-name', args.channel_name,
           '--ifo-tag', ifoTag, '--frame-cache', frameFileName]
cmdList += leftovers

# Run lalapps_tmpltbank
retCode = subprocess.call(cmdList, stderr=sys.stderr, stdout=sys.stdout,
          shell=False)
os.remove(frameFileName)
if not retCode:
    # Move the output and exit
    shutil.move(inspOutName, args.output_file)
    sys.exit(0)
else:
    # Cleanup handled by atexit
    sys.exit(retCode)