#!/usr/bin/env python

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

from optparse import OptionParser
try:
    import sqlite3
except ImportError:
    # pre 2.5.x
    from pysqlite2 import dbapi2 as sqlite3
import sys
import os
import copy

from pylal import ligolw_sqlutils as sqlutils
from pylal import ligolw_cbc_compute_durations as compute_dur

from glue import git_version
from glue import segments
from glue import segmentsUtils
from glue.ligolw import table
from glue.ligolw import lsctables
from glue.ligolw import dbtables
from glue.ligolw import utils
from glue.ligolw.utils import process

__prog__ = "ligolw_cbc_compute_durations"
__author__ = "Collin Capano <cdcapano@physics.syr.edu>"

description = \
"Computes durations for every row in the experiment_summary table in a " + \
"database and stores them."

# =============================================================================
#
#                                   Set Options
#
# =============================================================================


def parse_command_line():
    """
    Parse the command line, return options and check for consistency among the
    options.
    """
    parser = OptionParser(
        version = git_version.verbose_msg,
        usage   = "%prog [options]",
        description = description
        )

    parser.add_option( "-d", "--database", action = "store", type = "string", default = None,
        help = 
            "Input database to read. Can only input one at a time."
        )
    parser.add_option( "", "--livetime-program", action = "store", type = "string", default = None,
        help = 
            "Required; Set the name of the program whose entries in the " +
            "search summary table will set the live time. Ex. inspiral, ringdown"
        )
    parser.add_option( "-t", "--tmp-space", action = "store", type = "string", default = None,
        metavar = "PATH",
        help = 
            "Location of local disk on which to do work. This is optional; " +
            "it is only used to enhance performance in a networked " +
            "environment. "
        )
    parser.add_option( "-v", "--verbose", action = "store_true", default = False,
        help =
            "Be verbose."
        )

    (options, args) = parser.parse_args()

    # check for required options and for self-consistency
    if not options.database:
        raise ValueError, "No database specified."

    if not options.livetime_program:
        raise ValueError, "No livetime program specified."

    return options, sys.argv[1:]

# =============================================================================
#
#                              Function Definitions
#
# =============================================================================


def get_playground_sets_from_fulldata(full_data_dict):
    """
    Calculates playground segment sets using full_data single-ifo segments.
    Returns a dictionary of segments (which is of form {instrument:segmentlist})
    with just playground segments in them.
    
    @param full_data_dict: the full_data single-ifo analyzed segments dictionary
    """
    playground_listdict = segments.segmentlistdict()
    playground_seglist = segmentsUtils.S2playground(full_data_dict.extent_all())
    for instruments, seglist in full_data_dict.items():
        playground_listdict[instruments] = seglist.coalesce() & playground_seglist 

    return playground_listdict


class Durations:
    """
    Class to store and retrieve durations.
    self.durations has the structure:
    self.durations[(on_instruments, veto_def_name, datatype, time_slide_id)]
    """
    def __init__( self ):
        self.durations = {}

    def map_slides_to_durations( self, veto_def_name, datatype, livetime_dict ):
        for key, duration in livetime_dict.items():
            # the dictionary key is of the form (time_slide_id, on_instruments)
            self.durations[( key[1], veto_def_name, datatype, str(key[0]) )] = duration

    def retrieve_duration( self, on_instruments, veto_def_name, datatype, time_slide_id ):
        # when both simulations & full_data are in the same database, use full_data livetime
        if datatype == "simulation":
            datatype = "all_data"
        # when tuple from the experiment tables is a key of the durations dict, return duration
        if (on_instruments, veto_def_name, datatype, str(time_slide_id)) in self.durations.keys():
            return self.durations[( on_instruments, veto_def_name, datatype, str(time_slide_id) )]


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

#
#       Generic Initilization
#

options, args = parse_command_line()

# get input database filename
filename = options.database
if not os.path.isfile( filename ):
    raise ValueError, "The input database, %s, cannot be found." % filename

# Setup working databases and connections
if options.verbose: 
    print >> sys.stderr, "Opening database..."

working_filename = dbtables.get_connection_filename( 
    filename, tmp_path = options.tmp_space, verbose = options.verbose )
connection = sqlite3.connect( working_filename )
if options.tmp_space:
    dbtables.set_temp_store_directory(connection, options.tmp_space, verbose = options.verbose)
xmldoc = dbtables.get_xml(connection)

livetime_program = sqlutils.validate_option( options.livetime_program )

# Add program to process and process params table

# FIXME: remove the following two lines once boolean type
# has been properly handled
from glue.ligolw import types as ligolwtypes
ligolwtypes.FromPyType[type(True)] = ligolwtypes.FromPyType[type(8)]

proc_id = process.register_to_xmldoc(xmldoc, __prog__, options.__dict__, version = git_version.id)

#
#       Compute Durations
#

