#!/home/conda/feedstock_root/build_artifacts/bld/rattler-build_pycbc_1768936173/host_env_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehold_placehol/bin/python

# Copyright (C) 2015-2023 Ian Harry, Gareth Cabourn Davies
#
# 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 3 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.

"""
Take a coinc xml file containing a single event and upload to gracedb.
"""

import os
import sys
import argparse
import logging
from ligo.gracedb.rest import GraceDb

import pycbc
from pycbc.io.gracedb import gracedb_tag_with_version

parser = argparse.ArgumentParser(description=__doc__)
pycbc.add_common_pycbc_options(parser)
parser.add_argument('--xml-file-for-upload', required=True, type=str,
                    help='LIGOLW XML file containing the event.')
parser.add_argument('--log-message', type=str, metavar='MESSAGE',
                    help='Add a log entry to each upload with the given message')
parser.add_argument('--testing', action="store_true", default=False,
                    help="Upload event to the TEST group of gracedb.")
parser.add_argument('--production-server', action="store_true", default=False,
                    help="Upload event to production graceDB. If not given "
                         "events will be uploaded to playground server.")
parser.add_argument('--search-tag', default='AllSky',
                    help="Specify the search tag. Default: AllSky")
parser.add_argument('--snr-timeseries-plot',
                    help="SNR timeseries plot to be uploaded with the event.")
parser.add_argument('--asd-plot',
                    help="ASD plot to be uploaded with the event.")
parser.add_argument('--skymap-plot',
                    help="Skymap plot to be uploaded with the event.")
parser.add_argument('--skymap-fits-file',
                    help="Skymap .fits file to be uploaded with the event.")
parser.add_argument('--source-probabilities',
                    help="Source probabilities file to be uploaded with the "
                         "event.")
parser.add_argument('--source-probabilities-plot',
                    help="Source probabilities plot to be uploaded with the "
                          "event.")
parser.add_argument('--labels', nargs='+',
                    help="Labels to add to the event in GraceDB")

args = parser.parse_args()

# Default logging level is info: --verbose adds to this
pycbc.init_logging(args.verbose, default_level=1)
# Make the scitokens logger a little quieter
# (it is called through GraceDB)
logging.getLogger('scitokens').setLevel(logging.root.level + 10)

if args.production_server:
    gracedb = GraceDb()
else:
    gracedb = GraceDb(service_url='https://gracedb-playground.ligo.org/api/')

labels = [l.upper() for l in (args.labels or [])]
allowed_labels = gracedb.allowed_labels

if set(labels) - set(allowed_labels):
    err_msg = "One or more supplied labels is not available on the server. "
    err_msg += f"Supplied {','.join(labels)}, allowed "
    err_msg += f"{','.join(allowed_labels)}."
    raise RuntimeError(err_msg)

group_tag = 'Test' if args.testing else 'CBC'
r = gracedb.create_event(
    group_tag,
    'pycbc',
    args.xml_file_for_upload,
    filecontents=open(args.xml_file_for_upload, "rb").read(),
    search=args.search_tag,
    offline=True,
    labels=labels
).json()

logging.info("Uploaded event %s.", r["graceid"])

# add info for tracking code version
gracedb_tag_with_version(gracedb, r['graceid'])

# document the absolute path to the input file
input_file_str = 'Candidate uploaded from ' \
    + os.path.abspath(args.xml_file_for_upload)
gracedb.write_log(r['graceid'], input_file_str)

# document the command line used in the event log
log_str = 'Upload command: ' + ' '.join(sys.argv)

# add the custom log message, if provided
if args.log_message is not None:
    log_str += '. ' + args.log_message

gracedb.write_log(
    r['graceid'],
    log_str,
    tag_name=['analyst_comments']
)


def upload_file(upload_filename, displayname, comment, tag):
    """
    Helper function to upload files associated with the event.
    """
    logging.info("Uploading %s file %s to event %s.",
                 displayname, upload_filename, r["graceid"])
    gracedb.write_log(
        r["graceid"],
        comment,
        filename=upload_filename,
        tag_name=[tag],
        displayName=[displayname]
    )


if args.asd_plot:
    upload_file(
        args.asd_plot,
        "ASDs",
        "PyCBC ASD estimate from the time of event",
        "psd"
    )

if args.snr_timeseries_plot:
    upload_file(
        args.snr_timeseries_plot,
        "SNR timeseries",
        "SNR timeseries plot upload",
        "background"
    )

if args.skymap_plot:
    upload_file(
        args.skymap_plot,
        "",
        "Skymap plot upload",
        "sky_loc"
    )

if args.skymap_fits_file:
    upload_file(
        args.skymap_fits_file,
        "",
        "sky localization complete",
        "sky_loc"
    )

if args.source_probabilities:
    upload_file(
        args.source_probabilities,
        "",
        "Source probabilities JSON file upload",
        "pe"
    )

if args.source_probabilities_plot:
    upload_file(
        args.source_probabilities_plot,
        "",
        "Source probabilities plot upload",
        "pe"
    )

logging.info('Done!')
