#!/usr/bin/env python
#
# Copyright (C) 2010  Peter Couvares  <pfcouvar@syr.edu>
#
# 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.

"""
This intended as a drop-in replacement for ligolw_segments_from_cats
which splits the specified query into multiple shorter queries (by GPS
time) and recombines the results.  It supports a single additional
optional argument, --max_gps_duration, to specify the maximum gps time
of each constituent query (default is one week).
"""

from __future__ import print_function
import sys
import os
import pwd
import tempfile
import re
from optparse import OptionParser

import ligo.segments
from ligo.segments.utils import segmentlist_range
from glue import git_version

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


__author__ = "Peter Couvares <pfcouvar@syr.edu>"
__version__ = "git id %s" % git_version.id
__date__ = git_version.date

# lifted from ligolw_segments_from_cats, with --max-gps-duration added
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("-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("-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")

    parser.add_option("--max-gps-duration",   metavar = "max_gps_duration", help = "Maximum duration of an individual query; larger durations will be split into multiple queries and recombined.")

    options, others = parser.parse_args()

    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


# generates a modified version of the provided argv array, minus a
# given argument (and optionally its value)
def strip_arg(argv, long_name, short_name=None, remove_value=True):
    argn = 0
    while argn < len(argv):
        arg = argv[argn]

	# if the argument includes a value delimited by '=', separate it
        argsplit = arg.split('=')
        arg = argsplit[0]
	
        if arg == long_name or arg == short_name:
            # if we are to remove the associated value, and that value
            # is not already inside the arg (via '='), then skip the
            # subsequent argument
            if remove_value and len(argsplit)==1:
                argn += 1
        else:
            yield argv[argn]
        argn += 1


if __name__ == '__main__':
    options = parse_command_line()

    if options.max_gps_duration is None:
        max_gps_duration = 60*60*24*7
    else:
        max_gps_duration = int(options.max_gps_duration)

    gps_start_time = int(options.gps_start_time)
    gps_end_time = int(options.gps_end_time)
    gps_duration = gps_end_time - gps_start_time

    if max_gps_duration > gps_duration:
        max_gps_duration = gps_duration

#    print gps_start_time
#    print gps_end_time

    gps_segmentlist = list(segmentlist_range(gps_start_time, gps_end_time, max_gps_duration))
    if gps_segmentlist[-1][1] != int(gps_end_time):
        gps_segmentlist.append(ligo.segments.segment(gps_segmentlist[-1][1],gps_end_time))

    print("segmentlist =", gps_segmentlist)

    argv = sys.argv[1:]
#    print "original argv =", argv
    argv = list(strip_arg(argv, "-s", "--gps-start-time"))
    argv = list(strip_arg(argv, "-e", "--gps-end-time"))
    argv = list(strip_arg(argv, "-o", "--output-dir"))
    argv = list(strip_arg(argv, "--max-gps-duration"))

#    print "stripped argv =", argv

    temp_dirs = []

    for segment in gps_segmentlist:
        temp_dir = tempfile.mkdtemp()
        temp_dirs.append(temp_dir)
        temp_argv = list(argv)
        temp_argv.append("-s %s" % segment[0])
        temp_argv.append("-e %s" % segment[1])
        temp_argv.append("-o %s" % temp_dir)
#        print "temp_argv =", temp_argv
        os.system("ligolw_segments_from_cats %s" % ' '.join(temp_argv))

    # in theory, the same output files should be present in each output
    # dir, so we just need to look in one to figure out the names
    # e.g., "L1-VETOTIME_CAT2-970271943-259200.xml"

    prefix = re.compile(r"^[A-Z][0-9]-VETOTIME_CAT[0-9]")
    temp_files = [prefix.match(s).group() for s in os.listdir(temp_dirs[0])]

    for file in temp_files:
        file_instances = ["%s/%s-*-*.xml" % (dir, file) for dir in temp_dirs]
        os.chdir(options.output_dir)
#        os.system("pwd")
        result_file = "%s-%d-%d.xml" % (file, gps_start_time, gps_duration)
        os.system("ligolw_add --output %s %s" % (result_file, ' '.join(file_instances)))
#        os.system("rm -v %s" % ' '.join(file_instances))

#    for dir in temp_dirs:
#        os.rmdir(dir)
