#!/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_inspiral call. This code takes directly the output file and frame locations and translates these to lalapps_inspiral.')

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 inspiral 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 inspiral but needed to figure out file naming")
parser.add_argument("--channel-name", action="store", required=True,
                     metavar="STR", help="Channel name. Sent straight to inspiral 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_inspiral to make a random ifotag to guarantee uniqueness
ifoTag = randStr

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

# Pipedown also wants us to store the usertag. This is not robust, but should
# work until we can do this a better way.
try:
    userTag = os.path.basename(args.output_file).split('-')[1]
    userTag = userTag.split('_')[1:]
    if userTag[0] == 'FULL' and userTag[1] == 'DATA':
        userTag = 'FULL_DATA'
    elif userTag[0] == 'PLAYGROUND':
        userTag = 'PLAYGROUND'
    elif userTag[0].endswith("INJ"):
        userTag = userTag[0]
    else:
        userTag = '_'.join(userTag)
except:
    userTag = None

# 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

if userTag is not None:
    inspOutName = "%s-INSPIRAL_%s_%s-%d-%d" \
              %(ifo, ifoTag, userTag, startTime, duration)
else:
    inspOutName = "%s-INSPIRAL_%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_inspiral']
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]

if userTag is not None:
    cmdList += ['-userTag', userTag]

cmdList += leftovers

# Run lalapps_inspiral
retCode = subprocess.call(cmdList, stderr=sys.stderr, stdout=sys.stdout,
                          shell=False)
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)
