#!/usr/bin/env python

# Copyright (C) 2019 Ian 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.

"""
Reduce a template bank using some input parameter cuts
"""


from __future__ import print_function
from __future__ import division
import numpy
import logging
import imp
import argparse
import h5py
import pycbc
import pycbc.version

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--version", action="version",
                    version=pycbc.version.git_verbose_msg)
parser.add_argument("-V", "--verbose", action="store_true",
                    help="print extra debugging information", default=False)
parser.add_argument("--input-bank", required=True,
                    help="Input template bank HDF file.")
parser.add_argument("--output-bank", required=True,
                    help="Output template bank HDF file.")
parser.add_argument("--filter-func-file", required=True,
                    help="This can be provided to give a function to define "
                         "which points are covered by the template bank "
                         "bounds, and which are not. The file should contain "
                         "a function called filter_tmpltbank, which should "
                         "take as call profile the template bank HDF object "
                         "and return a boolean (accept=1/reject=0) array.")


opt = parser.parse_args()
pycbc.init_logging(opt.verbose)

bank_fd = h5py.File(opt.input_bank, 'r')

modl = imp.load_source('filter_func', opt.filter_func_file)
func = modl.filter_tmpltbank
bool_arr = func(bank_fd)

logging.info("Downselecting templates. Started with", len(bool_arr),
             "templates, now have ", bool_arr.sum(), "after downselecting.")
bank_ofd = h5py.File(opt.output_bank, 'w')
for name in bank_fd.keys():
    bank_ofd[name] = bank_fd[name][:][bool_arr]
bank_ofd.close()

