#!/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

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

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

#=============================================================================================
# ALCHEMICAL MODIFICATIONS
#=============================================================================================

def build_alchemically_modified_system(reference_system, receptor_atoms, ligand_atoms, annihilate=True):
    """
    Build alchemically-modified system where molecules do not interact.
    
    """

    # Create new system.
    system = openmm.System()

    # Add atoms.
    for atom_index in range(reference_system.getNumParticles()):
        mass = reference_system.getParticleMass(atom_index)
        system.addParticle(mass)

    # Add constraints
    for constraint_index in range(reference_system.getNumConstraints()):
        [iatom, jatom, r0] = reference_system.getConstraintParameters(constraint_index)
        system.addConstraint(iatom, jatom, r0)    

    # Perturb force terms.
    for force_index in range(reference_system.getNumForces()):
        reference_force = reference_system.getForce(force_index)
        # Dispatch forces
        if isinstance(reference_force, openmm.HarmonicBondForce):
            # HarmonicBondForce
            force = openmm.HarmonicBondForce()
            for bond_index in range(reference_force.getNumBonds()):
                # Retrieve parameters.
                [iatom, jatom, r0, K] = reference_force.getBondParameters(bond_index)
                # Annihilate if directed.
                if annihilate and (iatom in ligand_atoms) and (jatom in ligand_atoms):
                    K *= 0.0                    
                # Add bond parameters.
                force.addBond(iatom, jatom, r0, K)
            # Add force to new system.
            system.addForce(force)
        elif isinstance(reference_force, openmm.HarmonicAngleForce):
            # HarmonicAngleForce
            force = openmm.HarmonicAngleForce()
            for angle_index in range(reference_force.getNumAngles()):
                # Retrieve parameters.
                [iatom, jatom, katom, theta0, Ktheta] = reference_force.getAngleParameters(angle_index)
                # Annihilate if directed:
                if annihilate and (iatom in ligand_atoms) and (jatom in ligand_atoms) and (katom in ligand_atoms):
                    Ktheta *= 0.0                                    
                # Add parameters.
                force.addAngle(iatom, jatom, katom, theta0, Ktheta)
            # Add force to system.                
            system.addForce(force)
        elif isinstance(reference_force, openmm.PeriodicTorsionForce):
            # PeriodicTorsionForce
            force = openmm.PeriodicTorsionForce()
            for torsion_index in range(reference_force.getNumTorsions()):
                # Retrieve parmaeters.
                [particle1, particle2, particle3, particle4, periodicity, phase, k] = reference_force.getTorsionParameters(torsion_index)
                # Annihilate if directed:
                if annihilate and (particle1 in ligand_atoms) and (particle2 in ligand_atoms) and (particle3 in ligand_atoms) and (particle4 in ligand_atoms):
                    k *= 0.0                                                    
                # Add parameters.
                force.addTorsion(particle1, particle2, particle3, particle4, periodicity, phase, k)
            # Add force to system.
            system.addForce(force)            
        elif isinstance(reference_force, openmm.NonbondedForce):
            # NonbondedForce
            force = openmm.NonbondedSoftcoreForce()
            for particle_index in range(reference_force.getNumParticles()):
                # Retrieve parameters.
                [charge, sigma, epsilon] = reference_force.getParticleParameters(particle_index)
                # Alchemically modify parameters.
                alchemical_lambda = 1.0
                if particle_index in ligand_atoms:
                    alchemical_lambda = 0.0
                    charge *= 0.0
                    if annihilate:
                        epsilon *= 0.0
                # Add modified particle parameters.
                force.addParticle(charge, sigma, epsilon, alchemical_lambda)
            for exception_index in range(reference_force.getNumExceptions()):
                # Retrieve parameters.
                [iatom, jatom, chargeprod, sigma, epsilon] = reference_force.getExceptionParameters(exception_index)
                # TODO: Alchemically modify parameters.
                if (iatom in ligand_atoms) and (jatom in ligand_atoms):
                    chargeprod *= 0.0
                    if annihilate:
                        epsilon *= 0.0
                # Add modified exception parameters.
                force.addException(iatom, jatom, chargeprod, sigma, epsilon)
            # Set parameters.
            force.setNonbondedMethod( reference_force.getNonbondedMethod() )
            force.setCutoffDistance( reference_force.getCutoffDistance() )
            force.setReactionFieldDielectric( reference_force.getReactionFieldDielectric() )
            force.setEwaldErrorTolerance( reference_force.getEwaldErrorTolerance() )
            # Add force to new system.
            system.addForce(force)
        elif isinstance(reference_force, openmm.GBSAOBCForce):
            # GBSAOBCForce
            force = openmm.GBSAOBCSoftcoreForce()
            for particle_index in range(reference_force.getNumParticles()):
                # Retrieve parameters.
                [charge, radius, scaling_factor] = reference_force.getParticleParameters(particle_index)
                # Alchemically modify parameters.
                nonpolar_scaling_factor = 1.0
                if particle_index in ligand_atoms:                    
                    charge *= 0.0
                    #radius *= 0.0
                    #scaling_factor *= 0.0
                    nonpolar_scaling_factor = 0.0
                    pass
                # Add parameters.
                force.addParticle(charge, radius, scaling_factor, nonpolar_scaling_factor)
            force.setSolventDielectric( reference_force.getSolventDielectric() )
            force.setSoluteDielectric( reference_force.getSoluteDielectric() )
            #force.setCutoffDistance( reference_force.getCutoffDistance() )
            #force.setNonbondedMethod( reference_force.getNonbondedMethod() )
            # Add force to new system.
            system.addForce(force)
        else:
            # Don't add unrecognized forces.
            pass

    return system

