#!/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 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("--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()
# make it safe to contract by elimating segments that are smaller than 2 * trim value
for k in segs:
	segs[k] = segments.segmentlist([seg for seg in segs[k] if abs(seg) > 2 * trim])
segs.contract(trim)

xmldoc = ligolw.Document()
xmldoc.appendChild(ligolw.LIGO_LW())
seg_def = lsctables.New(lsctables.SegmentDefTable, ligolw_segments.LigolwSegmentList.segment_def_columns)
xmldoc.childNodes[0].appendChild(seg_def)
segtab = lsctables.New(lsctables.SegmentTable, ligolw_segments.LigolwSegmentList.segment_columns)
xmldoc.childNodes[0].appendChild(segtab)

# populate the segment definer table
ifo_ids = {}
for ifo in segs.keys():
	row = seg_def.RowType()
	id = seg_def.get_next_id()
	row.set_ifos([ifo])
	row.segment_def_id = id
	row.process_id = None
	row.name = options.segment_name
	row.version = None
	row.comment = None
	row.insertion_time = None
	ifo_ids[ifo] = id
	seg_def.append(row)

for ifo, segl in segs.items():
	for seg in segl:
		row = segtab.RowType()
		row.segment_def_id = ifo_ids[ifo]
		row.segment_id = segtab.get_next_id()
		row.process_id = None
		row.set(seg)
		if abs(seg) > options.min_length:
			segtab.append(row)

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

