#!/usr/bin/env python

from pmx.parser import read_and_format
from pmx.estimators import Jarz, JarzGauss, Crooks, BAR
from pmx.analysis import read_dgdl_files, plot_work_dist, ks_norm_test
from pmx.utils import natural_sort
from pmx import __version__
import sys
import os
import time
import numpy as np
import pickle
import argparse
import warnings
from pmx.scripts.cli import check_unknown_cmd

# Constants
kb = 0.00831447215   # kJ/(K*mol)


# ==============================================================================
#                               FUNCTIONS
# ==============================================================================
def _dump_integ_file(outfn, f_lst, w_lst):
    with open(outfn, 'w') as f:
        for fn, w in zip(f_lst, w_lst):
            f.write('{dhdl} {work}\n'.format(dhdl=fn, work=w))


def _data_from_file(fn):
    data = read_and_format(fn, 'sf')
    return map(lambda a: a[1], data)


def _tee(fp, s, quiet=False):
    print(s, file=fp)
    if quiet is False:
        print(s)


def _time_stats(seconds):
    m, s = divmod(seconds, 60)
    h, m = divmod(m, 60)
    return h, m, s


# ==============================================================================
#                      COMMAND LINE OPTIONS AND MAIN
# ==============================================================================
def parse_options():

    parser = argparse.ArgumentParser(description='Calculates free energies '
            'from fast growth thermodynamic integration simulations. '
            'Available methods for free energy estimation: '
            'Crooks Gaussian Intersection (CGI); '
            'Benett Acceptance Ratio (BAR); '
            'Jarzinski equality (JARZ); '
            'Jarzinski with a Gaussian approximatiuon (JARZ_Gauss).')

    exclu1 = parser.add_mutually_exclusive_group()
    # the following exclusion groups are for fA/iA and fB/iB:
    # one can either specify -fA or -iA, but one of them is required.
    exclu2 = parser.add_mutually_exclusive_group()#required=True)
    exclu3 = parser.add_mutually_exclusive_group()#required=True)

    exclu2.add_argument('-fA',
                        metavar='dhdl',
                        dest='filesAB',
                        type=str,
                        help='dhdl.xvg files for the A->B simulations. Use '
                        'wildcard to select multiple xvg files: e.g. "-fa '
                        './forward_results/dhdl*.xvg"',
                        nargs='+')
    exclu3.add_argument('-fB',
                        metavar='dhdl',
                        dest='filesBA',
                        type=str,
                        help='dhdl.xvg files for the B->A simulations Use '
                        'wildcard to select multiple xvg files: e.g. "-fb '
                        './backward_results/dhdl*.xvg"',
                        nargs='+')
    parser.add_argument('-m',
                        metavar='method',
                        type=str.lower,
                        dest='methods',
                        help='Choose one or more estimators to use from the '
                        'available ones: CGI, BAR, JARZ. Default is all.',
                        default=['cgi', 'bar', 'jarz'],
                        nargs='+')
    parser.add_argument('-t',
                        metavar='temperature',
                        dest='temperature',
                        type=float,
                        help='Temperature in Kelvin. Default is 298.15.',
                        default=298.15)
    parser.add_argument('-o',
                        metavar='result file',
                        dest='outfn',
                        type=str,
                        help='Filename of output result file. Default is '
                        '"results.txt."',
                        default='results.txt')
    parser.add_argument('-b',
                        metavar='nboots',
                        dest='nboots',
                        type=int,
                        help='Number of bootstrap samples to use for the '
                        'bootstrap estimate of the standard errors. Default '
                        'is 0 (no bootstrap).',
                        default=100)
    parser.add_argument('-n',
                        metavar='nblocks',
                        dest='nblocks',
                        type=int,
                        help='Number of blocks to divide the data into for '
                        'an estimate of the standard error. You can use this '
                        'when multiple independent equilibrium simulations'
                        'have been run so to estimate the error from the '
                        'repeats. Default is 1 (i.e. no repeats). It assumes '
                        'the dhdl files for each repeat are read in order and '
                        'are contiguous, e.g. dhdl_0 to dhdl_9 is the first '
                        'repeat, dhdl_10 to dhdl_19 is the second one, etc.',
                        default=1)
    parser.add_argument('-w',
                        metavar='plot',
                        dest='wplot',
                        type=str,
                        help='Name of image file showing the distribution of '
                        'work values. Default is "wplot.png". If you want to '
                        'avoid saving this plot, pass "none" to this flag. '
                        'If you choose to calculate the free energy with '
                        'multiple estimators, the dG values shown on the plot '
                        'will be chosen following the hierarchy '
                        'BAR > CGI > JARZ.',
                        default='wplot.png')
    parser.add_argument('--nbins',
                        metavar='',
                        dest='nbins',
                        type=int,
                        help='Number of histograms bins for the plot. '
                        'Default is 20.',
                        default=20)
    parser.add_argument('--dpi',
                        metavar='',
                        dest='dpi',
                        type=int,
                        help='Resolution of the plot. Default is 300.',
                        default=300)
    parser.add_argument('--reverseB',
                        dest='reverseB',
                        help='Whether to reverse the work values for the '
                        'backward (B->A) transformation. This is useful '
                        'when in Gromacs both forward and reverse simulations '
                        'were run from lambda zero to one. '
                        'Default is False.',
                        default=False,
                        action='store_true')
    parser.add_argument('--integ_only',
                        dest='integ_only',
                        help='Whether to do integration only; the integrated '
                        'values are computed and saved, and the program '
                        'terminated. Default is False.',
                        default=False,
                        action='store_true')
    exclu2.add_argument('-iA',
                        metavar='work input',
                        dest='iA',
                        type=str,
                        help='Two-column dat file containing the list of input'
                        ' files and their respective integrated work values '
                        'for the forward (A->B) tranformation.')
    exclu3.add_argument('-iB',
                        metavar='work input',
                        dest='iB',
                        type=str,
                        help='Two-column dat file containing the list of input'
                        ' files and their respective integrated work values '
                        'for the reverse (B->A) tranformation.')
    parser.add_argument('-oA',
                        metavar='work output',
                        dest='oA',
                        type=str,
                        help='File where to save the list of input dhdl'
                        ' files and their respective integrated work values '
                        'for the forward (A->B) tranformation. Default is '
                        '"integA.dat"',
                        default='integA.dat')
    parser.add_argument('-oB',
                        metavar='work output',
                        dest='oB',
                        type=str,
                        help='File where to save the list of input dhdl'
                        ' files and their respective integrated work values '
                        'for the reverse (B->A) tranformation. Default is '
                        '"integB.dat"',
                        default='integB.dat')
    # The following are mutually exclusive options
    exclu1.add_argument('--skip',
                        metavar='',
                        dest='skip',
                        type=int,
                        help='Skip files, i.e. pick every nth work value. '
                        'Default is 1 (all); with 2, every other work value '
                        'is discarded, etc.',
                        default=1)
    exclu1.add_argument('--slice',
                        metavar='',
                        dest='slice',
                        type=int,
                        help='Subset of trajectories to analyze. '
                        'Provide list slice, e.g. "10 50" will '
                        'result in selecting dhdl_files[10:50]. '
                        'Default is all.',
                        default=None,
                        nargs=2)
    exclu1.add_argument('--rand',
                        metavar='',
                        dest='rand',
                        type=int,
                        help='Take a random subset of trajectories. '
                        'Default is None (do not take random subset)',
                        default=None)
    exclu1.add_argument('--index',
                        metavar='',
                        dest='index',
                        type=int,
                        help='Zero-based index of files to analyze (e.g.'
                        ' 0 10 20 50 60). It keeps '
                        'the dhdl.xvg files according to their position in '
                        'the list, sorted according to the filenames. '
                        'Default is None (i.e. all dhdl are used).',
                        default=None,
                        nargs='+')
    parser.add_argument('--prec',
                        metavar='',
                        dest='precision',
                        type=int,
                        help='The decimal precision of the screen/file output.'
                        ' Default is 2.',
                        default=2)
    parser.add_argument('--units',
                        metavar='',
                        dest='units',
                        type=str.lower,
                        help='The units of the output. Choose from "kJ", '
                        '"kcal", "kT". Default is "kJ."',
                        default='kJ',
                        choices=['kj', 'kcal', 'kt'])
    parser.add_argument('--pickle',
                        dest='pickle',
                        help='Whether to save the free energy results from '
                        'the estimators in pickled files. Default is False.',
                        default=False,
                        action='store_true')
    parser.add_argument('--no_ks',
                        dest='do_ks_test',
                        help='Whether to do a Kolmogorov-Smirnov test '
                        'to check whether the Gaussian assumption for CGI '
                        'holds. Default is True; this flag turns it to False.',
                        default=True,
                        action='store_false')
    parser.add_argument('--quiet',
                        dest='quiet',
                        help='Minimal screen output.',
                        default=False,
                        action='store_true')
    parser.add_argument('--sigmoid',
                        metavar='',
                        dest='sigmoid',
                        type=float,
                        help='Define a sigmoidal lambda path.',
                        default=0.0)

    args, unknown = parser.parse_known_args()
    check_unknown_cmd(unknown)

    return args


