#!/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 gstlal import omega
from optparse import OptionParser, Option
from glue.ligolw import lsctables
from glue.ligolw import utils
from glue.ligolw import ligolw
from glue.ligolw.utils import process
from glue.ligolw.ilwd import ilwdchar
from glue import segments
import numpy
from pylal.xlal.datatypes.snglburst import from_buffer as sngl_bursts_from_buffer
from gstlal import pipeparts
from gstlal import lloidparts
from gstlal import simplehandler
from pylal.datatypes import LIGOTimeGPS
import pygtk
pygtk.require("2.0")
import gobject
gobject.threads_init()
import pygst
pygst.require("0.10")
import gst


#
# Utility functions
#

def row_init(tab, row):
	for k,v in tab.validcolumns.items():
		if "int" in v:
			setattr(row, k, 0)
		elif "real" in v:
			setattr(row, k, 0)
		elif "string" in v:
			setattr(row, k, "")
		elif "ilwd" in v:
			setattr(row, k, "")

def create_bank(sample_rate, ifo, output, fmin = 10, fmax = 2000, fnum = 200, qmin = 3, qmax = 11, qnum = 5):
	# setup template bank file
	xmldoc = ligolw.Document()
	lw = xmldoc.appendChild(ligolw.LIGO_LW())
	burst = lsctables.New(lsctables.SnglBurstTable)
	lw.appendChild(burst)
	# FIXME do process table stuff correctly
	procrow = process.register_to_xmldoc(xmldoc, "gstlal_omega", options.__dict__)

	# create template banks
	# FIXME dont hardcode, provide as command line options
	f_vec = numpy.linspace(fmin, fmax, fnum)
	q_vec = numpy.linspace(qmin, qmax, qnum)

	for f in f_vec:
		for q in q_vec:
			row = burst.RowType()
			row_init(burst, row)
			row.central_freq = f
			row.duration = omega.duration_from_f_and_q(f,q)
			row.ifo = ifo
			burst.append(row)
			row.process_id = procrow.process_id
			row.event_id = ilwdchar("sngl_burst:event_id:1")

	utils.write_filename(xmldoc, output)
	bank = omega.omega_bank(burst, sample_rate)
	return bank


def parse_options():
	parser = OptionParser(description = __doc__)
	parser.add_option("--gps-start-time", metavar = "seconds", type = "int", help = "Set the start time of the segment to analyze in GPS seconds (required).  Can be specified to nanosecond precision.")
	parser.add_option("--gps-end-time", metavar = "seconds", type = "int", help = "Set the end time of the segment to analyze in GPS seconds (required).  Can be specified to nanosecond precision.")
	parser.add_option("--instrument", metavar = "ifo", help = "Set the name of the instrument")
	parser.add_option("--bank-output", metavar = "filename", help = "Set the name of the LIGO light-weight XML file to output")
	parser.add_option("--sample-rate", type = "int", default = 2048, help = "sample rate")
	parser.add_option("--injections", metavar = "filename", help = "Set the name of the LIGO light-weight XML file from which to load injections (optional).")
	parser.add_option("--channel-name", metavar = "name", default = "LSC-STRAIN", help = "Set the name of the channel to process (optional).  The default is \"LSC-STRAIN\".")
	parser.add_option("-v", "--verbose", action = "store_true", help = "Be verbose (optional).")
	parser.add_option("--frame-cache", metavar = "name", help = "Set the name of the frame cache (required)")
	parser.add_option("--output", help = "output file name")


	return parser.parse_args()
#
# MAIN
#

# Setup
options, filenames = parse_options()
seg = segments.segment(LIGOTimeGPS(options.gps_start_time), LIGOTimeGPS(options.gps_end_time))
bank = create_bank(options.sample_rate, options.instrument, options.bank_output)
print bank.shape

# pipeline setup
pipeline = gst.Pipeline(sys.argv[0])
mainloop = gobject.MainLoop()
handler = simplehandler.Handler(mainloop, pipeline)
seekevent = gst.event_new_seek(1.0, gst.Format(gst.FORMAT_TIME), gst.SEEK_FLAG_FLUSH | gst.SEEK_FLAG_KEY_UNIT, gst.SEEK_TYPE_SET, seg[0].ns(), gst.SEEK_TYPE_SET, seg[1].ns())

# pipeline construction
head = lloidparts.mkLLOIDbasicsrc(pipeline, seekevent, options.instrument, lloidparts.DetectorData(options.frame_cache, options.channel_name), data_source = "frames", injection_filename = options.injections, frame_segments = None, verbose = options.verbose)
head = lloidparts.mkLLOIDsrc(pipeline, head, [options.sample_rate], options.instrument, seekevent = seekevent, track_psd = True)
head = pipeparts.mkprogressreport(pipeline, head[options.sample_rate], "src")
# FIXME don't hard code these values
head = pipeparts.mkfirbank(pipeline, head, latency = 0, fir_matrix = bank, block_stride = 4096)
head = pipeparts.mkbursttriggergen(pipeline, pipeparts.mktogglecomplex(pipeline, head), 4096, options.bank_output)

# plain txt file output stage
appsink = pipeparts.mkappsink(pipeline, pipeparts.mkqueue(pipeline, head))
outfile = open(options.output, "w")
def dump_triggers(elem, output = outfile):
	for row in sngl_bursts_from_buffer(elem.emit("pull-buffer")):
		# FIXME jamie make the columns you want
		print >>outfile, row.peak_time + row.peak_time_ns*1e-9, row.snr, row.central_freq, row.duration
appsink.connect_after("new-buffer", dump_triggers)

# away we go.
if pipeline.set_state(gst.STATE_PLAYING) == gst.STATE_CHANGE_FAILURE:
	raise RuntimeError("pipeline failed to enter PLAYING state")
mainloop.run() 

