#!/usr/bin/env python
#
# Copyright (C) 2011  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 os
import numpy
import sys

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

from gstlal import pipeparts
from gstlal import reference_psd
from gstlal import simplehandler
from gstlal import datasource
from gstlal import multirate_datasource
from glue.ligolw import ligolw
from glue.ligolw import array as ligolw_array
from glue.ligolw import param as ligolw_param
from glue.ligolw import utils as ligolw_utils
from pylal.datatypes import LIGOTimeGPS
from pylal import series as lalseries

class LIGOLWContentHandler(ligolw.LIGOLWContentHandler):
	pass
ligolw_array.use_in(LIGOLWContentHandler)
ligolw_param.use_in(LIGOLWContentHandler)

## @file gstlal_fake_frames
# This program will make fake data in a variety of ways; see gstlal_fake_frames for help and usage
#
# This program will make fake data in a variety of ways.  Its input is anything
# supported by datasource.py.  You can additionally whiten the data or apply a
# frequency domain coloring filter.  It writes its output to frame files.
#
# #### Overview graph of the pipeline ####
#
# - gray boxes are optional and depend on the command line given
#
# @dot
# digraph G {
#	// graph properties
#
#	rankdir=LR;
#	compound=true;
#	node [shape=record fontsize=10 fontname="Verdana"];	
#	edge [fontsize=8 fontname="Verdana"];
#
#	// nodes
#
#	"mkbasicsrc()" [URL="\ref datasource.mkbasicsrc()"];
#	"whitened_multirate_src()" [label="whitened_multirate_src()", URL="\ref multirate_datasource.mkwhitened_multirate_src()", style=filled, color=lightgrey];
#	capsfilter [style=filled, color=lightgrey, URL="\ref pipeparts.mkcapsfilter()"];
#	resample [style=filled, color=lightgrey, URL="\ref pipeparts.mkresample()"];
#	audioamplify [style=filled, color=lightgrey, URL="\ref pipeparts.mkaudioamplify()", label="audioamplify \n if --data-source=white \n in order to have unit variance \n after resampling"];
#	lal_shift[style=filled, color=lightgrey, URL="\ref pipeparts.mkshift()", label="lal_shift \n [iff --shift provided]"];
#	progressreport1 [URL="\ref pipeparts.mkprogressreport()", label="progressreport \n [iff --verbose provided]"];
#	progressreport2 [style=filled, color=lightgrey, URL="\ref pipeparts.mkprogressreport()", label="progressreport \n [iff --verbose provided]"];
#	lal_firbank [style=filled, color=lightgrey, URL="\ref pipeparts.mkfirbank()", label="lal_firbank \n [iff --color-psd provided]"];
#	taginject [URL="\ref pipeparts.mktaginject()"];
#	lal_simulation [style=filled, color=lightgrey, URL="\ref pipeparts.mkinjections()", label="lal_simulation \n [iff --injections provided]"];
#	framecppchannelmux [URL="\ref pipeparts.mkframecppchannelmux()"];
#	framecppfilesink [URL="\ref pipeparts.mkframecppfilesink()"];
#
#	// connect it up
#
#	"mkbasicsrc()" -> progressreport1-> lal_shift  -> progressreport2;
#	progressreport2 -> "whitened_multirate_src()" [label="if --whiten given"];
#	"whitened_multirate_src()" -> lal_firbank;
#	progressreport2 -> resample [label="if --whiten not given"];
#	resample -> capsfilter;
#	capsfilter -> audioamplify
#	audioamplify -> lal_firbank;
#	lal_firbank -> taginject -> lal_simulation -> framecppchannelmux -> framecppfilesink;
# }
# @enddot
#
#
# ### Usage cases ###
#
#
# -# Making initial LIGO colored noise, Advanced LIGO noise, etc for different
# instruments.   See datasource.append_options() for other choices
#
#		$ gstlal_fake_frames --data-source=LIGO --channel-name=H1=FAKE-STRAIN --frame-type=H1_FAKE --gps-start-time=900000000 --gps-end-time=900005000 --output-path=testframes  --verbose
#
#		$ gstlal_fake_frames --data-source=AdvLIGO --channel-name=L1=FAKE-STRAIN --frame-type=L1_FAKE --gps-start-time=900000000 --gps-end-time=900005000 --output-path=testframes  --verbose
#
#
# -# Custom colored noise.
#
#	Obviously it is not practical to code up every possible noise curve to
#	simulate as a custom data source.  However, it is possible to provide your
#	own custom noise curve as an ASCII file with frequency in one column and
#	strain/Hz in the second. e.g., early advanced ligo noise curve found <a
#	href=http://www.lsc-group.phys.uwm.edu/cgit/gstlal/plain/gstlal/share/v1_early_asd.txt>here</a>.
#	You will need to convert ASD text files to PSD xml files using
#	gstlal_psd_xml_from_asd_txt first.
#
#		$ gstlal_fake_frames --data-source=white --color-psd v1psd.xml.gz --channel-name=V1=FAKE-STRAIN --frame-type=V1_FAKE --gps-start-time=900000000 --gps-end-time=900005000 --output-path=testframes  --verbose
#
#
# -# Recoloring existing frame data
#
#	This procedure is very similar to the above except that instead of
#	using white noise to drive the fake frame generation, we start with
#	real frame data and whiten it.  Recoloring can be thought of as first
#	whitening data and then applying a coloring filter.  You must first
#	make a reference PSD of the data with e.g. gstlal_reference_psd.  You
#	will need to make a frame cache of the frames to recolor.
#
#		gstlal_fake_frames --whiten-reference-psd reference_psd.xml.gz --color-psd recolor_psd.xml.gz --data-source frames --output-path /path/to/output --output-channel-name h_16384Hz --gps-start-time 966384031 --frame-type T1300121_V1_EARLY_RECOLORED_V2 --gps-end-time 966389031 --frame-duration 16 --frames-per-file 256 --frame-cache frame.cache --channel-name=V1=h_16384Hz --verbose
#
#
# -# Creating injections into silence (assumes you have an injection xml file from e.g. lalapps_inspinj)
#
#		gstlal_fake_frames --data-source silence --output-path /path/to/output --gps-start-time 966384031 --frame-type V1_INJECTIONS --gps-end-time 966389031 --frame-duration 16 --frames-per-file 256 --verbose --channel-name=V1=INJECTIONS --injections injections.xml
#
#
# -# Other things are certainly possible, please add some!
#
# ### Command line options ###
#
# See datasource.append_options() for common command line options shared among different programs
#
#	+ `--shift` [int] (ns): Number of nanoseconds, \f$\tau\f$ to delay (negative) or advance (positive) the input time series \f$x\f$ relative to the output time series \f$y\f$. \f$ y(t) = x(t+\tau)\f$
#	+ `--sample-rate` [int] (Hz): Sample rate at which to generate the data, should be less than or equal to the sample rate of the measured psds provided.  Default = 16384 Hz, max 16384 Hz.
# 	+ `--whiten-reference-psd` [file name]: Set the name of psd xml file to whiten the data with.
#	+ `--whiten-track-psd` []: Calculate PSD from input data and track with time.
#	+ `--color-psd` [file name]: Set the name of psd xml file to color the data with
#	+ `--output-path` [file path]: Path to output frame files. Default = "."
#	+ `--output-channel-name` [string]: The name of the channel in the output frames. The default is the same as the channel name.
#	+ `--frame-type` [string]: Frame type, required.
#	+ `--frame-duration` [int] (s): Set the duration of the output frames.  The duration of the frame file will be multiplied by --frames-per-file.  Default is 16s.
#	+ `--frames-per-file` [int]: Set the number of frames per file.  Default is 256.
#	+ `--verbose` []: Be verbose.
#
# ### Debug
#
# It is possible to check the pipeline graph by interupting the running process
# with ctrl+c if you have the GST_DEBUG_DUMP_DOT_DIR evironment set.  A dot
# graph will be written to gstlal_fake_frames.  Here is an example:
#
#		$ GST_DEBUG_DUMP_DOT_DIR="." gstlal_fake_frames --data-source=LIGO --channel-name=H1=FAKE-STRAIN --frame-type=H1_FAKE --gps-start-time=900000000 --gps-end-time=900005000 --output-path=testframes  --verbose
#
# You should see:
#
#		Wrote pipeline to ./gstlal_fake_frames_PLAYING.dot
#
# After Ctrl+c you should see:
# 
#		^C*** SIG 2 attempting graceful shutdown (this might take several minutes) ... ***
#		Wrote pipeline to ./gstlal_fake_frames.dot
#
# You can then turn the pipeline graph into an image with dot, e.g.,
#
#		dot -Tpng gstlal_fake_frames_PLAYING.dot > test.png
#
# ### Review
# 
# 

