"""
Generated by CHARMM-GUI (http://www.charmm-gui.org)

omm_vfswitch.py

This module contains vdW force-switching fuctions for OpenMM.

Correspondance: jul316@lehigh.edu or wonpil@lehigh.edu
Last update: February 5, 2025
"""

from openmm import *
from openmm.app import *
from openmm.unit import *


def vfswitch(system, psf, inputs):
    r_on = inputs.r_on
    r_off = inputs.r_off

    # custom nonbonded force for force-switch
    chknbfix = False
    for force in system.getForces():
        if isinstance(force, NonbondedForce):
            nonbonded = force
        if (
            isinstance(force, CustomNonbondedForce)
            and force.getNumTabulatedFunctions() == 2
        ):
            nbfix = force
            chknbfix = True

    if not chknbfix:
        # vfswitch w/o NBFIX
        vfswitch = CustomNonbondedForce(
            "select(step(r-Ron),(cr12*rjunk12 - cr6*rjunk6),"
            "(ccnba/r^12-ccnbb/r^6-ccnba/onoff^6+ccnbb/onoff^3));"
            "cr12 = ccnba*ofdif6;"
            "cr6 = ccnbb*ofdif3;"
            "rjunk12 = (1.0/r^6-1.0/Roff6)^2;"
            "rjunk6 = (1.0/r^3-1.0/Roff3)^2;"
            "ccnba = 4.0*epsilon*sigma^12;"
            "ccnbb = 4.0*epsilon*sigma^6;"
            "sigma = sigma1+sigma2;"
            "epsilon = epsilon1*epsilon2;"
            "ofdif6 = Roff6/(Roff6 - Ron^6);"
            "ofdif3 = Roff3/(Roff3 - Ron^3);"
            "onoff = Roff * Ron;"
            "Roff6 = Roff^6;"
            "Roff3 = Roff^3;"
            "Ron = %f;"
            "Roff = %f;" % (r_on, r_off)
        )
        vfswitch.addPerParticleParameter("sigma")
        vfswitch.addPerParticleParameter("epsilon")
        vfswitch.setNonbondedMethod(vfswitch.CutoffPeriodic)
        vfswitch.setCutoffDistance(nonbonded.getCutoffDistance())
        for i in range(nonbonded.getNumParticles()):
            chg, sig, eps = nonbonded.getParticleParameters(i)
            nonbonded.setParticleParameters(i, chg, 0.0, 0.0)  # zero-out LJ
            sig = sig * 0.5
            eps = eps**0.5
            vfswitch.addParticle([sig, eps])
        for i in range(nonbonded.getNumExceptions()):
            atom1, atom2 = nonbonded.getExceptionParameters(i)[:2]
            vfswitch.addExclusion(atom1, atom2)
        vfswitch.setForceGroup(psf.NONBONDED_FORCE_GROUP)
        system.addForce(vfswitch)

    else:
        # vfswitch w/ NBFIX
        nbfix.setEnergyFunction(
            "select(step(r-Ron),(cr12*rjunk12 - cr6*rjunk6),"
            "(ccnba/r^12-ccnbb/r^6-ccnba/onoff^6+ccnbb/onoff^3));"
            "cr12 = ccnba*ofdif6;"
            "cr6  = ccnbb*ofdif3;"
            "rjunk12 = (1.0/r^6-1.0/Roff6)^2;"
            "rjunk6 = (1.0/r^3-1.0/Roff3)^2;"
            "ccnbb = bcoef(type1, type2);"
            "ccnba = acoef(type1, type2)^2;"
            "ofdif6 = Roff6/(Roff6 - Ron^6);"
            "ofdif3 = Roff3/(Roff3 - Ron^3);"
            "onoff = Roff * Ron;"
            "Roff6 = Roff^6;"
            "Roff3 = Roff^3;"
            "Ron  = %f;"
            "Roff = %f;" % (r_on, r_off)
        )

        # turn off long range correction (OpenMM Issues: #2353)
        nbfix.setUseLongRangeCorrection(False)

    # vfswitch14
    vfswitch14 = CustomBondForce(
        "select(step(r-Ron),(cr12*rjunk12 - cr6*rjunk6),"
        "(ccnba/r^12-ccnbb/r^6-ccnba/onoff^6+ccnbb/onoff^3));"
        "cr12 = ccnba*ofdif6;"
        "cr6 = ccnbb*ofdif3;"
        "rjunk12 = (1.0/r^6-1.0/Roff6)^2;"
        "rjunk6 = (1.0/r^3-1.0/Roff3)^2;"
        "ccnba = 4.0*epsilon*sigma^12;"
        "ccnbb = 4.0*epsilon*sigma^6;"
        "ofdif6 = Roff6/(Roff6 - Ron^6);"
        "ofdif3 = Roff3/(Roff3 - Ron^3);"
        "onoff = Roff * Ron;"
        "Roff6 = Roff^6;"
        "Roff3 = Roff^3;"
        "Ron = %f;"
        "Roff = %f;" % (r_on, r_off)
    )
    vfswitch14.addPerBondParameter("sigma")
    vfswitch14.addPerBondParameter("epsilon")
    for i in range(nonbonded.getNumExceptions()):
        atom1, atom2, chg, sig, eps = nonbonded.getExceptionParameters(i)
        nonbonded.setExceptionParameters(
            i, atom1, atom2, chg, 0.0, 0.0
        )  # zero-out LJ14
        vfswitch14.addBond(atom1, atom2, [sig, eps])
    system.addForce(vfswitch14)

    return system