#!/usr/bin/env python
#
# Copyright (C) 2011  Kipp Cannon
#
# 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.


#
# =============================================================================
#
#                                   Preamble
#
# =============================================================================
#


import math
from optparse import OptionParser
import random


import lal


from glue.ligolw import ligolw
from glue.ligolw import lsctables
from glue.ligolw import utils
from glue.ligolw.utils import process as ligolw_process

LIGOTimeGPS = lsctables.LIGOTimeGPS = lal.LIGOTimeGPS


#
# =============================================================================
#
#                                 Command Line
#
# =============================================================================
#


def parse_command_line():
	parser = OptionParser()
	parser.add_option("--gps-start-time", metavar = "seconds", 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", help = "Set the end time of the segment to analyze in GPS seconds (required).  Can be specified to nanosecond precision.")
	parser.add_option("--time-slide-instruments", metavar = "name[,name,...]", default = "H1,H2,L1,V1", help = "Comma-delimited list of the instruments to include in the time-slide vector (default = \"H1,H2,L1,V1\").")
	parser.add_option("--rate", metavar = "Hz", default = .005, help = "Set the mean glitch rate in Hz (default = 0.005 Hz).")
	parser.add_option("--slope", metavar = "value", default = -1.5, help = "Set the exponent of the glitch amplitude distribution's power law (default = -2.0).")
	parser.add_option("-v", "--verbose", action = "store_true", help = "Be verbose (optional).")

	options, filenames = parser.parse_args()

	required_options = ["gps_start_time", "gps_end_time", "time_slide_instruments"]
	missing_options = [option for option in required_options if getattr(options, option) is None]
	if missing_options:
		raise ValueError, "missing required option(s) %s" % ", ".join("--%s" % option.replace("_", "-") for option in sorted(missing_options))

	if len(filenames) != 1:
		raise ValueError, "must provide exactly 1 output filename"

	options.start_time = LIGOTimeGPS(options.gps_start_time)
	options.end_time = LIGOTimeGPS(options.gps_end_time)
	options.instruments = lsctables.instrument_set_from_ifos(options.time_slide_instruments)

	return options, filenames


#
# =============================================================================
#
#                                Random Numbers
#
# =============================================================================
#


def uniform_log(a, b):
	return math.exp(random.uniform(math.log(a), math.log(b)))


def power_law(a, b, n):
	if n == -1:
		raise ValueError, "requires n != -1"
	exponent = 1.0 / (n + 1.0)
	return random.uniform(a**(1.0 / exponent), b**(1.0 / exponent))**exponent


#
# =============================================================================
#
#                                     Main
#
# =============================================================================
#


#
# parse command line
#


options, filenames = parse_command_line()


#
# build document
#


xmldoc = ligolw.Document()
xmldoc.appendChild(ligolw.LIGO_LW())
process_table = xmldoc.childNodes[-1].appendChild(lsctables.New(lsctables.ProcessTable))
xmldoc.childNodes[-1].appendChild(lsctables.New(lsctables.ProcessParamsTable))
time_slide_table = xmldoc.childNodes[-1].appendChild(lsctables.New(lsctables.TimeSlideTable))
sim_burst_table = xmldoc.childNodes[-1].appendChild(lsctables.New(lsctables.SimBurstTable))

process = ligolw_process.register_to_xmldoc(xmldoc, "gstlal_glitch_population", {"gps_start_time": options.gps_start_time, "gps_end_time": options.gps_end_time, "time_slide_instruments": options.time_slide_instruments, "rate": options.rate, "slope": options.slope})

time_slide_table.append_offsetvector(dict.fromkeys(options.instruments, 0.0), process)


#
# generate glitches
#


# number to generate
N = int(round(float(options.end_time - options.start_time) * options.rate))

while len(sim_burst_table) < N:
	# initialize glitch
	glitch = lsctables.SimBurst()
	glitch.waveform = "BTLWNB"
	glitch.pol_ellipse_e = glitch.pol_ellipse_angle = glitch.q = None
	glitch.waveform_number = random.randint(0, 2**64 - 1)	# seed
	glitch.process_id = process.process_id
	glitch.time_slide_id = time_slide_id
	glitch.simulation_id = sim_burst_table.get_next_id()

	# uniform peak time
	glitch.set_time_geocent(options.start_time + random.uniform(0.0, float(options.end_time - options.start_time)))
	glitch.time_geocent_gmst = lal.GreenwichMeanSiderealTime(glitch.get_time_geocent())

	# uniform on sky (this is useless, but can't hurt)
	glitch.ra = random.uniform(0.0, 2 * math.pi)
	glitch.dec = math.asin(random.uniform(-1.0, +1.0))
	glitch.psi = random.uniform(0.0, 2 * math.pi)	# polarization orientation

	# uniform in log bandwidth
	glitch.bandwidth = uniform_log(10.0, 1000.0)

	# uniform in log centre frequency
	glitch.frequency = uniform_log(0.0 + glitch.bandwidth/2, 8000.0 - glitch.bandwidth/2)

	# duration, assume DOF = 2 (sine-Gaussian).  the extra .00001 is to
	# protect against round-off giving the waveform slightly fewer than
	# 2 degrees of freedom
	glitch.duration = 2.00001 / (math.pi * glitch.bandwidth)

	# amplitude power law (note that energy needs to be scaled by the
	# square of the frequency in order to maintain the peak amplitude
	# across the band;  the energy bounds are interpreted as being
	# those at 100 Hz)
	glitch.egw_over_rsquared = power_law(1e-11, 1e-4, options.slope**2) * (glitch.frequency / 100.0)**2
	glitch.amplitude = glitch.hrss = None	# not used

	# add to document
	sim_burst_table.append(glitch)


#
# done
#


utils.write_filename(xmldoc, filenames[0], gz = filenames[0].endswith(".gz"), verbose = options.verbose)
