#!/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-file1", metavar = "file", help = "Set the first segment file")
	parser.add_option("--segment-file2", metavar = "file", help = "Set the first segment file")
	parser.add_option("--segment-name1", metavar = "name", default = "datasegments", help = "Set the name of segments default datasegments")
	parser.add_option("--segment-name2", metavar = "name", default = "datasegments", help = "Set the name of segments default datasegments")
	parser.add_option("--intersection", action = "store_true", help = "intersection of --segment-name1 in segment-file1 and --segment-name2 in segment-file2")
	parser.add_option("--union", action = "store_true", help = "union of --segment-name1 in segment-file1 and --segment-name2 in segment-file2")
	parser.add_option("--diff", action = "store_true", help = "diff of --segment-name1 in segment-file1 and --segment-name2 in segment-file2")
	parser.add_option("--output-file", metavar = "file", default = ".", help = "Set the name of the output file")
	parser.add_option("--output-segment-name", metavar = "name", default = "datasegments", help = "Set the name of output segments default datasegments")
	parser.add_option("-v", "--verbose", action = "store_true", help = "Be verbose.")
	options, filenames = parser.parse_args()
	
	if len([op for op in (options.intersection, options.union, options.diff) if op is not None]) != 1:
		raise ValueError("must specify exactly one of --intersection, --union or --diff")

	return options, filenames

options, filenames = parse_command_line()

def extract_segs(fname, segname):
	xmldoc = utils.load_filename(fname, verbose = True)
	return ligolw_segments.segmenttable_get_by_name(xmldoc, segname).coalesce()

segs1 = extract_segs(options.segment_file1, options.segment_name1)
segs2 = extract_segs(options.segment_file2, options.segment_name2)

if options.diff:
	segs = segs1 - segs2
if options.intersection:
	segs = segs1 & segs2
if options.union:
	segs = segs1 | segs2

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.output_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)
		segtab.append(row)

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

