# Copyright (C) 2016 Miriam Cabero Mueller
#
# 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 3 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
#
# =============================================================================
#

"""
Create an HDF file with a ringdown injection
"""

import argparse
import h5py
from pycbc.waveform import ringdown

def mode_info_from_input(all_modes, option, input):
    info_modes = [info.split(':')[0] for info in input]
    x, y = {}, {}
    for mode in all_modes:
        if mode not in info_modes:
            raise ValueError('Number of modes in option lmns does not match '
                             'number of modes in option {}'.format(option))
    for info in input:
        try:
            x['%s' %info.split(':')[0]]=info.split(':')[1]
            y['%s' %info.split(':')[0]]=info.split(':')[2]
        except:
            raise ValueError('Information missing in {} '
                             'for mode {}'.format(option,mode))
    return x, y

parser = argparse.ArgumentParser()
parser.add_argument('--output', required=True,
                    help='Path to output hdf file.')
parser.add_argument('--approximant', required=True, 
                    help='Name of the ringdown approximant.')
parser.add_argument('--tc', required=True, type=float,
                    help='Geocentric coalescence time '
                         '(start time of the ringdown injection).')
parser.add_argument('--lmns', nargs='+',
                    help='Modes desired, choose only one for single-mode '
                         'ringdown (lm modes available: 22, 21, 33, 44, 55). '
                         'Example: 222 331 gives the modes 220, 221, and 330.')
parser.add_argument('--final-mass', type=float,
                    help='Mass of the final black hole for '
                         'MassSpin approximants.')
parser.add_argument('--final-spin', type=float,
                    help='Spin of the final black hole for '
                         'MassSpin approximants.')
parser.add_argument('--freqs-taus', nargs='+',
                    help='Central frequency (Hz) and damping time (s) of the '
                         'ringdown for each mode (FreqTau approximants).'
                         'Use format mode:freq:tau.'
                         'Example: 220:317:0.003 221:309:0.001 330:503:0.003')
parser.add_argument('--amps-phis', nargs='+',
                    help='Amplitudes and phases for each mode. '
                         'Use format mode:amplitude:phase. '
                         'Always give amplitude and phase for 220 mode, even if '
                         'not included in lmns, and specifiy amplitudes of '
                         'subdominant modes as fraction of 220 amplitude. '
                         'Example: 220:1e-21:0 221:0.1:3.14 330:0.05:3.14')
parser.add_argument('--t-final', type=float,
                    help='Ending time of the output time series '
                         'for time domain approximants.')
parser.add_argument('--f-lower', type=float,
                    help='Starting frequency of the output freqeucny series '
                         'for frequency domain approximants.')
parser.add_argument('--f-final', type=float, 
                    help='Ending frequency of the output frequency series '
                         'for frequency domain approximants.')
parser.add_argument('--ra', required=True, type=float,
                    help='Right ascension for detector frame waveform generator class.')
parser.add_argument('--dec', required=True, type=float,
                    help='Declination for detector frame waveform generator class.')
parser.add_argument('--polarization', required=True, type=float,
                    help='Polarization for detector frame waveform generator class.')
parser.add_argument('--inclination', type=float,
                    help='Inclination for the spherical harmonics. '
                         'Default is 0 (face on).')
parser.add_argument('--taper', default=None, type=float,
                    help='Taper at the beginning of the time-domain ringdown '
                         'waveform. Duration of the taper will be taper * tau,'
                         ' with tau the damping time of the ringdown.')

opts = parser.parse_args()

# Check that the approximant given is a valid rindown approximant
approxs = ringdown.ringdown_fd_approximants.keys() + \
          ringdown.ringdown_td_approximants.keys()
if opts.approximant not in approxs:
    raise ValueError('Invalid ringdown approximant')

# Check that the taper option is only given with time-domain approximan
if opts.taper is not None \
and opts.approximant in ringdown.ringdown_fd_approximants.keys():
    raise ValueError('The taper option can only be given for time-domain approximants.')

# Write hdf file with the parameters for the injection
injection = h5py.File('%s' %opts.output, 'w')
# Store common arguments
injection.create_dataset('approximant', data=opts.approximant)
injection.create_dataset('tc', data=opts.tc)
injection.create_dataset('lmns', data=opts.lmns)
injection.create_dataset('ra', data=opts.ra)
injection.create_dataset('dec', data=opts.dec)
injection.create_dataset('polarization', data=opts.polarization)
if opts.inclination is not None:
    injection.create_dataset('inclination', data=opts.inclination)
if opts.approximant in ringdown.ringdown_fd_approximants:
    if opts.f_lower is not None:
        injection.create_dataset('f_lower', data=opts.f_lower)
    if opts.f_final is not None:
        injection.create_dataset('f_final', data=opts.f_final)
else:
    if opts.taper is not None:
        injection.create_dataset('taper', data=opts.taper)
    if opts.t_final is not None:
        injection.create_dataset('t_final', data=opts.t_final)

# Amplitudes and phases for each mode have to be given
all_modes = []
for lmn in opts.lmns:
    l, m, nmodes = int(lmn[0]), int(lmn[1]), int(lmn[2])
    [all_modes.append('%d%d%d' %(l,m,n)) for n in range(nmodes)]
amps, phis = mode_info_from_input(all_modes, 'amps-phis', opts.amps_phis)
# Create datasets for each argument
for mode in all_modes:
    injection.create_dataset('amp%s' %mode, data=float(amps[mode]))
    injection.create_dataset('phi%s' %mode, data=float(phis[mode]))
# Amplitude of 220 mode is always required, even if 22 not included
if '220' not in all_modes:
    try:
        injection.create_dataset('amp220', data=float(amps['220']))
    except:
        raise ValueError('Please provide information of 220 mode')

if opts.approximant=='FdQNMfromFinalMassSpin' \
or opts.approximant=='TdQNMfromFinalMassSpin':
    # Check that the necessary arguments are given
    for parameter in ringdown.mass_spin_required_args:
        if getattr(opts, parameter) is None:
            raise ValueError('%s is required' %parameter)
    injection.create_dataset('final_mass', data=opts.final_mass)
    injection.create_dataset('final_spin', data=opts.final_spin)

if opts.approximant=='FdQNMfromFreqTau' \
or opts.approximant=='TdQNMfromFreqTau':
    # Frequencies and damping times for each mode have to be given
    freqs, taus = mode_info_from_input(all_modes, 'freqs-taus', opts.freqs_taus)
    # Create datasets for each argument
    for mode in all_modes:
        injection.create_dataset('f_%s' %mode, data=float(freqs[mode]))
        injection.create_dataset('tau_%s' %mode, data=float(taus[mode]))

injection.close()