#!/usr/bin/env python

# Copyright (C) 2011 Ian W. Harry
#
# 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.

"""
Read in a tmpltbank and create a file containing the list of chi parameters
"""

from __future__ import division

import pycbc
import pycbc.version

__author__  = "Ian Harry <ian.harry@astro.cf.ac.uk>"
__version__ = pycbc.version.git_verbose_msg
__date__    = pycbc.version.date
__program__ = "pycbc_bank_verification"

import argparse
import os, sys
import copy
import logging
import numpy
import h5py
from glue.ligolw import ligolw, table, lsctables, utils as ligolw_utils
from pycbc import tmpltbank, psd, strain, pnutils
import matplotlib
matplotlib.use('Agg')
import pylab

# Read command line option
_desc = __doc__[1:]
parser = argparse.ArgumentParser(description=_desc,
           formatter_class=tmpltbank.IndentedHelpFormatterWithNL)

# Begin with code specific options
parser.add_argument("--version", action="version", version=__version__)
parser.add_argument("--verbose", action="store_true", default=False,
                    help="verbose output")
parser.add_argument("--input-bank", action="store", required=True,
                    help="The template bank to use an input.")
parser.add_argument("--output-file", action="store", required=True,
                    help="The ASCII file to create as output. The columns in "
                    "this file will hold: mass1, mass2, spin1z, spin2z, chi_1, "
                    "chi_2, ....")
# Not sure that vary-fupper makes sense in this code.
#parser.add_argument("-V", "--vary-fupper", action="store_true", default=False,
#                    help="Use a variable upper frequency cutoff in laying "
#                    "out the bank.  OPTIONAL.")
#parser.add_argument("--bank-fupper-step", type=float, default=10.,
#                    help="Size of discrete frequency steps used when varying "
#                    "the fupper. If --calculate-ethinca-metric and "
#                    "--ethinca-freq-step are also given, the code will use "
#                    "the smaller of the two step values. OPTIONAL. Units=Hz")
#parser.add_argument("--bank-fupper-formula", default="SchwarzISCO",
#                    choices=["SchwarzISCO","LightRing","ERD"],
#                    help="Frequency cutoff formula for varying fupper. "
#                    "Frequencies will be rounded to the nearest discrete "
#                    "step. OPTIONAL.")
parser.add_argument("--random-seed", action="store", type=int,
                    default=None,
                    help="Random seed to use when calling numpy.random "
                    "functions used in obtaining the principal components in "
                    "parameter space and when translating points back to "
                    "physical space.  If given, the code should give the "
                    "same output when run with the same random seed.")

# Insert the metric calculation options
tmpltbank.insert_metric_calculation_options(parser)

# Insert the mass range options
tmpltbank.insert_mass_range_option_group(parser)

# Insert the PSD options
pycbc.psd.insert_psd_option_group(parser)

# Insert the data reading options
pycbc.strain.insert_strain_option_group(parser)

opts = parser.parse_args()

if opts.verbose:
    log_level = logging.DEBUG
else:
    log_level = logging.WARN
logging.basicConfig(format='%(asctime)s %(message)s', level=log_level)

# Sanity check options
tmpltbank.verify_metric_calculation_options(opts, parser)
metricParams=tmpltbank.metricParameters.from_argparse(opts)
tmpltbank.verify_mass_range_options(opts, parser)
massRangeParams=tmpltbank.massRangeParameters.from_argparse(opts)
psd.verify_psd_options(opts, parser)
if opts.psd_estimation:
    pycbc.strain.verify_strain_options(opts, parser)

# Set random seed if needed
if opts.random_seed is not None:
    numpy.random.seed(opts.random_seed)

# If we are going to use h(t) to estimate a PSD we need h(t)
if opts.psd_estimation:
    logging.info("Obtaining h(t) for PSD generation")
    strain = pycbc.strain.from_cli(opts, pycbc.DYN_RANGE_FAC)
else:
    strain = None

# Get the PSD using the pycbc interface
logging.info("Obtaining PSD")
# Want the number of samples to be a binary number and Nyquist must be above
# opts.f_upper. All this assumes that 1 / deltaF is a binary number
nyquistFreq = 2**numpy.ceil(numpy.log2(opts.f_upper))
numSamples = int(round(nyquistFreq / opts.delta_f)) + 1
psd = pycbc.psd.from_cli(opts, length=numSamples, delta_f=opts.delta_f,
                         low_frequency_cutoff=opts.f_low, strain=strain,
                         dyn_range_factor=pycbc.DYN_RANGE_FAC)
metricParams.psd = psd

# Begin by calculating a metric
logging.info("Calculating metric")
metricParams = tmpltbank.determine_eigen_directions(metricParams,
    vary_fmax=False, vary_density=None)

# Choose the frequency values to use for metric calculation
refFreq = metricParams.fUpper

logging.info("Calculating covariance matrix")

vals = tmpltbank.estimate_mass_range(
    1000000, massRangeParams, metricParams, refFreq, covary=False)
cov = numpy.cov(vals)
evalsCV,evecsCV = numpy.linalg.eig(cov)
metricParams.evecsCV = {}
metricParams.evecsCV[refFreq] = evecsCV

logging.info("Reading template bank.")
if opts.input_bank.endswith(('.xml','.xml.gz','.xmlgz')):
    # dummy class needed for loading LIGOLW files
    class LIGOLWContentHandler(ligolw.LIGOLWContentHandler):
        pass

    lsctables.use_in(LIGOLWContentHandler)

    indoc = ligolw_utils.load_filename(opts.input_bank,
                                       contenthandler=LIGOLWContentHandler)
    template_list = table.get_table(indoc,
                                    lsctables.SnglInspiralTable.tableName)
    mass1 = [template.mass1 for template in template_list]
    mass2 = [template.mass2 for template in template_list]
    spin1z = [template.spin1z for template in template_list]
    spin2z = [template.spin2z for template in template_list]
elif opts.input_bank.endswith(('.h5','.hdf','.hdf5')):
    h5_fp = h5py.File(opts.input_bank, 'r')
    mass1 = h5_fp['mass1'][:]
    mass2 = h5_fp['mass2'][:]
    spin1z = h5_fp['spin1z'][:]
    spin2z = h5_fp['spin2z'][:]
else:
    err_msg = "Don't know how to read extension {}.".format(opts.input_bank)
    raise NotImplementedError(err_msg)

fP = open(opts.output_file, "w")
for idx in xrange(len(mass1)):
    chi_coords = tmpltbank.get_cov_params(mass1[idx], mass2[idx], spin1z[idx],
                                          spin2z[idx], metricParams, refFreq)
    vals = " ".join([str(mass1[idx]), str(mass2[idx]), str(spin1z[idx]),
                     str(spin2z[idx])])
    vals += " "
    vals += " ".join([str(i) for i in chi_coords])
    fP.write("%s\n" %(vals))
fP.close()