#! /usr/bin/env python

# Copyright (C) 2016 Christopher M. Biwer
#
# 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.

import argparse
import h5py
import logging
import matplotlib as mpl; mpl.use("Agg")
import matplotlib.pyplot as plt
import sys
from pycbc import results
from pycbc.io.inference_hdf import InferenceFile

# command line usage
parser = argparse.ArgumentParser(usage="pycbc_inference_plot_acceptance_rate [--options]",
    description="Plots fraction of walkers that accepted each step.")

# add data options
parser.add_argument("--input-file", type=str, required=True,
    help="Path to input HDF file.")

# output plot options
parser.add_argument("--output-file", type=str, required=True,
    help="Path to output plot.")

# add thinning options
parser.add_argument("--thin-start", type=int, default=None,
    help="Sample number to start collecting samples to plot.")
parser.add_argument("--thin-interval", type=int, default=None,
    help="Interval to use for thinning samples.")
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.")

# verbose option
parser.add_argument("--verbose", action="store_true", default=False,
    help="")

# parse the command line
opts = parser.parse_args()

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

# read input file
logging.info("Reading input file")
fp = InferenceFile(opts.input_file, "r")
acceptance_fraction = fp.read_acceptance_fraction(thin_start=opts.thin_start,
                                              thin_interval=opts.thin_interval,
                                              thin_end=opts.thin_end)

# plot acceptance rate and drawn values
logging.info("Plotting acceptance fraction")
fig = plt.figure()
plt.plot(acceptance_fraction, 'k', alpha=.5)
plt.ylabel("Mean Acceptance Rate")
plt.xlabel("Iteration")

# save figure with meta-data
caption = """This plot shows the fraction of samples accepted at each
iteration."""
results.save_fig_with_metadata(fig, opts.output_file,
                               cmd=" ".join(sys.argv),
                               title="Acceptance Rate",
                               caption=caption)
plt.close()

# exit
fp.close()
logging.info("Done")