def write_file(filename, contents):
    outfile = open(filename, 'w')
    outfile.write(contents)
    outfile.close()
    return

def write_coordinates(filename, coordinates):
    [natoms,dim] = coordinates.shape
    outfile = open(filename, 'w')
    outfile.write('%12d\n' % natoms)
    for atom_index in range(natoms):
        outfile.write('%16.8f %16.8f %16.8f\n' % (coordinates[atom_index,0] / units.nanometers, coordinates[atom_index,1] / units.nanometers, coordinates[atom_index,2] / units.nanometers))
    outfile.close()
    return

#=============================================================================================
# 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
    
    gbmodel = 'OBC'
    #gbmodel = 'none'
    
    # Load amber system.
    prmtop_filename = 'complex.prmtop'
    crd_filename = 'complex.crd'
    import simtk.pyopenmm.amber.amber_file_parser as amber
    reference_system = amber.readAmberSystem(prmtop_filename, mm=openmm, gbmodel=gbmodel, shake='h-bonds')
    coordinates = amber.readAmberCoordinates(crd_filename)

    # Create alchemically-modified system.
    receptor_atoms = range(0,12)
    ligand_atoms = range(12,27)
    system = build_alchemically_modified_system(reference_system, receptor_atoms, ligand_atoms)
    
    # Add harmonic potential.
    iatom = 0
    jatom = 12
    sigma = 3.816371 * units.angstroms
    r0 = 0.0 * units.angstroms
    K = 1.0 / (beta * sigma**2)
    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))

    add_restraint = False
    if add_restraint:
        add_new = False
        if add_new:
            # Add new HarmonicBondForce
            print "Adding restraint as new HarmonicBondForce."
            force = openmm.HarmonicBondForce()
            force.addBond(iatom, jatom, r0, K)
            system.addForce(force)
        else:
            # Add to existing HarmonicBondForce
            print "Appending restraint to existing HarmonicBondForce."
            for force_index in range(system.getNumForces()):
                force = system.getForce(force_index)
                if isinstance(force, openmm.HarmonicBondForce):
                    # We found it!
                    force.addBond(iatom, jatom, r0, K)
                    break
            
    # Load receptor only.
    prmtop_filename = 'receptor.prmtop'
    crd_filename = 'receptor.crd'
    receptor_system = amber.readAmberSystem(prmtop_filename, mm=openmm, gbmodel=gbmodel, shake='h-bonds')
    receptor_coordinates = amber.readAmberCoordinates(crd_filename)

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

    # Create integrator.
    timestep = 1.0 * units.femtoseconds
    friction_coefficient = 1.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)

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

    # Write out .xml files.
    write_file('receptor.xml', simtk.openmm.XmlSerializer.serializeSystem(receptor_system))
    write_file('complex.xml', simtk.openmm.XmlSerializer.serializeSystem(system))

    # Write out coordinates.
    write_coordinates('receptor.xyz', coordinates[receptor_atoms,:])
    write_coordinates('complex.xyz', coordinates[:,:])    

    if not add_restraint:
        del context, integrator
            
        # Compare receptor energies.
        print "Comparing energies with receptor alone..."
        # Create context for receptor.
        reference_platform = openmm.Platform.getPlatformByName("Cuda")
        receptor_integrator = openmm.VerletIntegrator(timestep)
        receptor_context = openmm.Context(receptor_system, receptor_integrator, reference_platform)
        # Compute energy.
        receptor_context.setPositions(coordinates[receptor_atoms,:])
        state = receptor_context.getState(getEnergy=True)
        receptor_potential = state.getPotentialEnergy()
        print "Total potential = %.5f kcal/mol, receptor potential = %.5f kcal/mol, error = %.5f kcal/mol" % (potential / units.kilocalories_per_mole, receptor_potential / units.kilocalories_per_mole, (receptor_potential - potential) / units.kilocalories_per_mole)
        
        del receptor_context, receptor_integrator
        integrator = openmm.LangevinIntegrator(temperature, friction_coefficient, timestep)
        context = openmm.Context(system, integrator, platform)
        context.setPositions(coordinates)            


    # Run simulation
    distance_filename = 'distances.txt'
    niterations = 10000
    nsteps = 1000
    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[jatom,:] - coordinates[iatom,:]
        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()

        if not add_restraint:
            del context, integrator
            
            # Compare receptor energies.
            print "Comparing energies with receptor alone..."
            # Create context for receptor.
            reference_platform = openmm.Platform.getPlatformByName("Cuda")
            receptor_integrator = openmm.VerletIntegrator(timestep)
            receptor_context = openmm.Context(receptor_system, receptor_integrator, reference_platform)
            # Compute energy.
            receptor_context.setPositions(coordinates[receptor_atoms,:])
            state = receptor_context.getState(getEnergy=True)
            receptor_potential = state.getPotentialEnergy()
            print "Total potential = %.5f kcal/mol, receptor potential = %.5f kcal/mol, error = %.5f kcal/mol" % (potential / units.kilocalories_per_mole, receptor_potential / units.kilocalories_per_mole, (receptor_potential - potential) / units.kilocalories_per_mole)

            del receptor_context, receptor_integrator
            integrator = openmm.LangevinIntegrator(temperature, friction_coefficient, timestep)
            context = openmm.Context(system, integrator, platform)
            context.setPositions(coordinates)            
    
    outfile.close()