#!/usr/bin/env python

# Copyright (C) 2017 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.
""" Plots the Geweke convergence diagnositic statistic.
"""

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

from pycbc import __version__
from pycbc.inference import (io, geweke, option_utils)

# add options to command line
parser = io.ResultsArgumentParser(skip_args=['walkers'])

# program-specific

# verbose option
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 options
parser.add_argument("--output-file", type=str, required=True,
                    help="Path to output plot.")
parser.add_argument("--walkers", type=int, nargs="+", default=None,
                    help="Specific walkers to plot. Default is plot "
                         "all walkers.")

# Geweke calculation options
parser.add_argument("--segment-length", type=int, required=True,
                    help="Number of iterations to use in calculation.")
parser.add_argument("--segment-stride", type=int, required=True,
                    help="How many iterations to advance for next calculation "
                         "along chain.")
parser.add_argument("--segment-start", type=int, required=True,
                    help="Start iteration for calculating statistic.")
parser.add_argument("--segment-end", type=int, required=True,
                    help="End iteration for calculating statistic.")
parser.add_argument("--reference-start", type=int, required=True,
                    help="Start of reference segment for "
                         "calculating statistic.")
parser.add_argument("--reference-end", type=int, required=True,
                    help="End of reference segment for "
                         "calculating statistic.")


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

# setup log
pycbc.init_logging(opts.verbose)

# enfore that this is not a single iteration
# since that does not make sense for the Geweke convergence test
if opts.iteration is not None:
    raise ValueError("Cannot use --iteration")

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

# get walkers to plot
walkers = range(fp.nwalkers) if opts.walkers is None else opts.walkers
nwalkers = len(walkers)

# create Figure
fig = plt.figure()

# loop over each parameter
for param, label in zip(params, labels):
    logging.info("Plotting parameter %s", param)
    # loop over walkers
    for j in walkers:
        logging.info("Plotting walker %d of %d", j, nwalkers)

        # plot each walker
        y = fp.samples_from_cli(opts, parameters=param, walkers=j)

        # get samples for this parameter
        vals = y[param]

        # calculate the Geweke convergence statistic
        starts, ends, stats = geweke.geweke(
                                vals, opts.segment_length, opts.segment_stride,
                                opts.segment_end, opts.reference_start,
                                ref_end=opts.reference_end,
                                seg_start=opts.segment_start)

        # plot walker
        for start, end, stat in zip(starts, ends, stats):
            plt.plot([start, end], [stat, stat], "k")

# format plot
plt.ylabel(", ".join(labels))
plt.xlabel("Iteration")

# plot horizontal lines at -1 and 1
plt.hlines([-1, 1], 0, len(vals), "r", linestyles="dashed")

# save figure with meta-data
caption_kwargs = {
    "parameters" : ", ".join(labels),
}
caption = """The Geweke convergence diagnostic statistic for {parameters}
read from the input file.""".format(**caption_kwargs)
title = "Geweke Convergence 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")
