#!/usr/bin/env python
#
# Copyright (C) 2011--2013  Kipp Cannon, Chad Hanna, Drew Keppel
#
# 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.


from optparse import OptionParser
import sys

import pygtk
pygtk.require("2.0")
import gobject
gobject.threads_init()
import pygst
pygst.require('0.10')
import gst


from pylal import series as lalseries
from glue.ligolw import ligolw
from glue.ligolw import array
from glue.ligolw import param
from glue.ligolw import lsctables
array.use_in(ligolw.LIGOLWContentHandler)
param.use_in(ligolw.LIGOLWContentHandler)
lsctables.use_in(ligolw.LIGOLWContentHandler)
from glue.ligolw import utils
from gstlal import datasource
from gstlal import pipeparts
from gstlal import simplehandler
from glue.ligolw.utils import segments as ligolw_segments
from pylal import series as lalseries


def parse_command_line():
	parser = OptionParser(description = __doc__)
	parser.add_option("--output", metavar = "filename", help = "Set the filename in which to save the output.  If not given, output is sent to the default audio device.  The filename's extension determines the format, the following are recognized:  .wav, .flac, .ogg, .txt.")
	parser.add_option("--rate", metavar = "Hz", type = "int", help = "Downsample input to this sample rate (default = no resampling).")
	parser.add_option("--whiten", action = "store_true", help = "Whiten the time series (default = do not whiten).")
	parser.add_option("--reference-psd", metavar = "filename", help = "When whitening, normalize the time series to the spectrum described in this XML file.  If this option is not given, the spectrum is measured from the data.") 
	parser.add_option("--low-pass-filter", metavar = "Hz", type = "float", help = "Low pass filter frequency (default = no low-pass filter).  Low-pass filter is applied after whitening.")
	parser.add_option("--high-pass-filter", metavar = "Hz", type = "float", help = "High pass filter frequency (default = no high-pass filter).  High-pass filter is applied after whitening.")
	parser.add_option("--amplification", metavar = "factor", type = "float", default = 1.0, help = "Amplify the timeseries this much (default = no amplification).  Amplification is applied after low- and high-pass filtering.")
	parser.add_option("-v", "--verbose", action = "store_true", help = "Be verbose.")
	datasource.append_options(parser)

	options, filenames = parser.parse_args()

	if options.low_pass_filter is not None and options.high_pass_filter is not None:
		if options.low_pass_filter <= options.high_pass_filter:
			raise ValueError("--low-pass-filter must be > --high-pass-filter")
	if options.reference_psd and not options.whiten:
		raise ValueError("--reference-psd requires --whiten")
	if len(options.channel_name) > 1:
		raise ValueError("only one --channel-name allowed")

	return options, filenames


# parsing and setting up some core structures
options, filenames = parse_command_line()

gw_data_source_info = datasource.GWDataSourceInfo(options)
instrument, = gw_data_source_info.channel_dict

if options.reference_psd is not None:
	psd = lalseries.read_psd_xmldoc(utils.load_filename(options.reference_psd, verbose = options.verbose, contenthandler = ligolw.LIGOLWContentHandler))[instrument]
else:
	psd = None


# building the event loop and pipeline
gobject.threads_init()
mainloop = gobject.MainLoop()
pipeline = gst.Pipeline("gstlal_play_frames")
handler = simplehandler.Handler(mainloop, pipeline)

# the pipeline

head = datasource.mkbasicsrc(pipeline, gw_data_source_info, instrument, verbose = options.verbose)

if options.rate is not None:
	if options.whiten:
		# disallow upsampling.  upsampling screws with the whitener
		# because then there are frequency bins that have no
		# spectral content
		head = pipeparts.mkcapsfilter(pipeline, head, "audio/x-raw-float, rate=[%d,MAX]" % options.rate)
	head = pipeparts.mkresample(pipeline, head, quality = 9)
	head = pipeparts.mkcapsfilter(pipeline, head, "audio/x-raw-float, rate=%d" % options.rate)

if options.whiten:
	head = pipeparts.mkwhiten(pipeline, head)

	if psd[instrument] is None:
		# use running average PSD
		head.set_property("psd-mode", 0)
	else:
		head.set_property("psd-mode", 1)

		#
		# install signal handler to retrieve \Delta f and
		# f_{Nyquist} whenever they are known and/or change,
		# resample the user-supplied PSD, and install it into the
		# whitener.
		#

		def psd_resolution_changed(elem, pspec, psd):
			# get frequency resolution and number of bins
			delta_f = elem.get_property("delta-f")
			n = int(round(elem.get_property("f-nyquist") / delta_f) + 1)
			# interpolate and install PSD
			psd = reference_psd.interpolate_psd(psd, delta_f)
			elem.set_property("mean-psd", psd.data[:n])

		head.connect_after("notify::f-nyquist", psd_resolution_changed, psd)
		head.connect_after("notify::delta-f", psd_resolution_changed, psd)

if options.high_pass_filter is not None and options.low_pass_filter is not None:
	head = pipeparts.mkaudiochebband(pipeline, head, lower_frequency = options.high_pass_filter, upper_frequency = options.low_pass_filter)
else:
	if options.high_pass_filter is not None:
		head = pipeparts.mkaudiocheblimit(pipeline, head, cutoff = options.high_pass_filter, mode = "high-pass")
	elif options.low_pass_filter is not None:
		head = pipeparts.mkaudiocheblimit(pipeline, head, cutoff = options.low_pass_filter, mode = "low-pass")

head = pipeparts.mkaudioconvert(pipeline, pipeparts.mkaudioamplify(pipeline, head, options.amplification))

if options.output is None:
	pipeparts.mkgeneric(pipeline, pipeparts.mkqueue(pipeline, head), "autoaudiosink")
elif options.output.endswith(".wav"):
	pipeparts.mkfilesink(pipeline, pipeparts.mkwavenc(pipeline, head), options.output)
#elif options.output.endswith(".mp3"):
#	head = pipeparts.mkgeneric(pipeline, head, "lame")
#	pipeparts.mkfilesink(pipeline, head, options.output)
elif options.output.endswith(".flac"):
	pipeparts.mkfilesink(pipeline, pipeparts.mkflacenc(pipeline, head), options.output)
elif options.output.endswith(".ogg"):
	head = pipeparts.mkoggmux(pipeline, pipeparts.mkgeneric(pipeline, head, "vorbisenc"))
	pipeparts.mkfilesink(pipeline, head, options.output)
elif options.output.endswith(".txt"):
	pipeparts.mknxydumpsink(pipeline, head, options.output)
else:
	raise ValueError("unrecognized format for --output")


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