#!/usr/bin/env python
#
#
# Copyright (C) 2009  Larne Pekowsky
#
# 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.


#
# =============================================================================
#
#                                   Preamble
#
# =============================================================================
#


"""
Takes one or more segment XML files and a query file, runs the query against
the segments, and produces veto files as XML
"""


from optparse import OptionParser

try:
    import sqlite3
except ImportError:
    # pre 2.5.x
    from pysqlite2 import dbapi2 as sqlite3

import sys
import os
import pwd
import tempfile
import operator
import pycbc_glue.segments

import time
from pycbc_glue import gpstime

from pycbc_glue.ligolw import ligolw
from pycbc_glue.ligolw import lsctables
from pycbc_glue.ligolw import utils
from pycbc_glue.ligolw.utils import ligolw_sqlite
from pycbc_glue.ligolw import dbtables

from pycbc_glue.segmentdb import query_engine
from pycbc_glue.segmentdb import segmentdb_utils

from pycbc_glue.ligolw.utils import process

from pycbc_glue import git_version

import re
import urllib

PROGRAM_NAME = sys.argv[0].replace('./','')
PROGRAM_PID  = os.getpid()
try:
        USER_NAME = os.getlogin()
except:
        USER_NAME = pwd.getpwuid(os.getuid())[0]


__author__ = "Larne Pekowsky <lppekows@physics.syr.edu>"
__version__ = "git id %s" % git_version.id
__date__ = git_version.date


#
# =============================================================================
#
#                                 Command Line
#
# =============================================================================
#


def parse_command_line():
    """
    Parse the command line, return an options object
    """

    parser = OptionParser(
        version = "Name: %%prog\n%s" % git_version.verbose_msg,
        usage   = "%prog -v|--veto-file filename [options]",
        description = "Reads one or more segment files and a veto file and generates files of veto segments"        
	)
    
    parser.add_option("-v", "--veto-file",    metavar = "veto_file",    help = "veto XML file (required).")
    parser.add_option("-o", "--output-dir",   metavar = "output_dir",   default = '.',          help = "Directory to write output (default=cwd).")
    parser.add_option("-k", "--keep-db",      metavar = "keep_db",      action  = "store_true", help = "Keep sqlite database.")
    parser.add_option("-r", "--remove-definer-table",    action = "store_true", help = "If set the veto definer table will not be included in the output.")

    parser.add_option("-t", "--segment-url",  metavar = "segment_url", help = "Segment URL")
    parser.add_option("-d", "--database",     metavar = "use_database", action = "store_true", help = "use database specified by environment variable S6_SEGMENT_SERVER")
    parser.add_option("-f", "--dmt-file",     metavar = "use_files", action = "store_true", help = "use files in directory specified by environment variable ONLINEDQ")
    parser.add_option("-c", "--cumulative-categories",   action = "store_true", help = "If set the category N files will contain all segments in categories <= N")
    parser.add_option("-p", "--separate-categories",     action = "store_true", help = "If set the category N files will contain only category N")
    parser.add_option("-i", "--individual-results",      action = "store_true", help = "If set include information on each veto individually.")
    parser.add_option("-C", "--veto-categories", action="store", default=None, help="A comma separated list of all veto categories to consider. If not given then consider all categories in the veto-definer file.")
    parser.add_option("-I", "--ifo-list", action="store", default=None, help="A comma separated list of all ifos to consider. If not given then consider all ifos in the veto-definer file.")

    # Time options
    parser.add_option("-s", "--gps-start-time", metavar = "gps_start_time", help = "Start of GPS time range")
    parser.add_option("-e", "--gps-end-time",   metavar = "gps_end_time", help = "End of GPS time range")

    options, others = parser.parse_args()

    if not options.veto_file:
        raise ValueError, "missing required argument --veto-file"
   
    # User must specify how to treat the categories
    if not (options.cumulative_categories or options.separate_categories) or (options.cumulative_categories and options.separate_categories):
        print >>sys.stderr, "Must provide one of --cumulative-categories | --separate-categories"
        sys.exit(-1)

    tmp_dir = None

    # If the veto file is a URL, load it
    loaded_file = None

    if options.veto_file.find(':') != -1:
        tmp_dir = tempfile.mkdtemp()
        inurl   = urllib.urlopen(options.veto_file)

        # Grab the part of the name after the last slash
        pos     = options.veto_file[::-1].find('/')
        fname   = (pos > -1) and options.veto_file[ -1 * pos:] or "veto_definer.xml" 

        loaded_file = tmp_dir + "/" + fname

        outfile = open(loaded_file, 'w')
        for l in inurl:
            print >>outfile, l,

        inurl.close()
        outfile.close()

    # Make sure we have required arguments
    database_location = None
    file_location     = None

    # Make sure we know who to contact for data
    # (This is redundant with ligolw_segment_query, maybe move to
    # query_engine)
    if options.segment_url:
        if options.segment_url.startswith('ldbd') or options.segment_url.startswith('http'):
            database_location = options.segment_url
        elif options.segment_url.startswith('file:'):
            file_location = options.segment_url[len('file://'):]
        else:
            tmp_dir = tempfile.mkdtemp()
        
            # Grab the part of the name after the last slash
            pos     = options.segment_url[::-1].find('/')
            fname   = (pos > -1) and options.segment_url[ -1 * pos:] or "dmt.xml" 

            inurl   = urllib.urlopen(options.segment_url)
            outfile = open(tmp_dir + "/" + fname, 'w')
            for l in inurl:
                print >>outfile, l,

            inurl.close()
            outfile.close()
            file_location = tmp_dir
    elif options.database:
        if 'S6_SEGMENT_SERVER' not in os.environ:
            raise ValueError( "--database specified by S6_SEGMENT_SERVER not set" )
        database_location = os.environ['S6_SEGMENT_SERVER']
    elif options.dmt_file:
        if 'ONLINEDQ' not in os.environ:
            raise ValueError( "--dmt-file specified but ONLINEDQ not set" )

        tmp = os.environ['ONLINEDQ']
        if tmp.startswith('file://'):
            tmp = tmp[len('file://'):]
        file_location = tmp
    else:
        raise ValueError( "One of [ --segment_url | --database | --dmt-file ] must be provided" )


    if not options.gps_start_time:
        raise ValueError( "missing required argument --gps-start-time" )
    
    if not options.gps_end_time:
        raise ValueError( "missing required argument --gps-end-time" )

        
    return options, database_location, file_location, loaded_file




