#!/usr/bin/env python
#
# Copyright (C) 2011 Chad Hanna
#
# 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 2 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.

import sys
from glue import segments
from glue import segmentsUtils
from glue.ligolw import ligolw
from glue.ligolw import utils
from glue.ligolw import lsctables
from glue.ligolw.utils import segments as ligolw_segments
from glue.ligolw.utils import process as ligolw_process
from pylal.datatypes import LIGOTimeGPS
from optparse import OptionParser

def parse_command_line():
	parser = OptionParser()
	parser.add_option("--segment-name", metavar = "name", default = "datasegments", help = "Set the name of segments default datasegments")
	parser.add_option("--trim", metavar = "int", type = "int", default = 0, help = "Set the trim factor default 0")
	parser.add_option("--gps-start-time", metavar = "int", type = "int", default = None, help = "GPS start time to which segments are trimmed, default = None")
	parser.add_option("--gps-end-time", metavar = "int", type = "int", default = None, help = "GPS end time to which segments are trimmed, default = None")
	parser.add_option("--min-length", metavar = "int", type = "int", default = 0, help = "Set the min segment length, default 0")
	parser.add_option("--output", metavar = "file", default = ".", help = "Set the name of the output file")
	parser.add_option("-v", "--verbose", action = "store_true", help = "Be verbose.")
	options, filenames = parser.parse_args()

	return options, filenames

options, filenames = parse_command_line()

fname = filenames[0]
trim = options.trim

xmldoc = utils.load_filename(fname, verbose = True)

segs = ligolw_segments.segmenttable_get_by_name(xmldoc, options.segment_name).coalesce()
for k in segs:
	# Make it safe to contract by elimating segments that are smaller than
	# 2 * trim value.  Remove segments that don't satisfy the minimum length
	segs[k] = segments.segmentlist([seg for seg in segs[k] if (abs(seg) > 2 * trim) and (abs(seg) > options.min_length)])
	
segs.contract(trim)

# trim the segment list to certain GPS times if start and end GPS times are given
if options.gps_start_time is not None or options.gps_end_time is not None:
	start = LIGOTimeGPS(options.gps_start_time) if options.gps_start_time is not None else segments.NegInfinity
	stop = LIGOTimeGPS(options.gps_end_time) if options.gps_end_time is not None else segments.PosInfinity
	boundaries = segments.segmentlist([segments.segment(start, stop)])
	for seglist in segs.values():
		seglist &= boundaries

xmldoc = ligolw.Document()
xmldoc.appendChild(ligolw.LIGO_LW())

process = ligolw_process.register_to_xmldoc(xmldoc, sys.argv[0], options.__dict__)
ligolw_process.set_process_end_time(process)

lwseglists = ligolw_segments.LigolwSegments(xmldoc)
lwseglists.insert_from_segmentlistdict(segs, options.segment_name)
lwseglists.optimize()
lwseglists.finalize(process)

utils.write_filename(xmldoc, options.output, gz = options.output.endswith('.gz'), verbose = True)

