#!/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.openmm as openmm
import simtk.pyopenmm.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 equilibrate(system, temperature, sqrt_kT_over_m, coordinates, platform):
    collision_rate = 10.0 / units.picoseconds
    timestep = 1.0 * units.femtoseconds
    nsteps = 10000

    print "Equilibrating for %.3f ps..." % ((nsteps * timestep) / units.picoseconds)
    
    # Create integrator and context.
    integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep)
    context = openmm.Context(system, integrator, platform)

    # Set coordinates.
    context.setPositions(coordinates)

    # Set Maxwell-Boltzmann velocities
    velocities = sqrt_kT_over_m * numpy.random.standard_normal(size=sqrt_kT_over_m.shape)
    context.setVelocities(velocities)

    # Equilibrate.
    integrator.step(nsteps)

    # Compute energy
    print "Computing energy."
    state = context.getState(getEnergy=True)
    potential_energy = state.getPotentialEnergy()
    print "potential energy: %.3f kcal/mol" % (potential_energy / units.kilocalories_per_mole)

    # Get coordinates.
    state = context.getState(getPositions=True, getVelocities=True)    
    coordinates = state.getPositions(asNumpy=True)
    velocities = state.getVelocities(asNumpy=True)    
    box_vectors = state.getPeriodicBoxVectors()
    system.setDefaultPeriodicBoxVectors(*box_vectors)
    
    return [coordinates, velocities]

def compute_forces(platform, system, positions):
    """
    Compute forces for given positions.
    """

    timestep = 1.0 * units.femtoseconds
    integrator = openmm.VerletIntegrator(timestep)            
    context = openmm.Context(system, integrator, platform)    
    context.setPositions(positions)
    state = context.getState(getForces=True)
    forces = state.getForces(asNumpy=True)
    return forces

def compute_energy(platform, system, positions, velocities):
    """
    Compute total energy for positions and velocities.
    """
    # Create a context.
    timestep = 1.0 * units.femtoseconds
    integrator = openmm.VerletIntegrator(timestep)            
    context = openmm.Context(system, integrator, platform)
    # Set positions and velocities.
    context.setPositions(positions)
    context.setVelocities(velocities)
    # Compute total energy.
    state = context.getState(getEnergy=True)
    total_energy = state.getPotentialEnergy() + state.getKineticEnergy()
    return total_energy

def set_ion_charge(system, ion_charge, ion_atom_indices=[0]):
    """
    Set the ion charge in the given system.

    """

    # Find NonbondedForce.
    nonbonded_force = None
    for force_index in range(system.getNumForces()):
        force = system.getForce(force_index)
        if hasattr(force, 'getParticleParameters'):
            nonbonded_force = force
            break
    
    # Set ion charge.
    for atom_index in ion_atom_indices:
        [old_charge, sigma, epsilon] = nonbonded_force.getParticleParameters(atom_index)
        nonbonded_force.setParticleParameters(atom_index, ion_charge * units.elementary_charge, sigma, epsilon)    

    return

def minimize(platform, system, positions):
    # Create a Context.
    timestep = 1.0 * units.femtoseconds
    integrator = openmm.VerletIntegrator(timestep)
    context = openmm.Context(system, integrator, platform)
    # Set coordinates.
    context.setPositions(positions)
    # Compute initial energy.
    state = context.getState(getEnergy=True)
    initial_potential = state.getPotentialEnergy()
    print "initial potential: %12.3f kcal/mol" % (initial_potential / units.kilocalories_per_mole)
    # Minimize.
    openmm.LocalEnergyMinimizer.minimize(context)    
    # Compute final energy.
    state = context.getState(getEnergy=True, getPositions=True)
    final_potential = state.getPotentialEnergy()
    positions = state.getPositions(asNumpy=True)        
    # Report
    print "final potential  : %12.3f kcal/mol" % (final_potential / units.kilocalories_per_mole)

    return positions

