#!/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 equilibrate(system, temperature, sqrt_kT_over_m, coordinates, platform):
    collision_rate = 20.0 / units.picoseconds
    timestep = 2.0 * units.femtoseconds
    nsteps = 25000

    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)    

    print "Computing energy again."
    context.setPositions(coordinates)
    context.setVelocities(velocities)        
    state = context.getState(getEnergy=True)
    potential_energy = state.getPotentialEnergy()
    print "potential energy: %.3f kcal/mol" % (potential_energy / units.kilocalories_per_mole)
    
    total_energy = compute_energy(context, coordinates, velocities)    

    return [coordinates, velocities]

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

    context.setPositions(positions)
    state = context.getState(getForces=True)
    forces = state.getForces(asNumpy=True)
    return forces

def compute_energy(context, positions, velocities):
    """
    Compute total energy for positions and velocities.
    """
    # Set positions and velocities.
    context.setPositions(positions)
    context.setVelocities(velocities)
    # Compute total energy.
    state = context.getState(getEnergy=True)
    total_energy = state.getPotentialEnergy() + state.getKineticEnergy()

    #print "potential energy: %.3f kcal/mol" % (state.getPotentialEnergy() / units.kilocalories_per_mole)
    #print "kinetic   energy: %.3f kcal/mol" % (state.getKineticEnergy() / units.kilocalories_per_mole)    
    
    return total_energy

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

def norm(n01):
    return n01.unit * numpy.sqrt(numpy.dot(n01/n01.unit, n01/n01.unit))

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

if __name__ == "__main__":

    kB = units.BOLTZMANN_CONSTANT_kB * units.AVOGADRO_CONSTANT_NA

    # Lennard-Jones fluid parameters (argon).
    mass     = 39.9 * units.amu
    sigma    = 3.4 * units.angstrom
    epsilon  = 120.0 * units.kelvin * kB

    # Set temperature, pressure, and collision rate for stochastic thermostats.
    temperature = 0.9 / (kB / epsilon)
    pressure = 4.0 / (sigma**3 / epsilon) / units.AVOGADRO_CONSTANT_NA
    barostat_frequency = 25 # number of steps between MC volume adjustments
    collision_rate = 5.0 / units.picosecond # collision rate for Langevin integrator
    print "temperature = %.1f K, pressure = %.1f atm" % (temperature / units.kelvin, pressure / units.atmospheres)

    timestep = 5.0 * units.femtosecond # timestep for integrtion

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

    netcdf_filename = 'lennard-jones-dimer.nc'

    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 = [2**n for n in range(0,15)] # number of steps per switch
    print nsteps_to_try

    # Create system.     
    [system, coordinates] = testsystems.LennardJonesFluid(nx=6, ny=6, nz=6, sigma=sigma, epsilon=epsilon, mass=mass)    

    # Add dimer potential to first two particles.
    r0 = 2**(1.0/6.0) * sigma
    w = r0/2.0
    dimer_force = openmm.CustomBondForce('h*(1-((r-r0-w)/w)^2)^2;')
    dimer_force.addGlobalParameter('h', 6.0 * kT) # barrier height
    dimer_force.addGlobalParameter('r0', r0) # compact state separation
    dimer_force.addGlobalParameter('w', w) # second minimum is at r0 + 2*w
    dimer_force.addBond(0, 1, [])
    system.addForce(dimer_force)

    # 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

    # Add exclusion.
    nonbonded_force.addException(0, 1, 0.0 * units.elementary_charge**2, 1.0 * units.angstroms, 0.0 * units.kilocalories_per_mole)

    # Move initial particles.
    coordinates[0,:] *= 0.0
    coordinates[1,:] *= 0.0
    coordinates[1,0] = r0
    print coordinates[0,:]
    print coordinates[1,:]    
    print "distance = %s" % str(norm(coordinates[1,:] - coordinates[0,:]))

    # 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")

    # Minimize energy.
    print "Minimizing energy..."
    print system.getDefaultPeriodicBoxVectors()
    coordinates = minimize(platform, system, coordinates)
    print "distance = %s" % str(norm(coordinates[1,:] - coordinates[0,:]))
    
    # 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.
        barostat.setFrequency(barostat_frequency)
        [coordinates, velocities] = equilibrate(system, temperature, sqrt_kT_over_m, coordinates, platform)
            
        # Turn off barostat.
        barostat.setFrequency(32767)

        # Create a context.
        print "Creating new context..."
        integrator = openmm.VerletIntegrator(timestep)            
        context = openmm.Context(system, integrator, platform)

        # Draw new Maxwell-Boltzmann velocities.
        print "Randomizing velocities..."
        velocities = sqrt_kT_over_m * numpy.random.standard_normal(size=sqrt_kT_over_m.shape)
        total_energy = compute_energy(context, coordinates, velocities)
        
        # Choose neq move.
        initial_distance = norm(coordinates[1,:] - coordinates[0,:])
        print "initial distance = %s" % str(initial_distance)
        if (numpy.random.rand() < 0.5) and (initial_distance > 2*w):
            final_distance = initial_distance - 2*w
        else:
            final_distance = initial_distance + 2*w
            
        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. 
            #print "distance = %s" % str(norm(proposed_coordinates[1,:] - proposed_coordinates[0,:]))           
            total_energy = compute_energy(context, proposed_coordinates, proposed_velocities)
            initial_energy = total_energy            
            #print "initial energy = %f kT" % (initial_energy / kT)
            
            # Integrate nonequilibrium switching trajectory.
            for step in range(nsteps):
                # Accumulate work contribution.
                last_total_energy = total_energy
                distance = (initial_distance + (final_distance - initial_distance)*float(step+1)/float(nsteps))
                bond_midpoint = (proposed_coordinates[0,:] + proposed_coordinates[1,:]) / 2.0
                n01 = (proposed_coordinates[1,:] - proposed_coordinates[0,:]); n01 /= norm(n01);
                for k in range(3):
                    proposed_coordinates[0,k] = bond_midpoint[k] - float(n01[k]) * distance/2
                    proposed_coordinates[1,k] = bond_midpoint[k] + float(n01[k]) * distance/2
                #print "distance = %s" % str(norm(proposed_coordinates[1,:] - proposed_coordinates[0,:]))
                total_energy = compute_energy(context, 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(context, proposed_coordinates)
                proposed_velocities[2:,:] += 0.5 * forces[2:,:]/masses[2:,:] * timestep 

                # Full-kick positions.
                proposed_coordinates[2:,:] += proposed_velocities[2:,:] * timestep 

                # Half-kick velocities
                forces = compute_forces(context, proposed_coordinates)
                proposed_velocities[2:,:] += 0.5 * forces[2:,:]/masses[2:,:] * 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(context, 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()