##
# Return the number of digits in a number
#
def digits(num):
	return int(numpy.ceil(numpy.log10(float(num))))


##
# Extract and validate the command line options
#
def parse_command_line():
	parser = OptionParser(description = __doc__)

	#
	# Append data source options
	#

	datasource.append_options(parser)

	#
	# Append program specific options
	#

	parser.add_option("--shift", metavar = "ns", help = "Number of nanoseconds to delay (negative) or advance (positive) the time stream", type = "int")
	parser.add_option("--sample-rate", metavar = "Hz", default = 16384, type = "int", help = "Sample rate at which to generate the data, should be less than or equal to the sample rate of the measured psds provided, default = 16384 Hz, max 16384 Hz")
	parser.add_option("--whiten-reference-psd", metavar = "name", help = "Set the name of psd xml file to whiten the data with")
	parser.add_option("--whiten-track-psd", action = "store_true", help = "Calculate PSD from input data and track with time.")
	parser.add_option("--color-psd", metavar = "name", help = "Set the name of psd xml file to color the data with")
	parser.add_option("--output-path", metavar = "name", default = ".", help = "Path to output frame files (default = \".\").")
	parser.add_option("--output-channel-name", metavar = "name", help = "The name of the channel in the output frames. The default is the same as the channel name")
	parser.add_option("--frame-type", metavar = "name", help = "Frame type, required")
	parser.add_option("--frame-duration", metavar = "s", default = 16, type = "int", help = "Set the duration of the output frames.  The duration of the frame file will be multiplied by --frames-per-file.  Default: 16s")
	parser.add_option("--frames-per-file", metavar = "n", default = 256, type = "int", help = "Set the number of frames per file.  Default: 256")
	parser.add_option("-v", "--verbose", action = "store_true", help = "Be verbose (optional).")

	#
	# Parse options
	#

	options, filenames = parser.parse_args()

	if options.sample_rate > 16384:
		raise ValueError("--sample-rate must be <= 16384")

	if options.frame_type is None:
		raise ValueError("--frame-type is required")

	return options, filenames


