#ifndef OPENMM_ATMFORCE_H_
#define OPENMM_ATMFORCE_H_

/* -------------------------------------------------------------------------- *
 *                    OpenMM's Alchemical Transfer Force                      *
 * -------------------------------------------------------------------------- *
 * This is a Force of the OpenMM molecular simulation toolkit                 *
 * that implements the Alchemical Transfer Potential                          *
 * for absolute and relative binding free energy estimation                   *
 * (https://doi.org/10.1021/acs.jcim.1c01129). The code is derived from the   *
 * ATMMetaForce plugin                                                        *
 * https://github.com/Gallicchio-Lab/openmm-atmmetaforce-plugin               *
 * with support from the National Science Foundation CAREER 1750511           *
 *                                                                            *
 * Portions copyright (c) 2021-2023 by the Authors                            *
 * Authors: Emilio Gallicchio                                                 *
 * Contributors: Peter Eastman                                                *
 *                                                                            *
 * Permission is hereby granted, free of charge, to any person obtaining a    *
 * copy of this software and associated documentation files (the "Software"), *
 * to deal in the Software without restriction, including without limitation  *
 * the rights to use, copy, modify, merge, publish, distribute, sublicense,   *
 * and/or sell copies of the Software, and to permit persons to whom the      *
 * Software is furnished to do so, subject to the following conditions:       *
 *                                                                            *
 * The above copyright notice and this permission notice shall be included in *
 * all copies or substantial portions of the Software.                        *
 *                                                                            *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR *
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,   *
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL    *
 * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,    *
 * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR      *
 * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE  *
 * USE OR OTHER DEALINGS IN THE SOFTWARE.                                     *
 * -------------------------------------------------------------------------- */

#include "Force.h"
#include "internal/AssertionUtilities.h"
#include <openmm/Vec3.h>
#include <vector>
#include <string>
#include <map>
#include "internal/windowsExport.h"

