#!/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 gstlal import pipeparts
from gstlal import reference_psd
from gstlal import simplehandler
from gstlal import datasource
from gstlal import multirate_datasource
from optparse import OptionParser, Option
from glue import segments
from glue.ligolw import ligolw
from glue.ligolw import array
from glue.ligolw import param
array.use_in(ligolw.LIGOLWContentHandler)
param.use_in(ligolw.LIGOLWContentHandler)
from glue.ligolw import utils
from pylal.datatypes import LIGOTimeGPS
from pylal import series as lalseries
import pygtk
pygtk.require("2.0")
import gobject
gobject.threads_init()
import pygst
pygst.require("0.10")
import gst
import sys
import os
import numpy

__doc__ = """
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.

Here are some typical usage cases:

1) recoloring existing frame data

	Recoloring can be thought of as first whitening data and then applying a coloring filter.

	gstlal_fake_frames --whiten-reference-psd V1_median_psd.xml.gz --color-psd recolor_psd.xml.gz --data-source frames --output-path /home/channa/rctest/T1300121_V1_EARLY_RECOLORED_V2 --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.V1.cache --channel-name=V1=h_16384Hz --verbose

2) Producing colored Gaussian noise

	gstlal_fake_frames --color-psd recolor_psd.xml.gz --data-source white --output-path /home/channa/rctest/V1_T1300121_V1_EARLY_GAUSSIAN --gps-start-time 966384031 --frame-type V1_T1300121_V1_EARLY_GAUSSIAN --gps-end-time 966389031 --frame-duration 16 --frames-per-file 256 --verbose --channel-name=V1=GAUSSIAN

3) Creating injections into silence

	gstlal_fake_frames --data-source silence --output-path /home/channa/rctest/V1_INJECTIONS --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 ../BNS-MDC1-FIXEDMASS.xml

Other things are certainly possible.
"""

def digits(num):
	return int(numpy.ceil(numpy.log10(float(num))))

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")
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()
gw_data_source = datasource.GWDataSourceInfo(options)

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

#
# Assume instrument is the first and only key of the 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]

#
# don't do injections in mkbasicsrc(), we will do them later
#

injections, options.injections = options.injections, None

#
# Setup the pipeline
#

pipeline = gst.Pipeline(sys.argv[0])
mainloop = gobject.MainLoop()
handler = simplehandler.Handler(mainloop, pipeline)

#
# Get the basic input
#

head = datasource.mkbasicsrc(pipeline, gw_data_source, instrument, verbose = options.verbose)
head = pipeparts.mkprogressreport(pipeline, head, "frames")

#
# Apply a time shift if requested
#

if options.shift is not None:
	head = pipeparts.mkshift(pipeline, head, shift = options.shift)
	head = pipeparts.mkprogressreport(pipeline, head, "frames_shifted")

#
# first whiten and or downsample
#

if options.whiten_reference_psd:
	wpsd = lalseries.read_psd_xmldoc(utils.load_filename(options.whiten_reference_psd, verbose = options.verbose, contenthandler = ligolw.LIGOLWContentHandler))[instrument]
else:
	wpsd = None

if options.whiten_reference_psd or options.whiten_track_psd:
	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:
	head = pipeparts.mkcapsfilter(pipeline, pipeparts.mkresample(pipeline, head, quality = 9), "audio/x-raw-float, rate=%d" % options.sample_rate)

#
# Apply a coloring filter
#

if options.color_psd:

	#
	# read coloring psd file and convert to an FIR filter
	#

	rpsd = lalseries.read_psd_xmldoc(utils.load_filename(options.color_psd, verbose = options.verbose, contenthandler = ligolw.LIGOLWContentHandler))[instrument]
	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]
	fir_matrix, latency, measured_sample_rate = reference_psd.psd_to_fir_kernel(rpsd)
	head = pipeparts.mkfirbank(pipeline, head, latency = latency, fir_matrix = [fir_matrix], block_stride = 32 * options.sample_rate)

#
# Put the units back to strain before writing to frames.  Additionally,
# override the output channel name if provided from the command line
#

tagstr = "units=strain,channel-name=%s,instrument=%s" % (options.output_channel_name, instrument)
head = pipeparts.mktaginject(pipeline, head, tagstr)

#
# Do injections into the final fake data
#

print injections

if injections is not None:
	head = pipeparts.mkinjections(pipeline, head, injections)

#
# write the result to disk
#

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."

head = pipeparts.mkframecppchannelmux(pipeline, {"%s:%s" % (instrument, options.output_channel_name): head}, frame_duration = options.frame_duration, frames_per_file = options.frames_per_file)
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")
mainloop.run()

