#
# Copyright (C) 2013-2018  Leo Singer
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation; either version 2 of the License, or (at your
# option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General
# Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
#
r"""
Listen for new events from LVAlert and perform sky localization.

`bayestar_localize_lvalert` supports two modes of operation. You can
explicitly specify the GraceDb ID on the command line, as in:

    $ bayestar_localize_lvalert T90713

Or, `bayetar_localize_lvalert` can GraceDB IDs from stdin (e.g., from the
terminal, or redirected from a fifo):

    $ mkfifo /var/run/bayestar
    $ tail -F /var/run/bayestar | bayestar_localize_lvalert &
    $ echo T90713 > /var/run/bayestar
"""

from .. import command


def parser():
    parser = command.ArgumentParser(
        parents=[
            command.waveform_parser, command.prior_parser, command.mcmc_parser])
    parser.add_argument(
        '-d', '--disable-detector', metavar='X1', type=str, nargs='+',
        help='disable certain detectors')
    parser.add_argument(
        '-N', '--dry-run', action='store_true',
        help='Dry run; do not update GraceDB entry')
    parser.add_argument(
        '--no-tag', action='store_true',
        help='Do not set lvem tag for GraceDB entry')
    parser.add_argument(
        '-o', '--output', metavar='FILE.fits[.gz]', default='bayestar.fits.gz',
        help='Name for uploaded file')
    parser.add_argument(
        'graceid', metavar='G123456', nargs='*',
        help='Run on these GraceDB IDs. If no GraceDB IDs are listed on the '
        'command line, then read newline-separated GraceDB IDs from stdin.')
    return parser


def main(args=None):
    opts = parser().parse_args(args)

    import logging
    import os
    import re
    import sys
    from ..bayestar import localize, rasterize
    from ..io import fits
    from ..io import events
    from ..util.file import rename, TemporaryDirectory
    import ligo.gracedb.logging
    import ligo.gracedb.rest

    # Squelch annoying and uniformative LAL log messages.
    import lal
    lal.ClobberDebugLevel(lal.LALNDEBUG)

    logging.basicConfig(format='%(asctime)s %(levelname)s %(message)s')
    log = logging.getLogger('BAYESTAR')

    # If no GraceDB IDs were specified on the command line, then read them
    # from stdin line-by-line.
    graceids = opts.graceid if opts.graceid else command.iterlines(sys.stdin)

    # Fire up a GraceDb client
    # FIXME: Mimic the behavior of the GraceDb command line client, where the
    # environment variable GRACEDB_SERVICE_URL overrides the default service URL.
    # It would be nice to get this behavior into the gracedb package itself.
    gracedb = ligo.gracedb.rest.GraceDb(
        os.environ.get(
            'GRACEDB_SERVICE_URL', ligo.gracedb.rest.DEFAULT_SERVICE_URL))

    if opts.chain_dump:
        chain_dump = re.sub(r'.fits(.gz)?$', r'.hdf5', opts.output)
    else:
        chain_dump = None

    tags = ("sky_loc",)
    if not opts.no_tag:
        tags += ("lvem",)

    event_source = events.gracedb.open(graceids, gracedb)

    if opts.disable_detector:
        event_source = events.detector_disabled.open(
            event_source, opts.disable_detector)

    for graceid in event_source.keys():

        try:
            event = event_source[graceid]
        except:
            log.exception('failed to read event %s from GraceDB', graceid)
            continue

        # Send log messages to GraceDb too
        if not opts.dry_run:
            handler = ligo.gracedb.logging.GraceDbLogHandler(gracedb, graceid)
            handler.setLevel(logging.INFO)
            logging.root.addHandler(handler)

        # A little bit of Cylon humor
        log.info('by your command...')

        try:
            # perform sky localization
            log.info("starting sky localization")
            sky_map = rasterize(localize(
                event, opts.waveform, opts.f_low, opts.min_distance,
                opts.max_distance, opts.prior_distance_power, opts.cosmology,
                mcmc=opts.mcmc, chain_dump=chain_dump,
                enable_snr_series=opts.enable_snr_series,
                f_high_truncate=opts.f_high_truncate))
            sky_map.meta['objid'] = str(graceid)
            sky_map.meta['url'] = 'https://gracedb.ligo.org/events/{0}'.format(
                graceid)
            log.info("sky localization complete")

            # upload FITS file
            with TemporaryDirectory() as fitsdir:
                fitspath = os.path.join(fitsdir, opts.output)
                fits.write_sky_map(fitspath, sky_map, nest=True)
                log.debug('wrote FITS file: %s', opts.output)
                if opts.dry_run:
                    rename(fitspath, os.path.join('.', opts.output))
                else:
                    gracedb.writeLog(
                        graceid, "BAYESTAR rapid sky localization ready",
                        filename=fitspath, tagname=tags)
                log.debug('uploaded FITS file')
        except KeyboardInterrupt:
            # Produce log message and then exit if we receive SIGINT (ctrl-C).
            log.exception("sky localization failed")
            raise
        except:
            # Produce log message for any otherwise uncaught exception.
            # Unless we are in dry-run mode, keep going.
            log.exception("sky localization failed")
            if opts.dry_run:
                # Then re-raise the exception if we are in dry-run mode
                raise

        if not opts.dry_run:
            # Remove old log handler
            logging.root.removeHandler(handler)
            del handler