namespace OpenMM {

/**
 * 
 * The ATMForce class implements the Alchemical Transfer Method (ATM) for OpenMM.
 * ATM is used to compute the binding free energies of molecular complexes and of other equilibrium processes.
 * ATM and its implementation are described in the open access article:
 *
 *    Solmaz Azimi, Sheenam Khuttan, Joe Z. Wu, Rajat K. Pal, and Emilio  Gallicchio. 
 *    Relative Binding Free Energy Calculations for Ligands with Diverse Scaffolds with the Alchemical Transfer Method. 
 *    J. Chem. Inf. Model.  62, 309 (2022)
 *    https://doi.org/10.1021/acs.jcim.1c01129
 *
 * Refer to the publication above for a detailed description of the ATM method and the parameters used in this API 
 * and please cite it to support our work if you use this software in your research.
 *
 * The ATMForce implements an arbitrary potential energy function that depends on the potential
 * energies (u0 and u1) of the system before and after a set of atoms are displaced by some amount.
 * For example, you might displace a molecule from the solvent bulk to a receptor binding site to simulate 
 * a binding process.  The potential energy function typically also depends on one or more parameters that
 * are dialed to implement alchemical transformations.
 *
 * To use this class, create an ATMForce object, passing an algebraic expression to the
 * constructor that defines the potential energy. This expression can be any combination
 * of the variables u0 and u1. Then call addGlobalParameter() to define the parameters on which the potential energy expression depends.
 * The values of global parameters may be modified during a simulation by calling Context::setParameter().
 * Next, call addForce() to add Force objects that define the terms of the potential energy function
 * that change upon displacement. Finally, call addParticle() to specify the coordinate transformation applied to
 * each particle. Currently supported coordinate transformations consist of displacing the positions of particles by a fixed amount
 * or by the offset of the positions between two given particles. As any per-particle parameters, changes in particle coordinate 
 * transformations take effect only after calling updateParametersInContext().
 *
 * As an example, the following code creates an ATMForce based on the change in energy of
 * two particles when the second particle is displaced by 1 nm in the x direction.
 * The energy change is dialed using an alchemical parameter Lambda, which in this case is set to 1/2:
 *
 * \verbatim embed:rst:leading-asterisk
 * <c++>
 * ATMForce *atmforce = new ATMForce("u0 + Lambda*(u1 - u0)");
 * atm->addGlobalParameter("Lambda", 0.5);
 * atm->addParticle();
 * atm->addParticle(new ATMForce::FixedDisplacement(Vec3(1, 0, 0)));
 * CustomBondForce* force = new CustomBondForce("0.5*r^2");
 * atm->addForce(force);
 * </c++>
 * <python>
 * atmforce = ATMForce("u0 + Lambda*(u1 - u0)")
 * atm.addGlobalParameter("Lambda", 0.5)
 * atm.addParticle()
 * atm.addParticle(FixedDisplacement(Vec3(1, 0, 0)))
 * force = CustomBondForce("0.5*r^2")
 * atm.addForce(force)
 * </python>
 * \endverbatim
 *
 * Note that calling addParticle() without arguments is equivalent to a zero fixed displacement.
 *
 * In the example above, the displacement is specified by fixed lab-frame vector. ATMForce also supports variable displacements in internal
 * system coordinates in terms of vector distance between specified particles. For example, if pos[] is the internal array holding the positions
 * of the particles, the following code creates an ATMForce based on the change in energy when the first particle is displaced by the vector 
 * pos[2]-pos[1] going from the second particle to the third particle, 
 *
 * \verbatim embed:rst:leading-asterisk
 * <c++>
 * ATMForce *atmforce = new ATMForce("u0 + Lambda*(u1 - u0)");
 * atm->addGlobalParameter("Lambda", 0.5);
 * atm->addParticle(new ATMForce::ParticleOffsetDisplacement(2, 1));
 * atm->addParticle();
 * atm->addParticle();
 * CustomBondForce* force = new CustomBondForce("0.5*r^2");
 * atm->addForce(force);
 * </c++>
 * <python>
 * atmforce = ATMForce("u0 + Lambda*(u1 - u0)")
 * atm.addGlobalParameter("Lambda", 0.5)
 * atm.addParticle(ParticleOffsetDisplacement(2, 1))
 * atm.addParticle()
 * atm.addParticle()
 * force = CustomBondForce("0.5*r^2")
 * atm.addForce(force)
 * </python>
 * \endverbatim
 *
 * where ParticleOffsetDisplacement is the class that describes this particular type of coordinate transformation.
 *
 * Energy expressions may involve the operators + (add), - (subtract), * (multiply), / (divide), and ^ (power), and the following
 * functions: sqrt, exp, log, sin, cos, sec, csc, tan, cot, asin, acos, atan, atan2, sinh, cosh, tanh, erf, erfc, min, max, abs, floor, ceil, step, delta,
 * select.  All trigonometric functions
 * are defined in radians, and log is the natural logarithm.  step(x) = 0 if x is less than 0, 1 otherwise.  delta(x) = 1 if x is 0, 0 otherwise.
 * select(x,y,z) = z if x = 0, y otherwise.
 *
 * If instead of the energy expression the ATMForce constructor specifies the values of a series of parameters,
 * the default energy expression is used:
 *
 * \verbatim embed:rst:leading-asterisk
 * .. code-block::
 *
 *    select(step(Direction), u0, u1) + ((Lambda2-Lambda1)/Alpha)*log(1+exp(-Alpha*(usc-Uh))) + Lambda2*usc + W0;
 *    usc = select(step(u-Ubcore), (Umax-Ubcore)*fsc+Ubcore, u), u);
 *    fsc = (z^Acore-1)/(z^Acore+1);
 *    z = 1 + 2*(y/Acore) + 2*(y/Acore)^2;
 *    y = (u-Ubcore)/(Umax-Ubcore);
 *    u = select(step(Direction), 1, -1)*(u1-u0)
 * \endverbatim
 *
 * which is the same as the soft-core softplus alchemical potential energy function in the Azimi et al. paper above.
 *
 * The ATMForce is then added to the System as any other Force
 *
 * \verbatim embed:rst:leading-asterisk
 * <c++>
 * system.addForce(atmforce);
 * </c++>
 * <python>
 * system.addForce(atmforce)
 * </python>
 * \endverbatim
 *
 * after which it will be used for energy/force evaluations for molecular dynamics and energy optimization.
 * You can call getPerturbationEnergy() to query the values of u0 and u1, which are needed for computing
 * free energies.
 *
 * In most cases, particles are only displaced in one of the two states evaluated by this force.  It computes the
 * change in energy between the current particle coordinates (as stored in the Context) and the displaced coordinates.
 * In some cases, it is useful to apply displacements to both states.  You can do this by providing two displacement
 * vectors to the fixed displacement transformation given to addParticle(). For example, with:
 *
 * \verbatim embed:rst:leading-asterisk
 * <c++>
 * atm->addParticle(new ATMForce::FixedDisplacement(Vec3(1, 0, 0), Vec3(-1, 0, 0)));
 * </c++>
 * <python>
 * atm.addParticle(FixedDisplacement(Vec3(1, 0, 0), Vec3(-1, 0, 0)))
 * </python>
 * \endverbatim
 * 
 * the energy u1 will be computed after displacing the particle in the positive x direction, and
 * u0 will be computed after displacing it in the negative x direction. Similarly,
 *
 * \verbatim embed:rst:leading-asterisk
 * <c++>
 * atm->addParticle(new ATMForce::ParticleOffsetDisplacement(4, 3, 2, 1));
 * </c++>
 * <python>
 * atm.addParticle(ParticleOffsetDisplacement(4, 3, 2, 1))
 * </python>
 * \endverbatim
 *
 * adds a particle whose position is displaced by pos[4]-pos[3] before calculating u1 and by pos[2]-pos[1] before calculating u0.
 *
 * The ATMForce class has the ability to compute derivatives of the potential energy with respect to global parameters.
 * Call addEnergyParameterDerivative() to request that the derivative with respect to a particular parameter be
 * computed.  You can then query its value in a Context by calling getState() on it.
 */

class OPENMM_EXPORT ATMForce : public OpenMM::Force {
public:
    /**
     * Create an ATMForce object.
     *
     * @param energy   an algebraic expression giving the energy of the system as a function
     *                 of u0 and u1, the energies before and after displacement
     */
    explicit ATMForce(const std::string& energy);
    /**
     * Create an ATMForce object with the default softplus energy expression.  The values passed to
     * this constructor are the default values of the global parameters for newly created Contexts.
     * Their values can be changed by calling setParameter() on the Context using the parameter
     * names defined by the Lambda1(), Lambda2(), etc. methods below.
     *
     * @param lambda1    the default value of the Lambda1 parameter (dimensionless).  This should be
     *                   a number between 0 and 1.
     * @param lambda2    the default value of the Lambda2 parameter (dimensionless).  This should be
     *                   a number between 0 and 1.
     * @param alpha      the default value of the Alpha parameter (kJ/mol)^-1
     * @param uh         the default value of the Uh parameter (kJ/mol)
     * @param w0         the default value of the W0 parameter (kJ/mol)
     * @param umax       the default value of the Umax parameter (kJ/mol)
     * @param ubcore     the default value of the Ubcore parameter (kJ/mol)
     * @param acore      the default value of the Acore parameter dimensionless)
     * @param direction  the default value of the Direction parameter (dimensionless).  This should be
     *                   either 1 for the forward transfer, or -1 for the backward transfer.
     */
    ATMForce(double lambda1, double lambda2, double alpha, double uh, double w0, double umax, double ubcore, double acore, double direction);
    ~ATMForce();

