#!/bin/env python
# Copyright (C) 2015 Alexander Harvey Nitz
#
# 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.
""" Make tables describing a missed injection"""
import h5py, argparse, pycbc.version, pycbc.results, sys
import numpy, pycbc.pnutils
from pycbc.detector import Detector

parser = argparse.ArgumentParser()
parser.add_argument('--version', action='version',
    version=pycbc.version.git_verbose_msg)
parser.add_argument('--verbose', action='store_true')
parser.add_argument('--output-file')
parser.add_argument('--injection-file', required=True,
    help="The HDF format injection file. Required")
parser.add_argument('--injection-index', type=int, required=True,
    help="The index of the injection to print out. Required")
parser.add_argument('--n-nearest', type=int,
    help="Optional, used in the title")

args = parser.parse_args()

f = h5py.File(args.injection_file, 'r')
iidx = args.injection_index

# make a table for the coincident information #################################
params = []
headers = []
data = []

labels = {
    'tc': 'End&nbsp;time',
    'dec_chirp_dist': 'Dec.&nbsp;chirp dist (HL)',
    'eff_dist_h'    : 'D<sub>eff</sub>&nbsp;H',
    'eff_dist_l'    : 'D<sub>eff</sub>&nbsp;L',
    'eff_dist_v'    : 'D<sub>eff</sub>&nbsp;V',
    'mass1'      : 'm<sub>1</sub>',
    'mass2'      : 'm<sub>2</sub>',
    'mchirp'     : 'M<sub>c</sub>',
    'eta'        : '&eta;',
    'ra'  : 'RA',
    'dec'   : 'Dec',
    'inclination': '&iota;',
    'spin1x': 's<sub>1x</sub>',
    'spin1y': 's<sub>1y</sub>',
    'spin1z': 's<sub>1z</sub>',
    'spin2x': 's<sub>2x</sub>',
    'spin2y': 's<sub>2y</sub>',
    'spin2z': 's<sub>2z</sub>'
}

params += ['tc']

m1, m2 = f['injections']['mass1'][iidx], f['injections']['mass2'][iidx]
mchirp, eta = pycbc.pnutils.mass1_mass2_to_mchirp_eta(m1, m2)

if 'optimal_snr' in ' '.join(list(f['injections'].keys())):
    ifolist = f.attrs['ifos'].split(' ')
    for ifo in ifolist:
        labels['optimal_snr_%s' % ifo] = 'Opt.&nbsp;SNR %s' % ifo
        params += ['optimal_snr_%s' % ifo]
else:
    eff_dist = {}
    for ifo in ['H1', 'L1', 'V1']:
        eff_dist[ifo] = Detector(ifo).effective_distance(
                             f['injections/distance'][iidx],
                             f['injections/ra'][iidx],
                             f['injections/dec'][iidx],
                             f['injections/polarization'][iidx],
                             f['injections/tc'][iidx],
                             f['injections/inclination'][iidx])

    params += ['dec_chirp_dist', 'eff_dist_h', 'eff_dist_l', 'eff_dist_v']
    dec_dist = max(eff_dist['H1'], eff_dist['L1'])
    dec_chirp_dist = pycbc.pnutils.chirp_distance(dec_dist, mchirp)

params += ['mass1', 'mass2', 'mchirp', 'eta', 'ra', 'dec',
            'inclination', 'spin1x', 'spin1y', 'spin1z', 'spin2x', 'spin2y',
            'spin2z']

for p in params:

    if p in f['injections']:
        data += ["%.2f" % f['injections'][p][iidx]]
    elif 'eff_dist' in p:
        ifo = '%s1' % p.split('_')[-1]
        data += ["%.2f" % eff_dist[ifo.upper()]]
    elif p == 'mchirp':
        data += ["%.2f" % mchirp]
    elif p == 'eta':
        data += ["%.2f" % eta]
    elif p == 'dec_chirp_dist':
        data += ["%.2f" % dec_chirp_dist]
    else:
        print("No data for {}, so skipping".format(p))
        continue

    headers += [labels[p]]

table = numpy.array([data], dtype=str)
html = str(pycbc.results.static_table(table, headers))

tag = ''
if args.n_nearest is not None:
    tag = ':%s' % (args.n_nearest + 1)

pycbc.results.save_fig_with_metadata(html, args.output_file, {},
                        cmd = ' '.join(sys.argv),
                        title = 'Parameters of missed injection' + tag,
                        caption = "Parameters of this injection")
