#!/usr/local/bin/env python

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

"""
Umbrella sampling simulation of WCA dimer in dense solvent.

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
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
#=============================================================================================

kB = units.BOLTZMANN_CONSTANT_kB * units.AVOGADRO_CONSTANT_NA

#=============================================================================================
# WCA Fluid
#=============================================================================================

def WCAFluid(N=512, density=None, mm=None, mass=None, epsilon=None, sigma=None):
    """
    Create a Weeks-Chandler-Andersen system.

    OPTIONAL ARGUMENTS

    N (int) - total number of atoms (default: 150)
    density (float) - N sigma^3 / V (default: 0.96)
    sigma
    epsilon

    """
    
    # Choose OpenMM package.
    if mm is None:
        mm = simtk.chem.openmm

    # Set unit system based on Rowley, Nicholson, and Parsonage argon parameters.
    if mass    is None:  mass        = 39.948 * units.amu # arbitrary reference mass        
    if epsilon is None:  epsilon     = 119.8 * units.kelvin * units.BOLTZMANN_CONSTANT_kB * units.AVOGADRO_CONSTANT_NA # arbitrary reference energy    
    if sigma   is None:  sigma       = 0.3405 * units.nanometers # arbitrary reference lengthscale
    if density is None:  density     = 0.96 / sigma**3

    # Create system
    system = mm.System()

    # Compute total system volume.
    volume = N / density
    
    # Make system cubic in dimension.
    length = volume**(1.0/3.0)
    # TODO: Can we change this to use tuples or 3x3 array?
    a = units.Quantity(numpy.array([1.0, 0.0, 0.0], numpy.float32), units.nanometer) * length/units.nanometer
    b = units.Quantity(numpy.array([0.0, 1.0, 0.0], numpy.float32), units.nanometer) * length/units.nanometer
    c = units.Quantity(numpy.array([0.0, 0.0, 1.0], numpy.float32), units.nanometer) * length/units.nanometer
    print "box edge length = %s" % str(length)
    system.setDefaultPeriodicBoxVectors(a, b, c)

    # Add particles to system.
    for n in range(N):
        system.addParticle(mass)
            
    # WCA: Lennard-Jones truncated at minim and shifted so potential is zero at cutoff.
    energy_expression = '4.0*epsilon*((sigma/r)^12 - (sigma/r)^6) + epsilon;'
    
    # Create force.
    force = mm.CustomNonbondedForce(energy_expression)

    # Set epsilon and sigma global parameters.
    force.addGlobalParameter('epsilon', epsilon)
    force.addGlobalParameter('sigma', sigma)

    # Add particles
    for n in range(N):
        force.addParticle([])
    
    # Set periodic boundary conditions with cutoff.
    force.setNonbondedMethod(mm.CustomNonbondedForce.CutoffPeriodic)
    cutoff = 2.**(1./6.) * sigma # cutoff at minimum of potential
    print "setting cutoff distance to %s" % str(cutoff)
    force.setCutoffDistance(cutoff)    

    # Add nonbonded force term to the system.
    system.addForce(force)

    # Create initial coordinates using random positions.
    coordinates = units.Quantity(numpy.random.rand(N,3), units.nanometer) * (length / units.nanometer)
       
    # Return system and coordinates.
    return [system, coordinates]

#=============================================================================================
# WCA dimer plus flud
#=============================================================================================

def WCADimer(N=512, density=None, mm=None, mass=None, epsilon=None, sigma=None, h=None, r0=None, w=None):
    """
    Create a Weeks-Chandler-Andersen system.

    OPTIONAL ARGUMENTS

    N (int) - total number of atoms (default: 150)
    density (float) - N sigma^3 / V (default: 0.96)
    sigma
    epsilon

    """
    
    # Choose OpenMM package.
    if mm is None:
        mm = simtk.chem.openmm

    # Set unit system based on Rowley, Nicholson, and Parsonage argon parameters.
    if mass    is None:  mass        = 39.948 * units.amu # arbitrary reference mass        
    if epsilon is None:  epsilon     = 119.8 * units.kelvin * units.BOLTZMANN_CONSTANT_kB * units.AVOGADRO_CONSTANT_NA # arbitrary reference energy    
    if sigma   is None:  sigma       = 0.3405 * units.nanometers # arbitrary reference lengthscale
    if density is None:  density     = 0.96 / sigma**3

    # Create system
    system = mm.System()

    # Compute total system volume.
    volume = N / density
    
    # Make system cubic in dimension.
    length = volume**(1.0/3.0)
    # TODO: Can we change this to use tuples or 3x3 array?
    a = units.Quantity(numpy.array([1.0, 0.0, 0.0], numpy.float32), units.nanometer) * length/units.nanometer
    b = units.Quantity(numpy.array([0.0, 1.0, 0.0], numpy.float32), units.nanometer) * length/units.nanometer
    c = units.Quantity(numpy.array([0.0, 0.0, 1.0], numpy.float32), units.nanometer) * length/units.nanometer
    print "box edge length = %s" % str(length)
    system.setDefaultPeriodicBoxVectors(a, b, c)

    # Add particles to system.
    for n in range(N):
        system.addParticle(mass)
            
    # WCA: Lennard-Jones truncated at minim and shifted so potential is zero at cutoff.
    energy_expression = '4.0*epsilon*((sigma/r)^12 - (sigma/r)^6) + epsilon;'
    
    # Create force.
    force = mm.CustomNonbondedForce(energy_expression)

    # Set epsilon and sigma global parameters.
    force.addGlobalParameter('epsilon', epsilon)
    force.addGlobalParameter('sigma', sigma)

    # Add particles
    for n in range(N):
        force.addParticle([])

    # Add exclusion between bonded particles.
    force.addExclusion(0,1)
    
    # Set periodic boundary conditions with cutoff.
    force.setNonbondedMethod(mm.CustomNonbondedForce.CutoffPeriodic)
    cutoff = 2.**(1./6.) * sigma # cutoff at minimum of potential
    print "setting cutoff distance to %s" % str(cutoff)
    force.setCutoffDistance(cutoff)    

    # Add nonbonded force term to the system.
    system.addForce(force)

    # Add dimer potential to first two particles.
    dimer_force = openmm.CustomBondForce('h*(1-((r-r0-w)/w)^2)^2;')
    dimer_force.addGlobalParameter('h', h) # 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)

    # Create initial coordinates using random positions.
    coordinates = units.Quantity(numpy.random.rand(N,3), units.nanometer) * (length / units.nanometer)
       
    # Reposition dimer particles at compact minimum.
    coordinates[0,:] *= 0.0
    coordinates[1,:] *= 0.0
    coordinates[1,0] = r0

    # Return system and coordinates.
    return [system, coordinates]

#=============================================================================================
# WCA dimer plus fluid in umbrella potential
#=============================================================================================

def WCADimerUmbrella(N=512, density=None, mm=None, mass=None, epsilon=None, sigma=None, h=None, r0=None, w=None, rmin=None, rmax=None, kT=None):
    """
    Create a Weeks-Chandler-Andersen system.

    OPTIONAL ARGUMENTS

    N (int) - total number of atoms (default: 150)
    density (float) - N sigma^3 / V (default: 0.96)
    sigma
    epsilon

    """
    
    # Choose OpenMM package.
    if mm is None:
        mm = simtk.chem.openmm

    # Set unit system based on Rowley, Nicholson, and Parsonage argon parameters.
    if mass    is None:  mass        = 39.948 * units.amu # arbitrary reference mass        
    if epsilon is None:  epsilon     = 119.8 * units.kelvin * units.BOLTZMANN_CONSTANT_kB * units.AVOGADRO_CONSTANT_NA # arbitrary reference energy    
    if sigma   is None:  sigma       = 0.3405 * units.nanometers # arbitrary reference lengthscale
    if density is None:  density     = 0.96 / sigma**3

    # Create system
    system = mm.System()

    # Compute total system volume.
    volume = N / density
    
    # Make system cubic in dimension.
    length = volume**(1.0/3.0)
    # TODO: Can we change this to use tuples or 3x3 array?
    a = units.Quantity(numpy.array([1.0, 0.0, 0.0], numpy.float32), units.nanometer) * length/units.nanometer
    b = units.Quantity(numpy.array([0.0, 1.0, 0.0], numpy.float32), units.nanometer) * length/units.nanometer
    c = units.Quantity(numpy.array([0.0, 0.0, 1.0], numpy.float32), units.nanometer) * length/units.nanometer
    print "box edge length = %s" % str(length)
    system.setDefaultPeriodicBoxVectors(a, b, c)

    # Add particles to system.
    for n in range(N):
        system.addParticle(mass)
            
    # WCA: Lennard-Jones truncated at minim and shifted so potential is zero at cutoff.
    energy_expression = '4.0*epsilon*((sigma/r)^12 - (sigma/r)^6) + epsilon;'
    
    # Create force.
    force = mm.CustomNonbondedForce(energy_expression)

    # Set epsilon and sigma global parameters.
    force.addGlobalParameter('epsilon', epsilon)
    force.addGlobalParameter('sigma', sigma)

    # Add particles
    for n in range(N):
        force.addParticle([])

    # Add exclusion between bonded particles.
    force.addExclusion(0,1)
    
    # Set periodic boundary conditions with cutoff.
    force.setNonbondedMethod(mm.CustomNonbondedForce.CutoffPeriodic)
    cutoff = 2.**(1./6.) * sigma # cutoff at minimum of potential
    print "setting cutoff distance to %s" % str(cutoff)
    force.setCutoffDistance(cutoff)    

    # Add nonbonded force term to the system.
    system.addForce(force)

    # Add umbrella potential to first two particles.
    delta = (rmax-rmin) * 0.01 # allowed deviation per kT
    K = kT / delta**2
    gamma = 0.7 * kT / units.angstrom    
    print "delta = %s" % str(delta)
    print "K = %s" % str(K)
    print "gamma = %s" % str(gamma)
    #dimer_force = openmm.CustomBondForce('kT*2.*log(r) + step(r-rmax)*(K/2.)*(r-rmax)^2;')
    #dimer_force = openmm.CustomBondForce('step(r-rmax)*(K/2.)*(r-rmax)^2;')
    #dimer_force = openmm.CustomBondForce('(K/2.)*(r-rmax)^2;')        
    #dimer_force = openmm.CustomBondForce('kT*2.*log(r) + step(r-rmax)*(K/2.)*(r-rmax)^2 + step(rmin-r)*(K/2.)*(r-rmin)^2;')
    dimer_force = openmm.CustomBondForce('kT*2.*log(r) - gamma*r + step(r-rmax)*(K/2.)*(r-rmax)^2 + step(rmin-r)*(K/2.)*(r-rmin)^2;')
    dimer_force.addGlobalParameter('kT', kT) # thermal energy
    dimer_force.addGlobalParameter('gamma', gamma) # start of restraining potential    
    dimer_force.addGlobalParameter('rmin', rmin) # start of restraining potential    
    dimer_force.addGlobalParameter('rmax', rmax) # start of restraining potential    
    dimer_force.addGlobalParameter('K', K) # spring constant for restraining potential
    dimer_force.addBond(0, 1, [])
    system.addForce(dimer_force)

    # Create initial coordinates using random positions.
    coordinates = units.Quantity(numpy.random.rand(N,3), units.nanometer) * (length / units.nanometer)
       
    # Reposition dimer particles.
    coordinates[0,:] *= 0.0
    coordinates[1,:] *= 0.0
    coordinates[1,0] = rmax

    # Return system and coordinates.
    return [system, coordinates]

def equilibrate(system, timestep, collision_rate, temperature, sqrt_kT_over_m, coordinates, platform):
    nsteps = 5000

    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 compute_potential(system, positions, platform):
    """
    Compute potential energy for positions.
    """
    # 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)
    # Compute total energy.
    state = context.getState(getEnergy=True)
    potential_energy = state.getPotentialEnergy() 
    return potential_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 norm(n01):
    return n01.unit * numpy.sqrt(numpy.dot(n01/n01.unit, n01/n01.unit))

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

if __name__ == "__main__":
    # PARAMETERS
    netcdf_filename = 'wca-dimer-umbrella.nc'

    # WCA fluid parameters (argon).
    mass     = 39.9 * units.amu
    sigma    = 3.4 * units.angstrom
    epsilon  = 120.0 * units.kelvin * kB
    r_WCA    = 2.**(1./6.) * sigma
    
    # Compute characteristic timescale.
    tau = units.sqrt(sigma**2 * mass / epsilon)
    print "tau = %.3f ps" % (tau / units.picoseconds)

    # Compute timestep.
    equilibrate_timestep = 0.001 * tau
    print "equilibrate timestep = %.1f fs" % (equilibrate_timestep / units.femtoseconds)

    # Set temperature, pressure, and collision rate for stochastic thermostats.
    temperature = 0.824 / (kB / epsilon)
    print "temperature = %.1f K" % (temperature / units.kelvin)

    kT = kB * temperature # thermal energy
    beta = 1.0 / kT # inverse temperature    
    collision_rate = 1.0 / tau # collision rate for Langevin integrator

    niterations = 20000 # number of work samples to collect

    # Compute particle density.
    rho = 0.96 / sigma**3

    # Dimer potential parameters.
    h = 7.0 * kT
    r0 = r_WCA
    w = 0.5 * r_WCA

    # Create systems.     
    [original_system, coordinates] = WCADimer(N=512, sigma=sigma, epsilon=epsilon, mass=mass, density=rho, h=h, r0=r0, w=w)
    rmin = 2.5 * units.angstroms
    rmax = 9.0 * units.angstroms
    [system, coordinates] = WCADimerUmbrella(N=512, sigma=sigma, epsilon=epsilon, mass=mass, density=rho, h=h, r0=r0, w=w, rmin=rmin, rmax=rmax, kT=kT)    

    # 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

    # 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("OpenCL")
    deviceid = 1
    platform.setPropertyDefaultValue('OpenCLDeviceIndex', '%d' % deviceid)
    platform.setPropertyDefaultValue('CudaDeviceIndex', '%d' % deviceid)         

    # Initialize netcdf file.
    if not os.path.exists(netcdf_filename):
        # Open NetCDF file for writing
        ncfile = netcdf.Dataset(netcdf_filename, 'w') # for netCDF4
        
        # Create dimensions.
        ncfile.createDimension('iteration', 0) # unlimited number of iterations
        ncfile.createDimension('nparticles', nparticles) # number of particles
        ncfile.createDimension('ndim', 3) # number of dimensions    

        # Create variables.
        ncfile.createVariable('distance', 'd', ('iteration',))
        ncfile.createVariable('positions', 'd', ('iteration','nparticles','ndim'))        
        ncfile.createVariable('potential', 'd', ('iteration',))
        ncfile.createVariable('original_potential', 'd', ('iteration',))        
        
        # Force sync to disk to avoid data loss.
        ncfile.sync()

        # Minimize.
        print "Minimizing energy..."
        coordinates = minimize(platform, system, coordinates)
    
        # Equilibrate.
        print "Equilibrating..."
        [coordinates, velocities] = equilibrate(system, equilibrate_timestep, collision_rate, temperature, sqrt_kT_over_m, coordinates, platform)
        iteration = 0

    else:
        # Open NetCDF file for reading.
        ncfile = netcdf.Dataset(netcdf_filename, 'a') # for netCDF4

        # Read iteration and coordinates.
        iteration = ncfile.variables['distance'][:].size
        coordinates = units.Quantity(ncfile.variables['positions'][iteration-1,:,:], units.angstroms)

    # Continue
    while (iteration < niterations):
        print "iteration %5d / %5d" % (iteration, niterations)
        initial_time = time.time()
        
        # Generate a new configuration.
        initial_distance = norm(coordinates[1,:] - coordinates[0,:])
        [coordinates, velocities] = equilibrate(system, equilibrate_timestep, collision_rate, temperature, sqrt_kT_over_m, coordinates, platform)
        final_distance = norm(coordinates[1,:] - coordinates[0,:])            
        print "Dynamics %.1f A -> %.1f A (barrier at %.1f A)" % (initial_distance / units.angstroms, final_distance / units.angstroms, (r0+w)/units.angstroms)

        # Compute potential.
        potential = beta * compute_potential(system, coordinates, platform)                
        original_potential = beta * compute_potential(original_system, coordinates, platform)
        print "r = %8.3f A, U = %12.1f kT, U_original = %12.1f" % (final_distance / units.angstroms, potential, original_potential)

        # Record attempt
        ncfile.variables['distance'][iteration] = final_distance / units.angstroms
        ncfile.variables['positions'][iteration,:,:] = coordinates[:,:] / units.angstroms
        ncfile.variables['potential'][iteration] = potential 
        ncfile.variables['original_potential'][iteration] = original_potential
        ncfile.sync()

        # Debug.
        final_time = time.time()
        elapsed_time = final_time - initial_time
        print "%12.3f s elapsed" % elapsed_time

        # Increment iteration counter.
        iteration += 1

    # Close netcdf file.
    ncfile.close()