    /**
     * Get the number of particles managed by ATMForce.
     *
     * This should be the same number of particles as the System
     */
    int getNumParticles() const;
    /**
     * Get the number of Forces included in the ATMForce.
     */
    int getNumForces() const;
    /**
     * Get the number of global parameters that the interaction depends on.
     */
    int getNumGlobalParameters() const;
    /**
     * Get the number of global parameters with respect to which the derivative of the energy
     * should be computed.
     */
    int getNumEnergyParameterDerivatives() const;
    /**
     * Get the algebraic expression that gives the energy of the system
     */
    const std::string& getEnergyFunction() const;
    /**
     * Set the algebraic expression that gives the energy of the system
     */
    void setEnergyFunction(const std::string& energy);
    /**
     * Add a Force whose energy will be computed by the ATMForce.
     *
     * @param force  the Force to the be added, which should have been created on the heap with the
     *               "new" operator.  The ATMForce takes over ownership of it, and deletes the Force when the
     *               ATMForce itself is deleted.
     * @return       The index within ATMForce of the force that was added
     */
    int addForce(Force* force);
    /**
     * return the force from index
     */
    Force& getForce(int index) const;

    /**
     * Add a stationary particle: one whose coordinate is not transformed
     *
     * All of the particles in the System must be added to the ATMForce in the same order
     * as they appear in the System.
     *
     * @return                 the index of the particle that was added
     */
    int addParticle();