#
# Main
#

options, filenames = parse_command_line()

## Parse the command line options into a python.datasource.GWDataSourceInfo class instance
gw_data_source = datasource.GWDataSourceInfo(options)

## Assume instrument is the first and only key of the python.datasource.GWDataSourceInfo.channel_dict
instrument = gw_data_source.channel_dict.keys()[0]

# disable progress reports if not verbose
if not options.verbose:
	pipeparts.mkprogressreport = lambda pipeline, src, *args: src

# set default output channel if not set by user
if options.output_channel_name is None:
	options.output_channel_name = gw_data_source.channel_dict[instrument]

# do not do injections in datasource.mkbasicsrc(), we will do them later
injections, gw_data_source.injection_filename = options.injections, None

## Setup the pipeline
pipeline = gst.Pipeline(os.path.split(sys.argv[0])[1])

## Main loop 
mainloop = gobject.MainLoop()

## An instance of the python.simplehandler.Handler class
handler = simplehandler.Handler(mainloop, pipeline)

## 1) Set the pipeline head to basic input from datasource.mkbasicsrc()
head = datasource.mkbasicsrc(pipeline, gw_data_source, instrument, verbose = options.verbose)

## 2) Set the pipeline head to be verbose with pipeparts.mkprogressreport()
head = pipeparts.mkprogressreport(pipeline, head, "frames")

if options.shift is not None:
	## 3) Set the pipeline head to apply a time shift if requested with a pipeparts.mkshift() element
	head = pipeparts.mkshift(pipeline, head, shift = options.shift)

	## 4) Set the pipeline head to be verbose with a pipeparts.mkprogressreport() element
	head = pipeparts.mkprogressreport(pipeline, head, "frames_shifted")

