import simtk.unit as unit
import simtk.openmm as openMM

THICKNESS = 31.0 * unit.angstrom
SOLVENT_DIELECTRIC = 78.3
SOLUTE_DIELECTRIC = 1

def buildCustomMembraneGBForce(custom):
    custom.addGlobalParameter("thickness", THICKNESS)
    custom.addGlobalParameter("solventDielectric", SOLVENT_DIELECTRIC)
    custom.addGlobalParameter("soluteDielectric", SOLUTE_DIELECTRIC)
    custom.addPerParticleParameter("q")
    custom.addPerParticleParameter("radius")
    custom.addPerParticleParameter("scale")

    s = "step(r+sr2-or1)*0.5*(1/L-1/U+0.25*(1/U^2-1/L^2)*(r-sr2*sr2/r)+0.5*log(L/U)/r+C)"
    s+= "; U=r+sr2"
    s+= "; C=2*(1/or1-1/L)*step(sr2-r-or1)"
    s+= "; L=max(or1, D)"
    s+= "; D=abs(r-sr2)"
    s+= "; sr2 = scale2*or2"
    s+= "; or1 = radius1-0.009"
    s+= "; or2 = radius2-0.009"
    custom.addComputedValue("Imol", s,
                            openMM.CustomGBForce.ParticlePairNoExclusions)

    s = "(1/radius+2*log(2)/thickness)/(1+exp(7.2*(abs(z)+radius-0.5*thickness)))"
    custom.addComputedValue("Imem", s,
                            openMM.CustomGBForce.SingleParticle)

    s = "1/(1/or-tanh(1*psi-0.8*psi^2+4.85*psi^3)/radius)"
    s+= "; psi=max(Imol,Imem)*or"
    s+= "; or=radius-0.009"
    custom.addComputedValue("B", s,
                            openMM.CustomGBForce.SingleParticle)

    s = "28.3919551*(radius+0.14)^2*(radius/B)^6-0.5*138.935456*(1/soluteDielectric-1/solventDielectric)*q^2/B"
    custom.addEnergyTerm(s, openMM.CustomGBForce.SingleParticle)

    s = "-138.935456*(1/soluteDielectric-1/solventDielectric)*q1*q2/f"
    s+= "; f=sqrt(r^2+B1*B2*exp(-r^2/(4*B1*B2)))"
    custom.addEnergyTerm(s, openMM.CustomGBForce.ParticlePairNoExclusions)


#CustomGBForce* custom = new CustomGBForce();
#custom->addGlobalParameter("thickness", thickness);
#custom->addGlobalParameter("solventDielectric", 78.3);
#custom->addGlobalParameter("soluteDielectric", 1);
#custom->addPerParticleParameter("q");
#custom->addPerParticleParameter("radius");
#custom->addPerParticleParameter("scale");
#custom->addComputedValue("Imol", "step(r+sr2-or1)*0.5*(1/L-1/U+0.25*(1/U^2-1/L^2)*(r-sr2*sr2/r)+0.5*log(L/U)/r+C);"
#                             "U=r+sr2;"
#                             "C=2*(1/or1-1/L)*step(sr2-r-or1);"
#                             "L=max(or1, D);"
#                             "D=abs(r-sr2);"
#                             "sr2 = scale2*or2;"
#                             "or1 = radius1-0.009; or2 = radius2-0.009", CustomGBForce::ParticlePairNoExclusions);
#custom->addComputedValue("Imem", "(1/radius+2*log(2)/thickness)/(1+exp(7.2*(abs(z)+radius-0.5*thickness)))", CustomGBForce::SingleParticle);
#custom->addComputedValue("B", "1/(1/or-tanh(1*psi-0.8*psi^2+4.85*psi^3)/radius);"
#                             "psi=max(Imol,Imem)*or; or=radius-0.009", CustomGBForce::SingleParticle);
#custom->addEnergyTerm("28.3919551*(radius+0.14)^2*(radius/B)^6-0.5*138.935456*(1/soluteDielectric-1/solventDielectric)*q^2/B", CustomGBForce::SingleParticle);
#custom->addEnergyTerm("-138.935456*(1/soluteDielectric-1/solventDielectric)*q1*q2/f;"
#                     "f=sqrt(r^2+B1*B2*exp(-r^2/(4*B1*B2)))", CustomGBForce::ParticlePairNoExclusions);