    /**
     * Add a particle to the force with fixed lab frame displacements
     *
     * @deprecated This method exists only for backward compatibility. Use:
     *   addParticle(new ATMFixedDisplacement(displacement1, displacement0))
     *
     * All of the particles in the System must be added to the ATMForce in the same order
     * as they appear in the System.
     *
     * @param displacement1    the displacement of the particle for the target state in nm
     * @param displacement0    the displacement of the particle for the initial state in nm
     * @return                 the index of the particle that was added
     */
    int addParticle(const Vec3& displacement1, const Vec3& displacement0=Vec3());

    class CoordinateTransformation;
    class FixedDisplacement;
    class ParticleOffsetDisplacement;

    /**
     * Add a particle to the force with a coordinate transformation method
     *
     * All of the particles in the System must be added to the ATMForce in the same order
     * as they appear in the System.
     *
     * @param transformation  the pointer to the CoordinateTransformation object, which should have been
     *                        created on the heap with the "new" operator.  The ATMForce takes over
     *                        ownership of it, and deletes the CoordinateTransformation when the ATMForce
     *                        itself is deleted. Currently supported transformations are FixedDisplacement and
     *                        ParticleOffsetDisplacement.
     * @return                the index of the particle that was added
     */
    int addParticle(CoordinateTransformation* transformation);

    /**
     * Get the parameters for a particle
     *
     * @deprecated This method exists only for backward compatibility.  Use:
     *   const ATMForce::CoordinateTransformation& transformation = getParticleTransformation(index);
     *   Vec3 displacement1 = dynamic_cast<const ATMForce::FixedDisplacement*>(&transformation)->getFixedDisplacement1();
     *   Vec3 displacement0 = dynamic_cast<const ATMForce::FixedDisplacement*>(&transformation)->getFixedDisplacement0();
     *
     * @param index           the index in the force for the particle for which to get parameters
     * @param displacement1   the fixed lab-frame displacement of the particle for the target state in nm
     * @param displacement0   the fixed lab-frame displacement of the particle for the initial state in nm
     */
    void getParticleParameters(int index, Vec3& displacement1, Vec3& displacement0) const;

    /**
     * Set the displacements for a particle as fixed lab frame vectors
     *
     * @deprecated This method exists only for backward compatibility. Use:
     *    setParticleTransformation(index, new ATMForce::FixedDisplacement(displacement1, displacement0))
     *
     * @param index           the index in the force of the particle for which to set parameters
     * @param displacement1   the fixed lab-frame displacement of the particle for the target state in nm
     * @param displacement0   the fixed lab-frame displacement of the particle for the initial state in nm
     */
    void setParticleParameters(int index, const Vec3& displacement1, const Vec3& displacement0=Vec3());