if options.whiten_reference_psd:
	## If whitening read the PSD
	wpsd = lalseries.read_psd_xmldoc(ligolw_utils.load_filename(options.whiten_reference_psd, verbose = options.verbose, contenthandler = LIGOLWContentHandler))[instrument]
else:
	## else set wpsd to None
	wpsd = None

if options.whiten_reference_psd or options.whiten_track_psd:
	## 5) Set the pipeline head to a whitened data stream if requested using a multirate_datasource.mkwhitened_multirate_src()
	head = multirate_datasource.mkwhitened_multirate_src(pipeline, head, [options.sample_rate], instrument, psd = wpsd, seekevent = gw_data_source.seekevent, track_psd = options.whiten_track_psd)[options.sample_rate]
else:
	## 6) Otherwise simply add a pipeparts.mkcapsfilter() and pipeparts.mkresample()
	head = pipeparts.mkcapsfilter(pipeline, pipeparts.mkresample(pipeline, head, quality = 9), "audio/x-raw-float, rate=%d" % options.sample_rate)
	# FIXME this is a bit hacky, should datasource.mkbasicsrc be patched to change the sample rate?
	# FIXME don't hardcode sample rate (though this is what datasource generates for all fake data, period 
	# Renormalize if datasource is "white"
	if options.data_source == "white":
		head = pipeparts.mkaudioamplify(pipeline, head, 1.0 / (pipeparts.audioresample_variance_gain(9, 16384, options.sample_rate))**.5)

# Apply a coloring filter
if options.color_psd:

	## read coloring psd file and convert to an FIR filter
	rpsd = lalseries.read_psd_xmldoc(ligolw_utils.load_filename(options.color_psd, verbose = options.verbose, contenthandler = LIGOLWContentHandler))[instrument]
	
	## Calculate the maximum sample rate
	max_sample = int(round(1.0 / rpsd.deltaF * options.sample_rate / 2.0)) + 1

	# Truncate to requested output sample rate, if it is higher than the psd provides an assert will fail later
	rpsd.data = rpsd.data[:max_sample]
	
	# create the coloring FIR kernel from reference_psd.psd_to_fir_kernel()
	fir_matrix, latency, measured_sample_rate = reference_psd.psd_to_fir_kernel(rpsd)

	## 7) Set the pipeline head to a pipeparts.mkfirbank() recoloring filter
	head = pipeparts.mkfirbank(pipeline, head, latency = latency, fir_matrix = [fir_matrix], block_stride = 32 * options.sample_rate)


## Set the tags for the output data
tagstr = "units=strain,channel-name=%s,instrument=%s" % (options.output_channel_name, instrument)

## 8) Put the units back to strain before writing to frames.  Additionally, override the output channel name if provided from the command line
head = pipeparts.mktaginject(pipeline, head, tagstr)


if injections is not None:
	## 9) Do injections into the final fake data
	head = pipeparts.mkinjections(pipeline, head, injections)

if not os.path.isdir(options.output_path):
	try:
        	os.makedirs(options.output_path)
	except:
        	print >> sys.stderr, "Unable to make directory ", options.output_path
		pass
else:
	print "Target output directory already exists."

## 10) create frames
head = pipeparts.mkframecppchannelmux(pipeline, {"%s:%s" % (instrument, options.output_channel_name): head}, frame_duration = options.frame_duration, frames_per_file = options.frames_per_file)

## 11) Write the frames to disk
head = pipeparts.mkframecppfilesink(pipeline, head, frame_type = options.frame_type)

# Put O(100) frames in each directory
head.connect("notify::timestamp", pipeparts.framecpp_filesink_ldas_path_handler, (options.output_path, digits(options.gps_start_time) - digits(options.frame_duration * options.frames_per_file * 100)))

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

## Debugging output
if "GST_DEBUG_DUMP_DOT_DIR" in os.environ:
	pipeparts.write_dump_dot(pipeline, "%s_PLAYING" % pipeline.get_name(), verbose = True)

	## Setup a signal handler to intercept SIGINT in order to write the pipeline graph at ctrl+C before cleanly shutting down
	class SigHandler(simplehandler.OneTimeSignalHandler):
		def do_on_call(self, signum, frame):
			pipeparts.write_dump_dot(pipeline, "%s_SIGINT" % pipeline.get_name(), verbose = True)

	sighandler = SigHandler(pipeline)

mainloop.run()
