#! /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 numpy
import sys
from pycbc import results
from pycbc.filter import autocorrelation
from pycbc.inference import option_utils

# command line usage
parser = argparse.ArgumentParser(
    description="Histograms autocorrelation length from inference samples.")

# verbose option
parser.add_argument("--verbose", action="store_true", default=False,
    help="Print logging info.")
# output plot options
parser.add_argument("--output-file", type=str, required=True,
    help="Path to output plot.")
parser.add_argument("--bins", type=int, default=10,
    help="Number of bins in histogram.")

# add results group
option_utils.add_inference_results_option_group(parser)

# 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)

# load the results
fp, parameters, labels, _ = option_utils.results_from_cli(opts,
    load_samples=False)

# calculate autocorrelation length for each walker
logging.info("Calculating autocorrelation length")
acls = []
for param_name in parameters:

    # loop over walkers and save an autocorrelation length
    # for each walker
    for i in range(fp.nwalkers):
        y = fp.read_samples(param_name, walkers=i, thin_start=opts.thin_start,
                           thin_interval=opts.thin_interval)
        acl = autocorrelation.calculate_acl(y[param_name], dtype=int)
        if acl == numpy.inf:
            acl = fp.niterations
        acls.append( acl )

# plot autocorrelation length
logging.info("Plotting autocorrelation lengths")
fig = plt.figure()
range_max = max(fp.acl, max(acls))
y,x,patches = plt.hist(acls, opts.bins, range=(0,range_max),
                                        histtype="step")

# get histogram bin width
poly_xy = patches[0].get_xy()
step = poly_xy[2][0] - poly_xy[0][0]

plt.xlabel("Iteration")
plt.ylabel(r'Autocorrelation Length for %s'%', '.join(labels))
plt.ylim(0, int(1.1*y.max()))
x_min = max(0, x.min()-2*step)
plt.xlim(x_min, x.max()+2*step)

# plot autocorrelation length saved in InferenceFile
plt.vlines(fp.acl, 0, int(1.1*y.max()))

# save figure with meta-data
caption_kwargs = {
    "parameters" : ", ".join(labels),
}
caption = """ The histogram (blue) is the autocorrelation length (ACL) from all
 the walker chains for the parameters. The vertical black line is the ACL
read from the input file."""
title = "Autocorrelation Length for {parameters}".format(**caption_kwargs)
results.save_fig_with_metadata(fig, opts.output_file,
                               cmd=" ".join(sys.argv),
                               title=title,
                               caption=caption)
plt.close()

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