    /**
     * Add a new global parameter that the interaction may depend on.  The default value provided to
     * this method is the initial value of the parameter in newly created Contexts.  You can change
     * the value at any time by calling setParameter() on the Context.
     *
     * @param name             the name of the parameter
     * @param defaultValue     the default value of the parameter
     * @return the index of the parameter that was added
     */
    int addGlobalParameter(const std::string& name, double defaultValue);
    /**
     * Get the name of a global parameter.
     *
     * @param index     the index of the parameter for which to get the name
     * @return the parameter name
     */
    const std::string& getGlobalParameterName(int index) const;
    /**
     * Set the name of a global parameter.
     *
     * @param index          the index of the parameter for which to set the name
     * @param name           the name of the parameter
     */
    void setGlobalParameterName(int index, const std::string& name);
    /**
     * Get the default value of a global parameter.
     *
     * @param index     the index of the parameter for which to get the default value
     * @return the parameter default value
     */
    double getGlobalParameterDefaultValue(int index) const;
    /**
     * Set the default value of a global parameter.
     *
     * @param index          the index of the parameter for which to set the default value
     * @param defaultValue   the default value of the parameter
     */
    void setGlobalParameterDefaultValue(int index, double defaultValue);
    /**
     * Request that this Force compute the derivative of its energy with respect to a global parameter.
     * The parameter must have already been added with addGlobalParameter().
     *
     * @param name      the name of the parameter
     */
    void addEnergyParameterDerivative(const std::string& name);
    /**
     * Get the name of a global parameter with respect to which this Force should compute the
     * derivative of the energy.
     *
     * @param index     the index of the parameter derivative, between 0 and getNumEnergyParameterDerivatives()
     * @return the parameter name
     */
    const std::string& getEnergyParameterDerivativeName(int index) const;
    /**
     * Update the per-particle parameters in a Context to match those stored in this Force object.  This method 
     * should be called after updating parameters with setParticleParameters() to copy them over to the Context.
     * The only information this method updates is the values of per-particle parameters.  The number of particles
     * cannot be changed.
     */
    void updateParametersInContext(Context& context);
    /**
     * Returns whether or not this force makes use of periodic boundary conditions.
     */
    bool usesPeriodicBoundaryConditions() const;
    /**
     * Returns the current perturbation energy.
     *
     * @param context  the Context for which to return the energy
     * @param u1       on exit, the energy of the displaced state
     * @param u0       on exit, the energy of the non-displaced state
     * @param energy   on exit, the value of this force's energy function
     */
    void getPerturbationEnergy(Context& context, double& u1, double& u0, double& energy);
    /**
     * Returns the name of the global parameter corresponding to lambda1.  The value assigned to this
     * parameter should be a number between 0 and 1.
     */
    static const std::string& Lambda1() {
        static const std::string key = "Lambda1";
        return key;
    }

    /**
     * Returns the name of the global parameter corresponding to lambda2.  The value assigned to this
     * parameter should be a number between 0 and 1.
     */
    static const std::string& Lambda2() {
        static const std::string key = "Lambda2";
        return key;
    }

    /**
     * Returns the name of the global parameter corresponding to alpha.  The value assigned to this
     * parameter should be in units of (kJ/mol)^-1.
     */
    static const std::string& Alpha() {
        static const std::string key = "Alpha";
        return key;
    }

    /**
     * Returns the name of the global parameter corresponding to uh.  The value assigned to this
     * parameter should be in units of (kJ/mol).
     */
    static const std::string& Uh() {
        static const std::string key = "Uh";
        return key;
    }

    /**
     * Returns the name of the global parameter corresponding to w0.  The value assigned to this
     * parameter should be in units of (kJ/mol).
     */
    static const std::string& W0() {
        static const std::string key = "W0";
        return key;
    }

    /**
     * Returns the name of the global parameter corresponding to umax.  The value assigned to this
     * parameter should be in units of (kJ/mol).
     */
    static const std::string& Umax() {
        static const std::string key = "Umax";
        return key;
    }

    /**
     * Returns the name of the global parameter corresponding to ubcore.  The value assigned to this
     * parameter should be in units of (kJ/mol).
     */
    static const std::string& Ubcore() {
        static const std::string key = "Ubcore";
        return key;
    }

    /**
     * Returns the name of the global parameter corresponding to acore.
     */
    static const std::string& Acore() {
        static const std::string key = "Acore";
        return key;
    }

    /**
     * Returns the name of the global parameter corresponding to direction.  The value assigned to
     * this parameter should be either 1 for the forward transfer, or -1 for the backward transfer.
     */
    static const std::string& Direction() {
        static const std::string key = "Direction";
        return key;
    }

    /**
     * Change the coordinate transformation method for the specified particle
     *
     * @param index           the index of the particle
     * @param transformation  the pointer to the CoordinateTransformation object, which should have been 
     *                        created on the heap with the "new" operator.  The ATMForce takes over 
     *                        ownership of it, and deletes the CoordinateTransformation when the ATMForce 
     *                        itself is deleted.
     */
    void setParticleTransformation(int index, CoordinateTransformation* transformation);

