#!/usr/local/bin/env python

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

"""
Run Metropolis MC simulation of simple dimer interacting only by harmonic bond.

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.openmm as openmm
    
#import Scientific.IO.NetCDF as netcdf # for netcdf interface in Scientific
import netCDF4 as netcdf # for netcdf interface provided by netCDF4 in enthought python

import wcadimer
import sampling

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

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

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

if __name__ == "__main__":
    # Create system
    system = openmm.System()

    temperature = 300 * units.kelvin
    kB = units.BOLTZMANN_CONSTANT_kB * units.AVOGADRO_CONSTANT_NA
    kT = kB * temperature
    beta = 1.0 / kT

    # Add particles to system.
    mass = 12.0 * units.amu
    system.addParticle(mass)
    system.addParticle(mass)

    # Add harmonic potential.
    #r0 = 3.5 * units.angstroms
    #sigma = 0.2 * units.angstroms
    r0 = 0.0 * units.nanometers
    sigma = 10.0 * units.angstroms
    # Comptue K from sigma.
    # sigma^(-2) = (beta K)
    # K = sigma^(-2) / beta
    K = 1.0 / (beta * sigma**2)
    force = openmm.HarmonicBondForce()
    force.addBond(0, 1, r0, K)
    print "sigma = %.3f A, r0 = %.3f A, K = %.16f kcal/mol/A**2" % (sigma / units.angstroms, r0 / units.angstroms, K / (units.kilocalories_per_mole / units.angstrom**2))
    system.addForce(force)

    # Create initial coordinates using random positions.
    coordinates = units.Quantity(numpy.zeros([2,3], numpy.float64), units.nanometer)
       
    # Reposition dimer particles at compact minimum.
    coordinates[1,0] = r0+sigma

    # Select platform.
    platform = openmm.Platform.getPlatformByName("Reference")
    #platform = openmm.Platform.getPlatformByName("Cuda")
    
    # Create integrator.
    timestep = 1.0 * units.femtoseconds
    friction_coefficient = 10.0 / units.picosecond
    integrator = openmm.LangevinIntegrator(temperature, friction_coefficient, timestep)
    #integrator = openmm.BrownianIntegrator(temperature, friction_coefficient, timestep)
    #integrator = openmm.VerletIntegrator(timestep)

    # Create context.
    context = openmm.Context(system, integrator, platform)
    context.setPositions(coordinates)

    # Compute potential.
    npoints = 1000 # number of points to evaluate r at
    rmax = r0 + 6.0 * sigma
    potential_filename = 'potential.txt'
    outfile = open(potential_filename, 'w')
    for index in range(npoints):
        r = (float(index) / float(npoints)) * rmax
        coordinates[1,0] = r
        context.setPositions(coordinates)
        state = context.getState(getEnergy=True,getForces=True)
        potential_energy = state.getPotentialEnergy()
        forces = state.getForces(asNumpy=True)
        Fr = forces[1,0]
        outfile.write("%16.8f %16.8f %16.8f\n" % (r / units.angstroms, beta*potential_energy, beta*Fr*units.angstroms))
    outfile.close()
        
    # Reposition dimer particles.
    coordinates = units.Quantity(numpy.zeros([2,3], numpy.float64), units.nanometer)
    coordinates[1,0] = r0+sigma
    context.setPositions(coordinates)

    # Get initial energy.
    state = context.getState(getEnergy=True)
    potential = state.getPotentialEnergy()

    # Run simulation
    distance_filename = 'distances.txt'
    niterations = 10000
    nsteps = 10000
    outfile = open(distance_filename, 'w')
    
    for iteration in range(niterations):
        if numpy.mod(iteration,100)==0: print iteration
        integrator.step(nsteps)
        state = context.getState(getEnergy=True,getPositions=True)
        coordinates = state.getPositions(asNumpy=True)
        potential = state.getPotentialEnergy()
        
        # Compute properties.
        dr = coordinates[1,:] - coordinates[0,:]
        r = norm(dr)
        outfile.write('%16.8f %16.8f %16.8f %16.8f %16.8f\n' % (r / units.angstroms, potential / units.kilocalories_per_mole, dr[0] / units.angstroms, dr[1] / units.angstroms, dr[2] / units.angstroms))
        outfile.flush()
        
    outfile.close()