def initialize_netcdf_file(filename, nsteps_to_try):
    """
    Initialize NetCDF file for storage.
    
    """    

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

    # Create dimensions.
    ncfile.createDimension('nsteps_index', len(nsteps_to_try))
    ncfile.createDimension('iteration', 0) # unlimited number of iterations

    ncfile.createVariable('nsteps_to_try', 'i', ('nsteps_index',))
    ncfile.variables['nsteps_to_try'][:] = numpy.array(nsteps_to_try)

    ncfile.createVariable('work', 'd', ('nsteps_index', 'iteration'))
    ncfile.createVariable('heat', 'd', ('nsteps_index', 'iteration'))
    ncfile.createVariable('lechner_work', 'd', ('nsteps_index', 'iteration'))    

    # Force sync to disk to avoid data loss.
    ncfile.sync()

    return ncfile
    
#=============================================================================================
# MAIN AND TESTS
#=============================================================================================

if __name__ == "__main__":

    temperature = 298.0 * units.kelvin
    collision_rate = 20.0 / units.picoseconds
    timestep = 1.0 * units.femtoseconds
    pressure = 1.0 * units.atmospheres

    initial_ion_charge = +1.0
    final_ion_charge = -1.0

    netcdf_filename = 'ion-inversion.nc'

    pdb_filename = 'tip3p900.pdb' # water box PDB filename
    cutoff = 9.0*units.angstroms

    # Create system.
    print "Creating system..."
    [system, coordinates] = testsystems.WaterBox(constrain=False, flexible=True, cutoff=cutoff, nonbonded_method=openmm.NonbondedForce.CutoffPeriodic, filename=pdb_filename, charges=True)

    kB = units.BOLTZMANN_CONSTANT_kB * units.AVOGADRO_CONSTANT_NA
    kT = kB * temperature # thermal energy
    beta = 1.0 / kT # inverse temperature
    
    niterations = 1000 # number of work samples to collect
    #nsteps_to_try = [1, 5, 10, 50, 100, 500, 1000, 5000] # number of steps per switch
    #nsteps_to_try = [1, 5, 10, 50, 100] # number of steps per switch DEBUG
    #nsteps_to_try = [1, 5, 10, 50, 100, 500, 1000, 5000] # number of steps per switch
    nsteps_to_try = [2**n for n in range(0,16)] # number of steps per switch
    print nsteps_to_try
    
    # Initialize netcdf file.
    ncfile = initialize_netcdf_file(netcdf_filename, nsteps_to_try)

    # List all available platforms
    print "Available platforms:"
    for platform_index in range(openmm.Platform.getNumPlatforms()):
        platform = openmm.Platform.getPlatform(platform_index)
        print "%5d %s" % (platform_index, platform.getName())
    print ""

    # Select platform.
    platform = openmm.Platform.getPlatformByName("Cuda")

    # Set initial ion charge.
    set_ion_charge(system, initial_ion_charge)
    
    # Minimize energy.
    print "Minimizing energy..."
    print system.getDefaultPeriodicBoxVectors()
    coordinates = minimize(platform, system, coordinates)

    # Add barostat (which will only be used during equilibration.
    barostat_frequency = 25 # number of steps between MC volume adjustments
    barostat = openmm.MonteCarloBarostat(pressure, temperature, barostat_frequency)
    system.addForce(barostat) 

    # Form vectors of masses and sqrt(kT/m) for force propagation and velocity randomization.
    print "Creating masses array..."
    nparticles = system.getNumParticles()
    masses = units.Quantity(numpy.zeros([nparticles,3], numpy.float64), units.amu)
    for particle_index in range(nparticles):
        masses[particle_index,:] = system.getParticleMass(particle_index)
    kT = kB * temperature # thermal energy    
    sqrt_kT_over_m = units.Quantity(numpy.zeros([nparticles,3], numpy.float64), units.nanometers / units.picosecond)
    for particle_index in range(nparticles):
        sqrt_kT_over_m[particle_index,:] = units.sqrt(kT / masses[particle_index,0]) # standard deviation of velocity distribution for each coordinate for this atom

    # Generate and store work samples for switching from +1 charge to -1 charge.
    print "Opening work file..."
    outfile = open('work.out', 'w') # file for writing work values
    lechner_outfile = open('lechner_work.out', 'w') # file for writing work values
    for iteration in range(niterations):
        print "iteration %5d / %5d" % (iteration, niterations)
        
        # Generate a new configuration.
        set_ion_charge(system, initial_ion_charge)
        barostat.setFrequency(barostat_frequency)
        [coordinates, velocities] = equilibrate(system, temperature, sqrt_kT_over_m, coordinates, platform)

        # Turn off barostat.
        barostat.setFrequency(32767)

        # Draw new Maxwell-Boltzmann velocities.
        velocities = sqrt_kT_over_m * numpy.random.standard_normal(size=sqrt_kT_over_m.shape)
        
        for (nsteps_index, nsteps) in enumerate(nsteps_to_try):
            # Accumulate work and heat.
            work = 0.0 * units.kilocalories_per_mole
            heat = 0.0 * units.kilocalories_per_mole
            
            # Generate initial coordinates and velocities.
            proposed_coordinates = copy.deepcopy(coordinates)
            proposed_velocities = copy.deepcopy(velocities)

            # Compute initial total energy.
            set_ion_charge(system, initial_ion_charge)
            total_energy = compute_energy(platform, system, proposed_coordinates, proposed_velocities)
            initial_energy = total_energy

            # Integrate nonequilibrium switching trajectory.
            for step in range(nsteps):
                # Accumulate work contribution.
                last_total_energy = total_energy
                charge = (initial_ion_charge + (final_ion_charge - initial_ion_charge)*float(step+1)/float(nsteps)) * units.elementary_charge
                set_ion_charge(system, charge)
                total_energy = compute_energy(platform, system, proposed_coordinates, proposed_velocities)
                work += total_energy - last_total_energy
                
                #
                # Velocity Verlet step
                # 

                old_coordinates = copy.deepcopy(proposed_coordinates)
                old_velocities = copy.deepcopy(proposed_velocities)
                
                # Half-kick velocities.
                forces = compute_forces(platform, system, proposed_coordinates)
                proposed_velocities += 0.5 * forces/masses * timestep 

                # Full-kick positions.
                proposed_coordinates += proposed_velocities * timestep 

                # Half-kick velocities
                forces = compute_forces(platform, system, proposed_coordinates)
                proposed_velocities += 0.5 * forces/masses * timestep

                # DEBUG
                #dx = numpy.sqrt((((proposed_coordinates - old_coordinates) / units.nanometers)**2).mean())
                #dv = numpy.sqrt((((proposed_velocities - old_velocities) / (units.nanometers / units.picoseconds))**2).mean())
                #print "dx = %16.8f, dv = %16.8f" % (dx, dv)

                # Accumulate heat contribution.
                last_total_energy = total_energy
                total_energy = compute_energy(platform, system, proposed_coordinates, proposed_velocities)
                heat += total_energy - last_total_energy

                print "step %5d / %5d : energy = %8.1f kT, accumulated work = %8.1f kT, accumulated heat = %8.1f kT" % (step+1, nsteps, total_energy / kT, work / kT, heat / kT)

            # Compute Lechner work.
            final_energy = total_energy
            lechner_work = final_energy - initial_energy
            print "lechner_work = %8.1f kT, work+heat = %8.1f kT" % (lechner_work / kT, (work+heat) / kT)

            # Record data.
            ncfile.variables['work'][nsteps_index,iteration] = work / kT
            ncfile.variables['heat'][nsteps_index,iteration] = heat / kT
            ncfile.variables['lechner_work'][nsteps_index,iteration] = lechner_work / kT
            ncfile.sync()

    # Close netcdf file.
    ncfile.close()