#!/home/travis/build/pycbc-build/environment/bin/python

# Copyright (C) 2017 Collin Capano
#
# 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
#
# =============================================================================
#
"""Extracts some or all samples from an InferenceFile, writing results to a new
InferenceFile.
"""

import os
import argparse
import logging
import numpy
import pycbc
from pycbc.inference import option_utils
from pycbc.io.inference_hdf import InferenceFile

parser = argparse.ArgumentParser(description=__doc__)

parser.add_argument("--input-file", type=str, required=True,
    help="Path to input HDF file.")
parser.add_argument("--output-file", type=str, required=True,
                    help="Output file to create.")
parser.add_argument("--force", action="store_true", default=False,
                    help="If the output-file already exists, overwrite it. "
                         "Otherwise, an IOError is raised.")
parser.add_argument("--posterior-only", action="store_true", default=False,
                    help="Write copied parameter samples and likelihood stats "
                         "as flattened arrays. For example, if the sampler "
                         "wrote a parameter's samples to "
                         "{samples_group}/{param}/walker{x}/, then the "
                         "copied file will have all selected samples from "
                         "all walkers in {samples_group/{param}/. Default is "
                         "for the copied file to have the same structure "
                         "as the input file.")
# options for down selecting parameters
parser.add_argument("--parameters", type=str, nargs="+",
    metavar="PARAM[:NAME]",
    help="Name of parameters to copy. If none provided will copy all of "
         "the variable_args in the input file. If provided, the "
         "parameters can be any of the variable args in "
         "the input file, derived parameters from them, or any function "
         "of them. Syntax for functions is python; any math functions in "
         "the numpy libary may be used. Can optionally also specify a "
         "name for each parameter, in which case the parameter will be "
         "saved with the given name in the output file.")
parser.add_argument("--thin-start", type=int, default=None,
    help="Sample number to start collecting samples to plot. If none "
         "provided, will start at the end of the burn-in.")
parser.add_argument("--thin-interval", type=int,
    default=None,
    help="Interval to use for thinning samples. If none provided, will "
         "use the auto-correlation length found in the file.")
parser.add_argument("--thin-end", type=int, default=None,
    help="Sample number to stop collecting samples to plot. If none "
         "provided, will stop at the last sample from the sampler.")
parser.add_argument("--iteration", type=int, default=None,
    help="Only retrieve the given iteration. To load the last n-th sampe "
         "use -n, e.g., -1 will load the last iteration. This overrides "
         "the thin-start/interval/end options.")
# add additional, sampler-specific options
parser.add_argument("--temperature", nargs='+',
                    help="For parallel-tempered samplers, which temperature "
                         "chain(s) to extract. Options are 'all', or one or "
                         "more indices indicating the temperature chain "
                         "(0 = coldest). Default is to only extract from "
                         "the coldest chain. ")
parser.add_argument("--verbose", action="store_true", default=False,
                    help="Be verbose")

opts = parser.parse_args()

pycbc.init_logging(opts.verbose)

# check that the output doesn't already exist
if os.path.exists(opts.output_file) and not opts.force:
    raise IOError("output file already exists; use --force if you wish to "
                  "overwrite.")

# add any sampler specific options
sampler_kwargs = {}
if opts.temperature is not None:
    temps = opts.temperature
    if temps == ['all']:
        temps = temps[0]
    else:
        temps = map(int, temps)
    sampler_kwargs['temps'] = temps

# open the input and parse parameters
source = InferenceFile(opts.input_file)
parameters, pnames = option_utils.parse_parameters_opt(opts.parameters)

# copy to output
target = source.copy(opts.output_file, parameters=parameters,
                     parameter_names=pnames,
                     posterior_only=opts.posterior_only,
                     thin_start=opts.thin_start, thin_end=opts.thin_end,
                     thin_interval=opts.thin_interval,
                     iteration=opts.iteration,
                     **sampler_kwargs) 

target.close()
