#!/usr/bin/env python

"""
Simple wrapper around bayestar-localize-coincs to get sky localization for a
specific compact binary merger event.

Uses an XML containing SNR timeseries and PSD and calls BAYESTAR to produce a
sky localization.

The wrapper needs to exist as the XML format doesnt allow us to indicate
the low frequency cutoff or waveform used. We also want to be able to control
the output filename which is not possible in bayestar-localize-coincs

Unrecognised options will be passed straight to the bayestar subprocess
"""

import argparse
import subprocess
import shutil
import logging
import os
import glob
import sys
import random
import tempfile
import pycbc.version
from pycbc import init_logging

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--version', action="version",
                    version=pycbc.version.git_verbose_msg)
parser.add_argument('--verbose', action='count',
                    help="Make the logging increasingly more verbose")
parser.add_argument('--bayestar-executable',
                    help="The bayestar-localize-coinc executable to be run. "
                         "If not given, will use whatever is available in "
                         "the current environment.")
parser.add_argument('--event-xml', required=True,
                    help="XML file containing event information, SNR "
                         "timeseries and PSD to pass to bayestar")
parser.add_argument('--waveform', default='TaylorF2',
                    help="Waveform used in the matched-filtering "
                         "to generate the SNR timeseries.")
parser.add_argument('--low-frequency-cutoff', type=float, default=20,
                    help="Low-frequency cutoff used in the matched-filtering "
                         "to generate the SNR timeseries")
parser.add_argument('--output-file', required=True,
                    help="Filename to output the fits file to.")
args, unknown = parser.parse_known_args()

# Default logging is set higher than normal for this job
logging_level = args.verbose + 1 if args.verbose else None
init_logging(logging_level)
logging.info("Starting")

bayestar_exe = args.bayestar_executable or 'bayestar-localize-coincs'

tmpdir = tempfile.mkdtemp()

# Set up the command to pass to bayestar.
# The XML file path being passed twice is a legacy requirement, not a mistake.
cmd = [bayestar_exe,
       args.event_xml,
       args.event_xml,
       '--waveform', args.waveform,
       '--f-low', str(args.low_frequency_cutoff),
       '-o', tmpdir]

# Pass any unknown options straight to the subprocess
cmd += unknown

logging.info("Running %s", ' '.join(cmd))
subprocess.check_output(cmd)

# Find the fits file in the temporary directory:
# It would be nice to do this better! - maybe use the input xml
# file to find the number which is going to be used?
fits_filenames = glob.glob(os.path.join(tmpdir, '*.fits*'))
if len(fits_filenames) != 1:
    raise ValueError(f'argh, got {len(fits_filenames)} FITS files after running BAYESTAR, I want one!')

logging.info("Moving output to %s", args.output_file)
shutil.move(fits_filenames[0], args.output_file)

logging.info("Done")