# ==============================================================================
#                               FUNCTIONS
# ==============================================================================
def main(args):
    """Run the main script.

    Parameters
    ----------
    args : argparse.Namespace
        The command line arguments
    """

    # start timing
    stime = time.time()

    # input arguments
    out = open(args.outfn, 'w')
    T = args.temperature
    sigmoid = args.sigmoid
    skip = args.skip
    prec = args.precision
    methods = args.methods
    reverseB = args.reverseB
    integ_only = args.integ_only
    nboots = args.nboots
    nblocks = args.nblocks
    do_ks_test = args.do_ks_test
    quiet = args.quiet

    #---------- find out whether two or only one direction is provided --------
    filesAB = []
    filesBA = []
    statesProvided = 'AB'
    if (args.iA is None) and (args.iB is None): # no integrated values provided
        if (args.filesAB is None) and (args.filesBA is None): 
            exit('Need to provide dhdl.xvg files or integrated work values')
        elif args.filesAB is None:
                statesProvided = 'B'
                _tee(out, 'Only one directional Jarzynski estimator will be used')
                filesBA = natural_sort(args.filesBA)
        elif args.filesBA is None: 
                statesProvided = 'A'
                _tee(out, 'Only one directional Jarzynski estimator will be used')
                filesAB = natural_sort(args.filesAB)
        else:
            filesAB = natural_sort(args.filesAB)
            filesBA = natural_sort(args.filesBA)
    elif (args.iA is None) or (args.iB is None):
        if args.iA is None:
            statesProvided = 'B'
        else:
            statesProvided = 'A'
        _tee(out, 'Only one directional Jarzynski estimator will be used')

    # -------------------
    # Select output units
    # -------------------
    units = args.units
    if units.lower() == 'kj':
        # kJ is the input from GMX
        unit_fact = 1.
        units = 'kJ/mol'
    elif units == 'kcal':
        unit_fact = 1./4.184
        units = 'kcal/mol'
    elif units.lower() == 'kt':
        unit_fact = 1./(kb*T)
        units = 'kT'
    else:
        exit('No unit type \'%s\' available' % units)

    print("# analyze_dhdl.py, pmx version = %s" % __version__, file=out)
    print("# pwd = %s" % os.getcwd(), file=out)
    print("# %s (%s)" % (time.asctime(), os.environ.get('USER')), file=out)
    print("# command = %s" % ' '.join(sys.argv), file=out)
    _tee(out, "\n", quiet=quiet)

    # ==========
    # Parse Data
    # ==========

    # If list of dhdl.xvg files are provided, parse dhdl
    # --------------------------------------------------
    res_ab = []
    res_ba = []
    if args.iA is None and args.iB is None:
        # If random selection is chosen, do this before reading files and
        # calculating the work values.
        if args.rand is not None:
            if 'A' in statesProvided:
                filesAB = np.random.choice(filesAB, size=args.rand, replace=False)
            if 'B' in statesProvided:
                filesBA = np.random.choice(filesBA, size=args.rand, replace=False)
            _tee(out, 'Selected random subset of %d trajectories.' % args.rand,
                 quiet=quiet)

        # If slice values provided, select the files needed. Again before
        # reading files so speed up the process
        if args.slice is not None:
            first = args.slice[0]
            last = args.slice[1]
            _tee(out, ' First trajectory read: %i' % first, quiet=quiet)
            _tee(out, ' Last trajectory read: %i' % last, quiet=quiet)
            _tee(out, '', quiet=quiet)
            if 'A' in statesProvided:
                filesAB = filesAB[first:last]
            if 'B' in statesProvided:
                filesBA = filesBA[first:last]

        # If index values provided, select the files needed
        if args.index is not None:
            # Avoid index out of range error if "wrong" indices are provided
            if 'A' in statesProvided:
                filesAB = [filesAB[i] for i in args.index if i < len(filesAB)]
            if 'B' in statesProvided:
                filesBA = [filesBA[i] for i in args.index if i < len(filesBA)]
            # ...but warn if this happens
            if any(i > (len(filesAB) - 1) for i in args.index):
                warnings.warn('\nindex out of range for some of your chosen '
                              '\nindices for the forward work values. This means you are'
                              '\ntrying to select input files that are not present.')
            if any(i > (len(filesBA) - 1) for i in args.index):
                warnings.warn('\nindex out of range for some of your chosen'
                              '\nindices for the reverse work values. This means you are'
                              '\ntrying to select input files that are not present.')

        # when skipping start count from end: in this way the last frame is
        # always included, and what can change is the first one
        if 'A' in statesProvided:
            filesAB = list(reversed(filesAB[::-skip]))
        if 'B' in statesProvided:
            filesBA = list(reversed(filesBA[::-skip]))

        # --------------------
        # Now read in the data
        # --------------------
        if quiet is True:
            print('Processing the input data...')
        elif quiet is False:
            print(' ========================================================')
            print('                   PROCESSING THE DATA')
            print(' ========================================================')
            print('  Forward Data')
        if 'A' in statesProvided:
            res_ab = read_dgdl_files(filesAB, lambda0=0,
                                 invert_values=False, verbose=not quiet, sigmoid=sigmoid)
            _dump_integ_file(args.oA, filesAB, res_ab)

        if quiet is False:
            print('  Reverse Data')
        if 'B' in statesProvided:
            res_ba = read_dgdl_files(filesBA, lambda0=1,
                                 invert_values=reverseB, verbose=not quiet, sigmoid=sigmoid)
            _dump_integ_file(args.oB, filesBA, res_ba)

    # If integrated work values are given instead, read those
    # -------------------------------------------------------
    elif args.iA is not None and args.iB is not None:
        if 'A' in statesProvided:
            print('  Reading integrated values (A->B) from', args.iA)
            res_ab.extend(_data_from_file(args.iA))
        if 'B' in statesProvided:
            print('  Reading integrated values (B->A) from', args.iB)
            res_ba.extend(_data_from_file(args.iB))
        # If slice values provided, select the files needed.
        if args.slice is not None:
            first = args.slice[0]
            last = args.slice[1]
            _tee(out, '  Keeping work values from %i to %i.' % (first, last))
            _tee(out, '')
            if 'A' in statesProvided:
                res_ab = res_ab[first:last]
            if 'B' in statesProvided:
                res_ba = res_ba[first:last]

    # If a mix of dhdl.xvg files and integrated work values are given
    # ---------------------------------------------------------------
    # TODO: we are currently not handling this - do we want to add it?
    elif (args.iA is not None and args.filesBA is not None) or (args.iB is not None and args.filesAB is not None):
        raise ValueError('you need to provide either -fA and -fB, or -iA and '
                         '-iB. Mixing dhdl.xvg files and integrated work '
                         'values is not currently supported.')
    # The following code should never get executed, as the case where the input
    # is not provided should have already by handled by argparse within
    # parse_options
    else:
        raise ValueError('you need to provide dhdl.xvg files or integrated '
                         'work values for both forward (-fA or -iA) ',
                         'and reverse (-fB or -iB) transisions.')

    # If asked to only do the integration of dhdl.xvg, exit
    if integ_only is True:
        if quiet is False:
            print('\n    Integration done. Skipping analysis.')
            print('\n    ......done........\n')
        exit()

    # ==============
    # Begin Analysis
    # ==============
    _tee(out, ' ========================================================', quiet=quiet)
    _tee(out, '                       ANALYSIS', quiet=quiet)
    _tee(out, ' ========================================================', quiet=quiet)
    if 'A' in statesProvided:
        _tee(out, '  Number of forward (0->1) trajectories: %d' % len(res_ab), quiet=quiet)
    if 'B' in statesProvided:
        _tee(out, '  Number of reverse (1->0) trajectories: %d' % len(res_ba), quiet=quiet)
    _tee(out, '  Temperature : %.2f K' % T, quiet=quiet)

    # ============================
    # Crooks Gaussian Intersection
    # ============================
    if 'cgi' in methods and 'AB' in statesProvided:
        _tee(out, '\n --------------------------------------------------------', quiet=quiet)
        _tee(out, '             Crooks Gaussian Intersection     ', quiet=quiet)
        _tee(out, ' --------------------------------------------------------', quiet=quiet)

        if quiet is True:
            print('Running CGI analysis...')
        elif quiet is False:
            print('  Calculating Intersection...')

        cgi = Crooks(wf=res_ab, wr=res_ba, nboots=nboots, nblocks=nblocks)
        if args.pickle is True:
            pickle.dump(cgi, open("cgi_results.pkl", "wb"))

        _tee(out, '  CGI: Forward Gauss mean = {m:8.{p}f} {u} '
                  'std = {s:8.{p}f} {u}'.format(m=cgi.mf*unit_fact,
                                                s=cgi.devf*unit_fact,
                                                p=prec, u=units),
             quiet=quiet)
        _tee(out, '  CGI: Reverse Gauss mean = {m:8.{p}f} {u} '
                  'std = {s:8.{p}f} {u}'.format(m=cgi.mr*unit_fact,
                                                s=cgi.devr*unit_fact,
                                                p=prec, u=units),
             quiet=quiet)

        if cgi.inters_bool is False:
            _tee(out, '\n  Gaussians too close for intersection calculation', quiet=quiet)
            _tee(out, '   --> Taking difference of mean values', quiet=quiet)

        _tee(out, '  CGI: dG = {dg:8.{p}f} {u}'.format(dg=cgi.dg*unit_fact,
                                                       p=prec, u=units), quiet=quiet)

        if nboots > 0:
            _tee(out, '  CGI: Std Err (bootstrap:parametric) = {e:8.{p}f} {u}'.format(e=cgi.err_boot1*unit_fact,
                                                                                      p=prec, u=units), quiet=quiet)
            _tee(out, '  CGI: Std Err (bootstrap) = {e:8.{p}f} {u}'.format(e=cgi.err_boot2*unit_fact,
                                                                           p=prec, u=units), quiet=quiet)

        if nblocks > 1:
            _tee(out, '  CGI: Std Err (blocks) = {e:8.{p}f} {u}'.format(e=cgi.err_blocks*unit_fact,
                                                                        p=prec, u=units), quiet=quiet)

    # --------------
    # Normality test
    # --------------
    if do_ks_test and 'AB' in statesProvided:
        if quiet is False:
            print('\n  Running KS-test...')
        q0, lam00, check0, bOk0 = ks_norm_test(res_ab)
        q1, lam01, check1, bOk1 = ks_norm_test(res_ba)

        _tee(out, '    Forward: gaussian quality = %3.2f' % q0, quiet=quiet)
        if bOk0:
            _tee(out, '             ---> KS-Test Ok', quiet=quiet)
        else:
            _tee(out, '             ---> KS-Test Failed. sqrt(N)*Dmax = %4.2f,'
                      ' lambda0 = %4.2f' % (q0, check0), quiet=quiet)
        _tee(out, '    Reverse: gaussian quality = %3.2f' % q1, quiet=quiet)
        if bOk1:
            _tee(out, '             ---> KS-Test Ok', quiet=quiet)
        else:
            _tee(out, '             ---> KS-Test Failed. sqrt(N)*Dmax = %4.2f,'
                      ' lambda0 = %4.2f' % (q1, check1), quiet=quiet)

    # ========================
    # Bennett Acceptance Ratio
    # ========================
    if 'bar' in methods and 'AB' in statesProvided:
        _tee(out, '\n --------------------------------------------------------', quiet=quiet)
        _tee(out, '             Bennett Acceptance Ratio     ', quiet=quiet)
        _tee(out, ' --------------------------------------------------------', quiet=quiet)

        if quiet is True:
            print('Running BAR analysis...')
        elif quiet is False:
            print('  Running Nelder-Mead Simplex algorithm... ')

        bar = BAR(res_ab, res_ba, T=T, nboots=nboots, nblocks=nblocks)
        if args.pickle:
            pickle.dump(bar, open("bar_results.pkl", "wb"))

        _tee(out, '  BAR: dG = {dg:8.{p}f} {u}'.format(dg=bar.dg*unit_fact, p=prec, u=units), quiet=quiet)
        _tee(out, '  BAR: Std Err (analytical) = {e:8.{p}f} {u}'.format(e=bar.err*unit_fact, p=prec, u=units), quiet=quiet)

        if nboots > 0:
            _tee(out, '  BAR: Std Err (bootstrap)  = {e:8.{p}f} {u}'.format(e=bar.err_boot*unit_fact, p=prec, u=units), quiet=quiet)
        if nblocks > 1:
            _tee(out, '  BAR: Std Err (blocks)  = {e:8.{p}f} {u}'.format(e=bar.err_blocks*unit_fact, p=prec, u=units), quiet=quiet)

        _tee(out, '  BAR: Conv = %8.2f' % bar.conv, quiet=quiet)

        if nboots > 0:
            _tee(out, '  BAR: Conv Std Err (bootstrap) = %8.2f' % bar.conv_err_boot, quiet=quiet)
        if nblocks > 1:
            _tee(out, '  BAR: Conv Std Err (blocks) = %8.2f' % bar.conv_err_blocks, quiet=quiet)

    # =========
    # Jarzynski
    # =========
    if 'jarz' in methods:
        _tee(out, '\n --------------------------------------------------------', quiet=quiet)
        _tee(out, '             Jarzynski estimator     ', quiet=quiet)
        _tee(out, ' --------------------------------------------------------', quiet=quiet)

        if quiet is True:
            print('Running Jarz analysis...')
        jarz = Jarz(wf=res_ab, wr=res_ba, T=T, nboots=nboots, nblocks=nblocks, statesProvided=statesProvided)
        if args.pickle:
            pickle.dump(jarz, open("jarz_results.pkl", "wb"))

        if 'A' in statesProvided:
            _tee(out, '  JARZ: dG Forward = {dg:8.{p}f} {u}'.format(dg=jarz.dg_for*unit_fact,
                                                                p=prec, u=units), quiet=quiet)
        if 'B' in statesProvided:
            _tee(out, '  JARZ: dG Reverse = {dg:8.{p}f} {u}'.format(dg=jarz.dg_rev*unit_fact,
                                                                p=prec, u=units), quiet=quiet)
        if 'AB' in statesProvided:
            _tee(out, '  JARZ: dG Mean    = {dg:8.{p}f} {u}'.format(dg=jarz.dg_mean*unit_fact,
                                                                p=prec, u=units), quiet=quiet)
        if nboots > 0:
            if 'A' in statesProvided:
                _tee(out, '  JARZ: Std Err Forward (bootstrap) = {e:8.{p}f} {u}'.format(e=jarz.err_boot_for*unit_fact, p=prec, u=units), quiet=quiet)
            if 'B' in statesProvided:
                _tee(out, '  JARZ: Std Err Reverse (bootstrap) = {e:8.{p}f} {u}'.format(e=jarz.err_boot_rev*unit_fact, p=prec, u=units), quiet=quiet)

        if nblocks > 1:
            if 'A' in statesProvided:
                _tee(out, '  JARZ: Std Err Forward (blocks) = {e:8.{p}f} {u}'.format(e=jarz.err_blocks_for*unit_fact, p=prec, u=units), quiet=quiet)
            if 'B' in statesProvided:
                _tee(out, '  JARZ: Std Err Reverse (blocks) = {e:8.{p}f} {u}'.format(e=jarz.err_blocks_rev*unit_fact, p=prec, u=units), quiet=quiet)

        # -------------------------------------
        # Jarzynski with Gaussian approximation
        # -------------------------------------
        if quiet is True:
            print('Running Jarzynski Gaussian approximation analysis...')
        jarzGauss = JarzGauss(wf=res_ab, wr=res_ba, T=T, nboots=nboots, nblocks=nblocks, statesProvided=statesProvided)
        if args.pickle:
            pickle.dump(jarzGauss, open("jarz_gauss_results.pkl", "wb"))

        if 'A' in statesProvided:
            _tee(out, '  JARZ_Gauss: dG Forward = {dg:8.{p}f} {u}'.format(dg=jarzGauss.dg_for*unit_fact, p=prec, u=units), quiet=quiet)
        if 'B' in statesProvided:
            _tee(out, '  JARZ_Gauss: dG Reverse = {dg:8.{p}f} {u}'.format(dg=jarzGauss.dg_rev*unit_fact,
p=prec, u=units), quiet=quiet)
        if 'AB' in statesProvided:
            _tee(out, '  JARZ_Gauss: dG Mean    = {dg:8.{p}f} {u}'.format(dg=(jarzGauss.dg_for+jarzGauss.dg_rev)/2.0*unit_fact, p=prec, u=units), quiet=quiet)
        if 'A' in statesProvided:
            _tee(out, '  JARZ_Gauss: Std Err (analytical) Forward = {dg:8.{p}f} {u}'.format(dg=jarzGauss.err_for*unit_fact, p=prec, u=units), quiet=quiet)
        if 'B' in statesProvided:
            _tee(out, '  JARZ_Gauss: Std Err (analytical) Reverse = {dg:8.{p}f} {u}'.format(dg=jarzGauss.err_rev*unit_fact, p=prec, u=units), quiet=quiet)
        if nboots > 0:
            if 'A' in statesProvided:
                _tee(out, '  JARZ_Gauss: Std Err Forward (bootstrap) = {e:8.{p}f} {u}'.format(e=jarzGauss.err_boot_for*unit_fact, p=prec, u=units), quiet=quiet)
            if 'B' in statesProvided:
                _tee(out, '  JARZ_Gauss: Std Err Reverse (bootstrap) = {e:8.{p}f} {u}'.format(e=jarzGauss.err_boot_rev*unit_fact, p=prec, u=units), quiet=quiet)

        if nblocks > 1:
            if 'A' in statesProvided:
                _tee(out, '  JARZ_Gauss: Std Err Forward (blocks) = {e:8.{p}f} {u}'.format(e=jarzGauss.err_blocks_for*unit_fact, p=prec, u=units), quiet=quiet)
            if 'B' in statesProvided:
                _tee(out, '  JARZ_Gauss: Std Err Reverse (blocks) = {e:8.{p}f} {u}'.format(e=jarzGauss.err_blocks_rev*unit_fact, p=prec, u=units), quiet=quiet)

    _tee(out, ' ========================================================', quiet=quiet)
    
    #close the outputfile, forcing a flush to disk, so that the unit test can read it properly
    out.close()

    # -----------------------
    # plot work distributions
    # -----------------------
    if args.wplot.lower() != 'none':
        if quiet is False:
            print('\n   Plotting histograms......')
        # hierarchy of estimators: BAR > Crooks > Jarz
        if 'bar' in locals():
            show_dg = bar.dg * unit_fact
            # hierarchy of error estimates : blocks > boots > analytical
            if hasattr(bar, 'err_blocks'):
                show_err = bar.err_blocks * unit_fact
            elif hasattr(bar, 'err_boot') and not hasattr(bar, 'err_blocks'):
                show_err = bar.err_boot * unit_fact
            else:
                show_err = bar.err * unit_fact
            # plot
            plot_work_dist(fname=args.wplot, wf=res_ab, wr=res_ba, dG=show_dg,
                           dGerr=show_err, nbins=args.nbins, dpi=args.dpi,
                           units=units)
        elif 'bar' not in locals() and 'cgi' in locals():
            show_dg = cgi.dg * unit_fact
            # hierarchy of error estimates : blocks > boots
            if hasattr(cgi, 'err_blocks'):
                show_err = cgi.err_blocks * unit_fact
            elif hasattr(cgi, 'err_boot2') and not hasattr(cgi, 'err_blocks'):
                show_err = cgi.err_boot2 * unit_fact
            else:
                show_err = None
            # plot
            plot_work_dist(fname=args.wplot, wf=res_ab, wr=res_ba, dG=show_dg,
                           dGerr=show_err, nbins=args.nbins, dpi=args.dpi,
                           units=units)
        elif 'bar' not in locals() and 'cgi' not in locals() and 'jarz' in locals():
            # for the moment, show values only under specific circumstances
            if hasattr(jarz, 'dg_mean'):
                show_dg = jarz.dg_mean * unit_fact
            elif 'A' in statesProvided:
                show_dg = jarz.dg_for
            elif 'B' in statesProvided:
                show_dg = jarz.dg_rev
            else:
                show_dg = None
            show_err = None
            # plot
            plot_work_dist(fname=args.wplot, wf=res_ab, wr=res_ba, dG=show_dg,
                           dGerr=show_err, nbins=args.nbins, dpi=args.dpi,
                           units=units,statesProvided=statesProvided)

    if quiet is True:
        print('Done')
    if quiet is False:
        print('\n   ......done...........\n')

    if args.pickle and quiet is False:
        print('   NOTE: units of results in pickled files are as in the\n'
              '   provided dhdl.xvg or integ.dat files. These are typically\n'
              '   in kJ/mol when using dhdl.xvg files from Gromacs.\n')
    # execution time
    etime = time.time()
    h, m, s = _time_stats(etime-stime)
    if quiet is False:
        print("   Execution time = %02d:%02d:%02d\n" % (h, m, s))


def entry_point():
    args = parse_options()
    main(args)

if __name__ == '__main__':
    entry_point()