#
# =============================================================================
#
#                                 Library Code
#
# =============================================================================
#

#
# Do to some sort of version skew in sqlite at LHO dbtables.get_column_info
# doesn't work if and only if a new table has been created, which we do
# in order to make the segment_backup table.  This method is a replcement
# that works in this case.
# TODO: once the environment is fixed, remove this.
#
def kludge_get_column_info(connection, table_name):
    """
    Return an in order list of (name, type) tuples describing the
    columns in the given table.
    """
    statement = None

    table_name_str = str(table_name)

    for name,sql in connection.cursor().execute("SELECT name, sql FROM sqlite_master WHERE type == 'table'"):
        if str(name) == table_name_str:
            statement = sql

    coldefs = re.match(dbtables._sql_create_table_pattern, statement).groupdict()["coldefs"]
    return [(coldef.groupdict()["name"], coldef.groupdict()["type"]) for coldef in re.finditer(dbtables._sql_coldef_pattern, coldefs) if coldef.groupdict()["name"].upper() not in ("PRIMARY", "UNIQUE", "CHECK")]

class ContentHandler(ligolw.LIGOLWContentHandler):
    connection = None


#
# =============================================================================
#
#                                     Main
#
# =============================================================================
#

if __name__ == '__main__':
    # Settings

    dbtables.get_column_info = kludge_get_column_info

    options, db_location, file_location, loaded_file = parse_command_line()    

    now = gpstime.GpsSecondsFromPyUTC(time.time())

    # 1. Load the veto file into sqlite    
    handle, temp_db    = tempfile.mkstemp(suffix='.sqlite')
    os.close(handle)
    target          = dbtables.get_connection_filename(temp_db, None, True, False)
    veto_connection = sqlite3.connect(target)
    dbtables.idmap_sync(veto_connection)


    # Create an engine to query veto information
    veto_query_engine = query_engine.SqliteQueryEngine(veto_connection)
    xml_files         = [loaded_file or options.veto_file]

    # 2. If we're working with DMT files load them into sqlite as well
    min_start_time = int(options.gps_start_time)
    max_end_time   = int(options.gps_end_time)
    
    if file_location:
        xml_files += segmentdb_utils.get_all_files_in_range(file_location, min_start_time, max_end_time)
        segment_query_engine = veto_query_engine
    else:
        segment_connection   = segmentdb_utils.setup_database(db_location)
        segment_query_engine = query_engine.LdbdQueryEngine(segment_connection)

    ContentHandler.connection = veto_connection
    dbtables.use_in(ContentHandler)
    ligolw_sqlite.insert_from_urls(xml_files, ContentHandler)

    # If we're working with DMT files make sure we have a segment table
    if file_location:
        segmentdb_utils.ensure_segment_table(veto_connection)
    
    cli_interval = pycbc_glue.segments.segmentlist([pycbc_glue.segments.segment(int(options.gps_start_time), int(options.gps_end_time))])
            
    #
    # Get the set of veto categories and ifos for which we have segments 
    #
    if options.veto_categories:
        categories=[int(x) for x in options.veto_categories.split(",")]
    else:
        categories = [row[0] for row in veto_query_engine.query("select distinct(category) from veto_definer")]
    if options.ifo_list:
        ifos = [x for x in options.ifo_list.split(",")]
    else:
        ifos = [row[0] for row in veto_query_engine.query("select distinct(ifo) from veto_definer")]

    category_selection_operator = options.cumulative_categories and ' <= ' or ' = '

    for ifo in ifos:
        for category_num in categories:
            # Start a new document
            doc = ligolw.Document()
            doc.appendChild(ligolw.LIGO_LW())

            # Copy veto definer's process table
            doc.childNodes[0].appendChild(dbtables.get_xml(veto_connection, ['process']).childNodes[0])

            # Register ourselves
            proc_id = process.register_to_xmldoc(doc, PROGRAM_NAME, options.__dict__, version=git_version.id).process_id

            # Keep the veto_definer for tracking purposes
            if not options.remove_definer_table:
                doc.childNodes[0].appendChild(dbtables.get_xml(veto_connection, ['veto_definer']).childNodes[0])

            # loop over select of name, version where ifo=current ifo and category <= current category
            cursor = veto_connection.cursor()
            sql = """SELECT ifo, name, version, start_time, end_time, start_pad, end_pad
                       FROM veto_definer
                      WHERE ifo = '%s' AND category %s %s""" % (ifo, category_selection_operator, category_num)

            segdefs = []
            for ifo, name, version, start_time, end_time, start_pad, end_pad in cursor.execute(sql):
                if end_time == 0:
                    end_time = now

                veto_interval  = pycbc_glue.segments.segmentlist([pycbc_glue.segments.segment(start_time, end_time)])
                veto_interval.coalesce()
                veto_interval &= cli_interval


                if len(veto_interval) > 0:
                    segdefs.append( (ifo, name, version, veto_interval[0][0], veto_interval[0][1], start_pad, end_pad) )


            vetoed_segments   = segmentdb_utils.query_segments(segment_query_engine, 'segment', segdefs)
            segment_summaries = segmentdb_utils.query_segments(segment_query_engine, 'segment_summary', segdefs)

            # We could write out everything we found.  This might help tracking down why a time
            # or trigger was vetoed.
            if options.individual_results:
              segmentdb_utils.add_segment_info(doc, proc_id, segdefs, vetoed_segments, segment_summaries)

            # Form the result
            vetoed_segments = reduce(operator.or_, vetoed_segments).coalesce()

            # Add the result type to the segment definer table
            seg_name   = 'VETO_CAT%d%s' % (category_num, options.cumulative_categories and '_CUMULATIVE' or '')
            seg_def_id = segmentdb_utils.add_to_segment_definer(doc, proc_id, ifo, seg_name, 1)

            # and segment summary
            segmentdb_utils.add_to_segment_summary(doc, proc_id, seg_def_id, [[min_start_time, max_end_time]])

            # and store the segments
            segmentdb_utils.add_to_segment(doc, proc_id, seg_def_id, vetoed_segments)
            
            # Dump to XML
            output_name = "%s/%s-VETOTIME_CAT%d-%d-%d.xml" % (options.output_dir, ifo, category_num, min_start_time, (max_end_time - min_start_time))
            utils.write_filename(doc, output_name)

    # Clean up
    if loaded_file:
        os.remove(loaded_file)
        pos   = loaded_file[::-1].find('/') + 1
        os.rmdir(loaded_file[:-1 * pos])


    if not options.keep_db:
        os.remove(temp_db)

