#!/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 logging
import sys

import numpy

from matplotlib import use
use('agg')
from matplotlib import pyplot as plt

import pycbc
from pycbc import results
from pycbc.filter import autocorrelation

from pycbc import __version__
from pycbc.inference import io
from pycbc.inference.sampler import samplers

# command line usage
# turn off thin-interval and temps since they won't be used
parser = io.ResultsArgumentParser(skip_args=['thin-interval', 'temps'],
                                  description="Histograms autocorrelation "
                                              "length per walker from an MCMC "
                                              "sampler.")
parser.add_argument("--verbose", action="store_true", default=False,
    help="Print logging info.")
parser.add_argument('--version', action='version', version=__version__,
                    help='show version number and exit')
# 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.")

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

# setup log
pycbc.init_logging(opts.verbose)

# load the results
fp, parameters, labels, _ = io.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_end=opts.thin_end)
        acl = autocorrelation.calculate_acl(y[param_name], dtype=int)
        if acl == numpy.inf:
            acl = fp.niterations
        acls.append(acl)

# get the file's acl
fpacl = fp[fp.sampler_group].attrs['acl']

# plot autocorrelation length
logging.info("Plotting autocorrelation lengths")
fig = plt.figure()
range_max = max(fpacl, 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 hdf file
plt.vlines(fpacl, 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")