    /**
     * Returns the Transformation object associated with the particle
     *
     * @param index           the index of the particle
     * @return                the CoordinateTransformation object associated with the particle
     */
    const CoordinateTransformation& getParticleTransformation(int index) const;

protected:
  ForceImpl* createImpl() const;
private:
    class ParticleInfo;
    class GlobalParameterInfo;
    std::string energyExpression;
    std::vector<GlobalParameterInfo> globalParameters;
    std::vector<Force *> forces;
    std::vector<ParticleInfo> particles;
    std::vector<int> energyParameterDerivatives;
};

/**
 * This is an internal class used to record information about a particle.
 * @private
 */
class ATMForce::ParticleInfo {
public:
    int index;
    CoordinateTransformation* transformation;
    ParticleInfo() : index(-1), transformation(NULL) {
    }
    ParticleInfo(int index) : index(index), transformation(NULL) {
    }
    ParticleInfo(int index, CoordinateTransformation* transformation) : index(index), transformation(transformation) {
    }
};

/**
 * This is an internal class used to record information about a global parameter.
 * @private
 */
class ATMForce::GlobalParameterInfo {
public:
    std::string name;
    double defaultValue;
    GlobalParameterInfo() {
    }
    GlobalParameterInfo(const std::string& name, double defaultValue) : name(name), defaultValue(defaultValue) {
    }
};

/**
 * The CoordinateTransformation class describes a generic coordinate transformation applied
 * to a particle. It is a virtual base class. Use the derived classes FixedDisplacement and
 * ParticleOffsetDisplacement to define actual coordinate transformations.
 */
class ATMForce::CoordinateTransformation {
protected:
    CoordinateTransformation(){}
public:
    virtual ~CoordinateTransformation(){};
};

/**
 * The FixedDisplacement class describes a coordinate transformation where a particle is displaced by
 * a fixed amount. To use it, create a FixedDisplacement object passing the displacement vectors for the two
 * states evaluated by ATMForce. The first displacement applies to the target state, and the second 
 * to the reference state. The second displacement can be omitted, in which case it is set to zero.
 */
class ATMForce::FixedDisplacement : public ATMForce::CoordinateTransformation {
public:
    FixedDisplacement(const Vec3& displacement1, const Vec3& displacement0=Vec3()) : displ1(displacement1), displ0(displacement0) {
    }
    ~FixedDisplacement() override {}
    const Vec3& getFixedDisplacement1() const {
        return displ1;
    }
    const Vec3& getFixedDisplacement0() const {
        return displ0;
    }
private:
    Vec3 displ1, displ0;
};

/**
 * The ParticleOffsetDisplacement class describes a coordinate transformation in which a particle is displaced by the
 * vector distance between two particles. The displacement is variable because it changes as the two particles move.
 * To use it, create a ParticleOffsetDisplacement passing the indexes, pDestination1 and pOrigin1, respectively, of 
 * the two particles, resulting in the variable displacement pos[pDestination1]-pos[pOrigin1] if the array pos holds the
 * particles' positions. Optionally, a second set of particles, pDestination0 and pOrigin0, can be specified to apply
 * a similar variable displacement at the reference state of the ATMForce.
 */
class ATMForce::ParticleOffsetDisplacement : public ATMForce::CoordinateTransformation {
public:
    ParticleOffsetDisplacement(int pDestination1, int pOrigin1, int  pDestination0 = -1, int pOrigin0 = -1) : pDestination1(pDestination1), pOrigin1(pOrigin1), pDestination0(pDestination0), pOrigin0(pOrigin0) {
    }
    ~ParticleOffsetDisplacement() override {}
    int getDestinationParticle1() const {
        return pDestination1;
    }
    int getOriginParticle1() const {
        return pOrigin1;
    }
    int getDestinationParticle0() const {
        return pDestination0;
    }
    int getOriginParticle0() const {
        return pOrigin0;
    }
 private:
    int pDestination1, pOrigin1, pDestination0, pOrigin0;
};

} // namespace OpenMM

#endif /*OPENMM_ATMFORCE_H_*/
