#!/usr/bin/env python

import os
import pwd
import logging
import json
import argparse
import shutil
try:
    from urllib.request import urlopen
except ImportError:   # python < 3
    from urllib2 import urlopen

import ligo.segments
from glue import git_version
from glue.ligolw import ligolw
from glue.ligolw import utils
from glue.ligolw.utils import process
from glue.segmentdb import segmentdb_utils
from dqsegdb import clientutils

# Logging formatting from pycbc optimal snr
log_fmt = '%(asctime)s %(message)s'
log_date_fmt = '%Y-%m-%d %H:%M:%S'
logging.basicConfig(level=logging.INFO, format=log_fmt,
                    datefmt=log_date_fmt)


# Function to query json segment data from LOSC
def query_losc(ifo, segment_name, gps_start_time, duration):
    """
    Function that queries the O1 LOSC data from json to xml

    Parameters
    ----------
    ifo: string
        The interferometer to query (H1, L1).
    segment_name: string
        The veto group or science group to query from LOSC.
    gps_start_time: int / string
        The starting gps time to begin querying from the O1 LOSC data set.
    duration: int / string
        The amount of time in seconds after the gps start time.

    Returns
    ---------
    segment_list :  ligo.segments.segmentlist
        The inverval returned by LOSC
    segment_summary :  ligo.segments.segmentlist
        The segments returned by LOSC
    """

    response = urlopen(
        'https://www.gw-openscience.org/timeline/segments/json/O1/{}_{}/{}/{}/'.format(
        ifo, segment_name, gps_start_time, duration))

    logging.info(response.info())
    json_segment_data = json.loads(response.read())

    summary_segment = ligo.segments.segmentlist([ligo.segments.segment(
                                                json_segment_data['start'],
                                                json_segment_data['end'])])

    segments = ligo.segments.segmentlist([ligo.segments.segment(
                                         x[0],x[1]) for x in json_segment_data['segments']])

    return summary_segment, segments


def write_xml_file(ifo, summary_segment, segments, filename):
    version = 1

    PROGRAM_NAME = 'pycbc_losc_segment_query'
    PROGRAM_PID  = os.getpid()
    USER_NAME = pwd.getpwuid(os.getuid())[0]

    __author__  = "Duncan Brown <dabrown@syr.edu>"
    __version__ = "git id %s" % git_version.id
    __date__ = git_version.date

    doc = ligolw.Document()
    doc.appendChild(ligolw.LIGO_LW())
    process_id = process.register_to_xmldoc(doc, PROGRAM_NAME, {},
                                            version = git_version.id,
                                            cvs_entry_time = __date__,
                                            comment='LOSC segments').process_id

    seg_def_id = segmentdb_utils.add_to_segment_definer(doc, process_id, ifo,
                                                        'RESULT', 1, comment='LOSC query result')

    clientutils.add_to_segment_summary_ns(doc, process_id, seg_def_id,
                                          summary_segment,
                                          comment='start and end time from losc query')

    clientutils.add_to_segment_ns(doc, process_id, seg_def_id, segments)

    utils.write_filename(doc, filename, gz=False)



parser = argparse.ArgumentParser()
parser.add_argument('--gps-start-time', type=int)
parser.add_argument('--gps-end-time', type=int)
parser.add_argument('--query-segments', action='store_true')
parser.add_argument('--segment-url', type=str)
parser.add_argument('--include-segments', type=str)
parser.add_argument('--output-file', type=str)
parser.add_argument('--protract-hw-inj', type=int, default=0)
args = parser.parse_args()

gps_start_time = args.gps_start_time
gps_end_time = args.gps_end_time
duration = gps_end_time - gps_start_time

logging.info("Reading in LOSC files from {} to {}.".format(gps_start_time,
                                                           gps_end_time))
detector=args.include_segments.split(':')[0]
logging.info("Querying for {}".format(detector))

file_list = []

logging.info("Querying science segments")
sci_summ, sci_segs = query_losc(detector, "DATA", gps_start_time, duration)
sci_segs.coalesce()

sci_file_name = "{}-SCIENCE_SEGMENTS.xml".format(detector)
write_xml_file(detector, sci_summ, sci_segs, sci_file_name)
file_list.append(sci_file_name)

logging.info("Calculating CAT1 veto time")
not_cat1_summ, not_cat1_segs = query_losc(detector, "CBC_CAT1", gps_start_time, duration)
not_cat1_segs.coalesce()

cat1_segs = ~not_cat1_segs
cat1_segs &= sci_segs

cat1_file_name = "{}-VETOTIME_CAT1-{}-{}.xml".format(detector,
                                                     gps_start_time, duration)
write_xml_file(detector, not_cat1_summ, cat1_segs, cat1_file_name)
file_list.append(cat1_file_name)

logging.info("Calculating CAT2 veto time")
not_cat2_summ, not_cat2_segs = query_losc(detector, "CBC_CAT2", gps_start_time, duration)
not_cat2_segs.coalesce()

cat2_segs = ~not_cat2_segs
cat2_segs &= sci_segs

cat2_file_name = "{}-VETOTIME_CAT2-{}-{}.xml".format(detector,
                                                     gps_start_time, duration)
write_xml_file(detector, not_cat2_summ, cat2_segs, cat2_file_name)
file_list.append(cat2_file_name)

logging.info("Calculating HW injection veto time")
not_hw_inj_summ, not_hw_inj_segs = query_losc(detector, "NO_CBC_HW_INJ", gps_start_time, duration)
not_hw_inj_segs.coalesce()

hw_inj_segs = ~not_hw_inj_segs
hw_inj_segs.protract(args.protract_hw_inj)
hw_inj_segs.coalesce()
hw_inj_segs &= sci_segs

hw_inj_file_name = "{}-VETOTIME_CAT3-{}-{}.xml".format(detector,
                                                       gps_start_time, duration)
write_xml_file(detector, not_hw_inj_summ, hw_inj_segs, hw_inj_file_name)
file_list.append(hw_inj_file_name)

destination_path = os.path.dirname(os.path.abspath(args.output_file))

for f in file_list:
    d = os.path.join(destination_path,f)
    logging.info("Copying {} to {}".format(f,d))
    shutil.copy2(f, os.path.join(destination_path,f))
    os.unlink(f)

logging.info("Science and Veto files written. Done.")