#!/usr/local/bin/env python

#=============================================================================================
# MODULE DOCSTRING
#=============================================================================================

"""
Compare platforms using different tolerances.

DESCRIPTION

COPYRIGHT

@author John D. Chodera <jchodera@gmail.com>

All code in this repository is released under the GNU General Public License.

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, see <http://www.gnu.org/licenses/>.

TODO

"""

#=============================================================================================
# GLOBAL IMPORTS
#=============================================================================================

import os
import os.path
import sys
import math
import copy
import time

import numpy

import simtk.unit as units
import simtk.chem.openmm as openmm
import simtk.chem.openmm.extras.testsystems as testsystems
    
#import Scientific.IO.NetCDF as netcdf # for netcdf interface in Scientific
import netCDF4 as netcdf # for netcdf interface provided by netCDF4 in enthought python

#=============================================================================================
# SUBROUTINES
#=============================================================================================

def logsum(a_n):
    """
    Compute log(sum(exp(a_n))) in a numerically stable manner.

    """
    
    a_n = a_n.copy()
    max_arg = numpy.max(a_n)
    return numpy.log(numpy.sum(numpy.exp(a_n - max_arg))) + max_arg

def log_mean_metropolis(u_n):
    """
    Compute log of mean Metropolis acceptance criteria.

    """
    log_acceptance = - u_n
    log_acceptance[log_acceptance > 0.0] = 0.0
    log_mean_acceptance = logsum(log_acceptance) - numpy.log(float(u_n.size))
    return log_mean_acceptance

#=============================================================================================
# MAIN AND TESTS
#=============================================================================================

if __name__ == "__main__":

    filename = 'ion-inversion.nc'
    filename = 'ion-inversion-anion.nc'    

    # Open NetCDF file for reading
    # ncfile = netcdf.NetCDFFile(filename, 'r') # for Scientific.IO.NetCDF
    ncfile = netcdf.Dataset(filename, 'r') # for netCDF4

    heat = ncfile.variables['heat'][:,:].copy()
    work = ncfile.variables['work'][:,:].copy()
    lechner_work = ncfile.variables['lechner_work'][:,:].copy()
    nsteps_to_try = ncfile.variables['nsteps_to_try'][:].copy()
    [K, niterations] = heat.shape
    niterations -= 1

    # Header
    print "HEAT"
    print "%8s" % "",
    for k in range(K):
        print "%8d" % (nsteps_to_try[k]),
    print ""
    # Data
    for iteration in range(niterations):
        print "%8d" % iteration,
        for k in range(K):
            print "%8.1f" % heat[k,iteration],
        print ""                

    # Header
    print "WORK"
    print "%8s" % "",
    for k in range(K):
        print "%8d" % (nsteps_to_try[k]),
    print ""
    # Data
    for iteration in range(niterations):
        print "%8d" % iteration,
        for k in range(K):
            print "%8.1f" % work[k,iteration],
        print ""                

    # Header
    print "LECHNER WORK"
    print "%8s" % "",
    for k in range(K):
        print "%8d" % (nsteps_to_try[k]),
    print ""
    # Data
    for iteration in range(niterations):
        print "%8d" % iteration,
        for k in range(K):
            print "%8.1f" % lechner_work[k,iteration],
        print ""

    # Compute instantaneous acceptance probability.
    log_instantaneous_acceptance = log_mean_metropolis(numpy.squeeze(work[0,:]))

    # Log mean acceptance
    print "LOG MEAN ACCEPTANCE OF LECHNER WORK"
    print "%8s" % "",
    for k in range(K):
        print "%8d" % (nsteps_to_try[k]),
    print ""
    print "%8s" % "",
    log_mean_acceptance = numpy.zeros([K], numpy.float64)
    for k in range(K):
        log_mean_acceptance[k] = log_mean_metropolis(numpy.squeeze(lechner_work[k,:]))                
        print "%8.1f" % (log_mean_acceptance[k]),
    print ""
    
    # Log mean acceptance
    print "LOG RELATIVE EFFICIENCY"
    print "%8s" % "",
    for k in range(K):
        print "%8d" % (nsteps_to_try[k]),
    print ""
    print "%8s" % "",
    log_mean_efficiency = numpy.zeros([K], numpy.float64)
    for k in range(K):
        nsteps = nsteps_to_try[k]
        log_mean_efficiency[k] = log_mean_acceptance[k] - numpy.log(nsteps) - log_instantaneous_acceptance
        print "%8.1f" % (log_mean_efficiency[k]),
    print ""

    # Close netcdf file.
    ncfile.close()