# turn the time slide table into a dictionary
time_slide_dict = table.get_table(xmldoc, lsctables.TimeSlideTable.tableName).as_dict()
zero_lag_dict = dict([dict_entry for dict_entry in time_slide_dict.items() if not any( dict_entry[1].values() )])
del time_slide_dict[ zero_lag_dict.keys()[0] ]

# initialize duration bank
dur_bank = Durations()

sqlquery = """
    SELECT value
    FROM process_params
    WHERE param == "-userTag"
    GROUP BY value
"""
usertags = set(usertag[0] for usertag in connection.cursor().execute(sqlquery) )

# determine single-ifo segments for each usertag
ifo_segments = segments.segmentlistdict()
datatypes = {}

# FIXME: Not really sure why we're going round the houses to get FULL_DATA and
# FIXME: PLAYGROUND. I guess the point is you want to see if a separate PLAYGROUND
# FIXME: run was done and use it if so. Given that a -userTag command is pretty
# FIXME: explicitly hardcoding lalapps_inspiral, I will set this to FULL_DATA if
# FIXME: usertags returns garbage

if not (set(usertags) & set(["FULL_DATA","PLAYGROUND"])):
    usertags = ["FULL_DATA"]

for tag in (set(usertags) & set(["FULL_DATA","PLAYGROUND"])):
    ifo_segments[tag] = compute_dur.get_single_ifo_segments(connection, program_name = livetime_program, usertag = tag)
    if tag == "FULL_DATA":
        datatypes[tag] = ["all_data","slide"]
    elif tag == "PLAYGROUND":
        datatypes[tag] = ["playground","slide"]

if "FULL_DATA" in usertags:
    # find playground segments despite a playground analysis not being done
    if "PLAYGROUND" not in usertags:
        tag = unicode("PLAYGROUND")
        ifo_segments[tag] = get_playground_sets_from_fulldata(ifo_segments["FULL_DATA"])
        datatypes[tag] = ["playground"]

    tag = unicode("EXCLUDE_PLAY")
    ifo_segments[tag] = ifo_segments["FULL_DATA"] - ifo_segments["PLAYGROUND"]
    datatypes[tag] = ["exclude_play"]

if options.verbose:
    print >> sys.stderr, "Getting all veto categories in the experiment_summary table..."

# get veto_segments
veto_segments = compute_dur.get_veto_segments(xmldoc, options.verbose)

for veto_def_name, veto_seg_dict in veto_segments.items():
    if options.verbose:
        print >> sys.stderr, "\n\tThe DQ vetoes applied to the single-ifo segments are %s" % veto_def_name
    for usertag, ifo_seg_dict in ifo_segments.items():
        # compute the durations (or livetimes) for every possible instrument combo for every
        # slide in the time-slide table; the resulting durations dictionary has the following form:
        # durations[(on_instruments, veto_def_name, datatype, time_slide_id)] = livetime (in seconds)

        if options.verbose:
            print >> sys.stderr, "\n\t\tCalculating live-time for %s zerolag time" % usertag
        # determine the durations for the zerolag time
        livetime_dict = compute_dur.get_livetimes(
            ifo_seg_dict - veto_seg_dict,
            zero_lag_dict,
            verbose = options.verbose)
        dur_bank.map_slides_to_durations(
            veto_def_name,
            datatypes[usertag][0],
            livetime_dict)

        # determine the durations for each separate time-slide
        if len(datatypes[usertag]) > 1:
            if options.verbose:
                print >> sys.stderr, "\n\t\tCalculating live-time for each %s time slide" % usertag
            livetime_dict = compute_dur.get_livetimes(
                ifo_seg_dict - veto_seg_dict,
                time_slide_dict,
                verbose = options.verbose)
            dur_bank.map_slides_to_durations(
                veto_def_name,
                datatypes[usertag][1],
                livetime_dict)

#
# finished getting all durations, now populate the experiment_summary table
# with them
#
if options.verbose:
    print >> sys.stderr, "\nPopulating the experiment_summary table with results..."

connection.create_function("retrieve_duration", 4, dur_bank.retrieve_duration)

# populate the experiment_summary table with the appropiate duration
sqlquery = """
    UPDATE experiment_summary
    SET duration = (
        SELECT retrieve_duration(
            experiment.instruments,
            experiment_summary.veto_def_name,
            experiment_summary.datatype,
            experiment_summary.time_slide_id
            )
        FROM
            experiment
        WHERE
            experiment.experiment_id == experiment_summary.experiment_id
        )"""
connection.cursor().execute( sqlquery )

#
#       Close database and exit
#

connection.commit()
connection.cursor().close()
dbtables.put_connection_filename(filename, working_filename, verbose = options.verbose)

if options.verbose:
    print >> sys.stderr, "Finished!"

# set process end time
process.set_process_end_time(proc_id)
sys.exit(0)


