#!/usr/bin/python3.13
#
# 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.

import numpy
from optparse import OptionParser
import os
import sys

import lal

from gstlal import kernels
from gstlal import pipeparts
from gstlal import datasource
from gstlal.psd import read_psd
from gstlal.stream import Stream

### 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
###
### .. graphviz::
###
###    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;
###    }
###
###
### Usage cases
### -----------
###
### 1. 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
###
###
### 2. 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
### https://git.ligo.org/lscsoft/gstlal/raw/master/gstlal/share/early_aligo_asd.txt
### 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
###
###
### 3. 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
###
###
### 4. 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
###
###
### 5. Other things are certainly possible, please add some!
###
### 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
### ------
###


##
# 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 = "Whiten the data by tracking PSD changes, can be combined with --whiten-reference-psd to seed the process.")
	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("--output-frame-type", metavar = "name", help = "Frame type, required")
	parser.add_option("--output-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("--output-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.output_frame_type is None:
		raise ValueError("--frame-type is required")


	return options, filenames


#
# Main
#

options, filenames = parse_command_line()

## Create output directory
os.makedirs(options.output_path, exist_ok=True)

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

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

# 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, we will do them later
injections, gw_data_source.injection_filename = options.injection_file, None

## 1) Initialize stream from data source
# FIXME: fake source causes problems when making large buffers, so block_size needs to be overwritten
gw_data_source.block_size = 8 * options.sample_rate
stream = Stream.from_datasource(
	gw_data_source,
	instrument,
	name=os.path.split(sys.argv[0])[1],
	verbose=options.verbose
)

## 2) Add progress report to stream if verbose
if options.verbose:
	stream = stream.progressreport("frames")

if options.shift is not None:
	## 3) Apply a time shift if requested with a shift() element
	stream = stream.shift(shift=options.shift)

	## 4) Add progress report to stream if verbose
	stream = stream.progressreport("frames_shifted")

if options.whiten_reference_psd:
	## If whitening read the PSD
	wpsd = read_psd(options.whiten_reference_psd, verbose = options.verbose)[instrument]
else:
	## else set wpsd to None
	wpsd = None

if options.whiten_reference_psd or options.whiten_track_psd:
	## 5) Whiten the data stream if requested with a condition() element
	stream = stream.condition(options.sample_rate, instrument, psd=wpsd, track_psd=options.whiten_track_psd)
else:
	## 6) Otherwise, resample the stream with resample() and add a capsfilter()
	stream = stream.resample(quality=9).capsfilter(f"audio/x-raw, rate={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":
		renormalize_factor = 1.0 / (pipeparts.audioresample_variance_gain(9, 16384, options.sample_rate))**.5
		stream = stream.audioamplify(renormalize_factor)

# Apply a coloring filter
if options.color_psd:

	## read coloring psd file and convert to an FIR filter
	rpsd = read_psd(options.color_psd, verbose = options.verbose)[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 = numpy.array(rpsd.data.data[:max_sample])
	if len(rpsd_data) < max_sample: # FIXME this is off by one sample, but shouldn't be.
		rpsd_data = numpy.append(rpsd_data, numpy.zeros(max_sample - len(rpsd_data)))

	rpsd = lal.CreateCOMPLEX16FrequencySeries(name = rpsd.name, epoch = rpsd.epoch, f0 = rpsd.f0, deltaF = rpsd.deltaF, length = max_sample, sampleUnits = rpsd.sampleUnits)
	rpsd.data.data = rpsd_data
	
	# create the coloring FIR kernel from reference_psd.psd_to_fir_kernel()
	FIRKernel = kernels.PSDFirKernel()
	fir_matrix, latency, measured_sample_rate = FIRKernel.psd_to_linear_phase_whitening_fir_kernel(rpsd, invert = False)

	## 7) Recolor the data with an FIR recoloring filter
	stream = stream.firbank(latency=latency, fir_matrix=[fir_matrix], block_stride=(1 * options.sample_rate))


## Set the tags for the output data
tagstr = f"units=strain,channel-name={options.output_channel_name},instrument={instrument}"

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

## 9) Add injections into the final fake data
if injections is not None:
	stream = stream.injections(injections)

## 10) create frames
stream = stream.framecppchannelmux(
	channels=f"{instrument}:{options.output_channel_name}",
	frame_duration=options.output_frame_duration,
	frames_per_file=options.output_frames_per_file
)

## 11) Write the frames to disk
framesink = stream.framecppfilesink(frame_type=options.output_frame_type)

# Put O(100000 s) frames in each directory
framesink.connect("notify::timestamp", pipeparts.framecpp_filesink_ldas_path_handler, (options.output_path, 5))

# Run it
stream.start()
